From 7d7c6cbc68b0edae2a5214d41929c40098fe7af1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 00:06:35 +0000 Subject: [PATCH 01/26] Scope hosted monetization, usage metering, and evaluator BYOK Research note, not an implementation. ROADMAP carries "Multi-tenancy and billing" as an open operational item resting on cost_records as its metering foundation; this writes the design that foundation needs before it grows a payment integration. The position it lands on: sell operations and license terms, not code. The AGPL makes any paywall removable in one commit, so nothing here depends on it being unremovable, and no feature already shipped moves behind a wall. Hosted capacity is metered on the four things that actually cost the operator money -- model tokens, sandbox time, deploy slots, stored bytes -- and the rest of the product stays free. Enforcement generalizes the src/beta/gate.ts pattern: an entitlements interface in the AGPL tree, defaulting to unlimited so self-hosters see no change, with plan definitions and billing behind a service URL that is not in this repository. It notes two things the metering needs that cost_records lacks (owner_id, and a platform/byok source column), and that spend counters belong in a Durable Object rather than the rate limiter's racy KV read-modify-write. BYOK ships before any cap rather than after it, on project_secrets, which already has the AAD-bound write-only shape the keys need. The security work is the policy file being repository content: an attacker-editable baseUrl would forge requests carrying the diff, so providers are selected by name against an operator allowlist. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- docs/research/monetization-and-byok.md | 397 +++++++++++++++++++++++++ 1 file changed, 397 insertions(+) create mode 100644 docs/research/monetization-and-byok.md diff --git a/docs/research/monetization-and-byok.md b/docs/research/monetization-and-byok.md new file mode 100644 index 00000000..2a5eef9a --- /dev/null +++ b/docs/research/monetization-and-byok.md @@ -0,0 +1,397 @@ +# Monetization and BYOK + +**Status:** research / proposal. Nothing here is implemented, and nothing here +changes what a self-hoster gets today. + +This document scopes how Stratum charges for hosted usage without withholding +code from the open-source repository, what the free tier costs the operator, and +how a project supplies its own model credentials (BYOK) instead of consuming a +hosted allowance. + +It exists because [`ROADMAP.md`](../../ROADMAP.md) carries **Multi-tenancy and +billing** as an open operational item with "per-change cost tracking already +exists and provides the metering foundation", and that foundation needs a design +before it grows a payment integration around it. + +--- + +## 1. The constraint the license already set + +[`LICENSING.md`](../../LICENSING.md) settled two things that decide most of this +design, and re-litigating them here would be wasted effort. + +The server is **AGPL-3.0-or-later**. Anyone may fork it, delete a paywall, and +run it — commercially. §13 only obliges them to publish the modified source. +That is not a loophole to close; it is the deal the project made deliberately at +v0.3.0. Any monetization plan that depends on a self-hoster being unable to +remove a check is already broken, so no plan below depends on that. + +There is a **CLA** ([`CLA.md`](../../CLA.md)), so dual licensing is available, +and `LICENSING.md` already advertises commercial terms. That is a second revenue +line that requires no code split at all, and it is worth more than it looks: +a large share of enterprises run an OSS policy that flatly refuses AGPL +dependencies. For those buyers the product is not a feature — it is the license. + +What follows from that: **sell operations and terms, not code.** The hosted +instance sells the thing a fork cannot copy — running it, on someone else's +Cloudflare bill, with backups tested and an SLA. The commercial license sells +the thing the AGPL cannot give. Neither requires a proprietary `ee/` directory, +and both survive a fork. + +### The one gap to close first + +Trademark is the only enforcement lever the AGPL leaves against a hosted +competitor, and the repository does not currently exercise it. There is no +`TRADEMARK.md`, and "Stratum" is not reserved anywhere in the licensing docs. +A fork may run the code; it may not call itself Stratum or use the mark to sell +against you. Registering the mark and adding a short trademark policy is cheap, +takes a week, and should land before any of the engineering below. + +--- + +## 2. What actually costs money + +The honest basis for a paywall is marginal cost to the operator plus scarcity of +a shared resource. Sorted by both, using what the code already meters and bounds: + +**Unbounded per-change spend.** The LLM evaluator (`src/evaluation/llm-evaluator.ts`) +runs on the Workers AI binding — the operator's neurons, not the author's. It is +invoked from `runEvaluation` in `src/services/change-flow.ts`, **synchronously at +change creation**, with a per-call input bound of `MAX_DIFF_CHARS_CEILING` +(100,000 chars) and no bound at all on how many times a user may create a change. +This is the single largest exposure and the obvious first thing to meter. + +**Shared scarce capacity.** Deployments are worse than expensive — they are +serialized. `MAX_DEPLOY_CONCURRENCY = 1` in `src/deploy/limits.ts`, and the +comment there explains it cannot be raised without first cutting the per-deploy +memory peak. One free user's deploy queue is therefore in front of a paying +customer's. Same shape, less severe, for the import queue and the merge queue +Durable Object. + +**Sandbox milliseconds.** Bounded per run by `MAX_TOTAL_BUDGET_MS` +(`src/evaluation/limits.ts`), unbounded in count. Currently theoretical: the +`[[sandboxes]]` binding is commented out at `wrangler.toml:156`, so nothing +executes code on any deploy today. It becomes the second-largest exposure the +day that binding is enabled. + +**Bytes that never go away.** Git objects on Artifacts and R2, backup blobs +(`BACKUP_RETENTION`, default 14 runs), and D1 rows — `events`, `audit_log`, +`provenance`, `cost_records`. The roadmap's "D1 hot/cold rotation" item is the +same problem seen from the operations side; retention tiers are the same lever +seen from the pricing side, and they should be designed together rather than +twice. + +**Near zero.** Everything else. The change flow itself, reviews, issues, orgs and +teams, the CLI, `/mcp`, git clone and push, webhooks, API tokens, and the +deterministic evaluators (`diff-evaluator`, `secret-scanner`) are Worker CPU on +requests that are already rate-limited. Gating any of these recovers no money +and costs the project its reputation. They stay free, permanently, for everyone. + +That last paragraph is the load-bearing one. The credible version of this +paywall meters four things — model tokens, sandbox time, deploy slots, stored +bytes — and leaves the product alone. + +--- + +## 3. Tier shape + +Concrete allowances are deliberately omitted; §7 says to measure before setting +them. What matters now is which axis each tier moves along. + +**Free.** Unlimited public projects and unlimited collaborators on them. A small +number of private projects. A monthly platform-LLM allowance sized to a few dozen +real changes. Sandbox off. One deploy target. A storage cap. Short retention on +events and audit entries. Community support. **BYOK available** — see §5. + +**Pro**, per seat. Private projects unlimited, a real platform-LLM allowance, +sandbox on, deploys, 90-day retention, a higher rate-limit bucket, priority in +the deploy and import queues, email support. + +**Team**, per seat. Orgs, teams, per-project permission grants (a roadmap item +today), SSO, one-year audit retention with export, self-serve backup restore, +higher concurrency, an SLA. + +**Enterprise self-managed.** A commercial license per `LICENSING.md`, support, +and no usage metering whatsoever, because it runs on their infrastructure. + +Two notes on the boundaries. Do not put SSO alone behind the top tier as the +only reason to reach it — the "SSO tax" is the most reliably resented pattern in +this market. It belongs in Team alongside SCIM, audit export, and permission +grants, as one item in a coherent administrative bundle rather than the toll. +And free-tier caps must be **hard**, with an explicit top-up, not silent +overage: Workers AI and queue time are real spend and a loop is one script away. +Paid tiers get soft overage with alerts, because a blocked merge on a paying +customer's release day is worth more than the overage. + +--- + +## 4. Enforcement architecture + +The pattern to copy is already in the repository. `src/beta/gate.ts` is a +cloud-only hook that lives in the AGPL tree, calls out to a service that is not +in this repository, and is **completely inert when its environment variables are +unset** — which is the default, so self-hosters are unaffected and the code that +does the gating is still readable by everyone. Generalize exactly that. + +### 4.1 Entitlements + +A new `src/billing/entitlements.ts` defines the interface and ships a default +implementation that grants everything: + +```ts +export type LimitKey = + | "llm_tokens_month" + | "sandbox_ms_month" + | "deploys_month" + | "storage_bytes" + | "private_projects" + | "requests_per_minute" + | "retention_days"; + +export interface Entitlements { + plan: string; // "self-hosted" by default + limits: Record; // -1 = unlimited +} + +export interface EntitlementsProvider { + forOwner(ownerId: string, kind: "user" | "org"): Promise; +} +``` + +`UnlimitedEntitlements` is the default and returns `-1` for every key. A +`RemoteEntitlements` provider activates only when `BILLING_SERVICE_URL` and +`BILLING_SERVICE_SECRET` are set, mirroring `betaGateEnabled`. Plan definitions, +Stripe, invoices, dunning, and the customer portal all live behind that URL, in +the same private cloud layer that already hosts the referral service. None of it +belongs in this repository, and none of it needs to be. + +**Fail direction, and why it differs from the beta gate.** `validateInviteCode` +fails *closed* — a service outage must not let ungated users through a signup +wall. Entitlements must fail *open*, to the last cached value or to the free +tier, because a billing-service outage that blocks a paying customer's merge is +a strictly worse failure than a few minutes of unmetered usage. Cache in +`env.STATE` with a short TTL and a stale-while-error read. Write this asymmetry +into the module's doc comment; it will otherwise be "fixed" by someone applying +the beta gate's convention. + +### 4.2 Metering + +There are two counters with different jobs, and conflating them is the classic +way to build a billing system that is both slow and wrong. + +`cost_records` (migration 021) stays the **durable ledger**. It is already +written by `recordCosts` from `runEvaluation`, already carries +`llm_tokens | sandbox_ms | git_ops` and an `estimated` flag, and it is what an +invoice or a dispute is reconciled against. It needs two additive migrations: + +- **`owner_id`** — the table records `project`, `project_id`, `change_id`, and + `workspace`, but nothing identifying who pays. Every per-customer aggregate is + impossible until this exists, and backfilling it later means joining through + namespaces for historical rows. Add it before the volume grows. +- **`source`** (`'platform' | 'byok'`) — §5 needs to attribute usage that must + not be charged. Without it, BYOK tokens are indistinguishable from hosted ones + in the ledger. + +The **hot counter** that gates a request is separate. It must not lose +increments under concurrency, which rules out the read-modify-write on KV that +`src/middleware/rate-limit.ts` currently performs — that race is acceptable for a +per-minute request limiter and is not acceptable for anything that gates spend. +The codebase already has the right precedent: `MagicLinkRateLimiter`, described +in `src/types.ts` as "serialized magic-link send counters". Add a `UsageMeter` +Durable Object per owner on the same model, reconciled against `cost_records` by +the existing scheduled handler. + +The rule to state once and follow: **Durable Object for anything that gates +spend, KV for anything that gates request rate.** + +### 4.3 Where the checks go + +Four call sites, and one design decision inside the first. + +*Before the LLM evaluator runs*, in `buildEvaluators` or at the top of +`LLMEvaluator.evaluate`. The decision is what an exhausted allowance means for a +merge *gate*. Silently skipping the evaluator is a governance hole — a policy +that requires AI review would stop requiring it the moment someone ran out of +credit, which is the worst possible failure for the product's entire premise. +Instead, mirror the existing behavior in that file for an oversize policy +context: return `ok()` with `passed: false` and a reason naming the cause and +both remedies (upgrade, or add your own key). That is already the established +shape at `llm-evaluator.ts` and it converts honestly. + +*Before claiming a deploy message*, in the queue consumer, so a free-tier deploy +never occupies the single concurrency slot ahead of a paid one. + +*At import and at `git-receive-pack`*, for storage bytes. Pushes already flow +through the rate limiter deliberately (`isExemptGitRead` exempts only the read +side), so the hook point exists. + +*At private-project creation*, for the count. + +And `rateLimitMiddleware`'s `defaultLimit` becomes entitlement-derived rather +than the current binary `isAuthenticated ? 1000 : 60`. + +--- + +## 5. BYOK + +BYOK is the feature that makes this a capacity limit rather than a paywall on +the open-source project, and it should ship *before* any cap is enforced. A free +user who supplies a key gets unlimited AI review, bounded only by request rate. +That framing is both true and the honest thing to say publicly. + +### 5.1 The agent already does it + +`agent/src/llm.ts` calls the Anthropic API directly with the user's own +`ANTHROPIC_API_KEY`, from the user's own machine, under Apache-2.0. The most +expensive model work in the whole system — actually writing the code — already +never touches the server and never will. Say so in the pricing page; it removes +the first objection before it is raised. + +### 5.2 What is missing: server-side evaluator BYOK + +`LLMEvaluator` takes an `AiBinding` in its constructor and is hard-wired to +Workers AI. Introducing a provider seam is a small refactor that is worth doing +for its own sake, independent of billing — it is what lets a self-hoster use a +frontier model, or a local vLLM, for their merge gate. + +```ts +export interface LlmProvider { + run(model: string, messages: Message[]): Promise<{ text: string; usage?: Usage }>; +} +``` + +Three implementations: `workers-ai` wrapping the existing binding and remaining +the default so nothing changes for anyone; `anthropic`; and `openai-compatible`, +which covers OpenAI, OpenRouter, Groq, Together, and self-hosted vLLM in one. +A real provider also returns real token counts, which retires the +`~4 chars/token` estimate and the `estimated: true` flag on those cost samples. + +### 5.3 Key storage — already built + +`project_secrets` (migration 047) is the right store and needs no redesign. It +is AES-GCM with `(project_id, name)` bound as additional authenticated data, so +a row copied to another project fails to decrypt rather than authorizing against +the wrong account; it is write-only by construction with no read path reachable +from any route; and `SECRET_NAME_PATTERN` already accepts `ANTHROPIC_API_KEY` +and `OPENAI_API_KEY` unchanged. The evaluator gets a `loadSecretValues` path +alongside the deploy runner's, held to the same rule: not reachable from a +route. + +The one extension needed is scope. `SecretScope` in `src/utils/crypto.ts` is +project-shaped, and a team will want one key for an org rather than one per +project. Add an org-scoped variant of the AAD; do not widen the existing one. + +### 5.4 The security problem, which is the real work here + +`.stratum/policy.yaml` is **repository content**. Anyone who can open a change +can propose an edit to it. Adding a `provider` and `baseUrl` to the evaluator +config therefore hands an attacker a request forgery primitive on the Worker, +and the body of that request is the diff — the customer's source code, plus +whatever the secret scanner was about to catch. + +Three mitigations, all mandatory: + +1. **`baseUrl` is never taken from the policy file.** The operator sets an + allowlist in the environment (`LLM_BYOK_ALLOWED_HOSTS`), and the policy may + only *select* a configured provider by name. This is not negotiable and is + the reason the config surface is `provider: anthropic` rather than a URL. +2. **Sanitize like `deploys:` already does.** `sanitizeDeploys` in + `src/deploy/config.ts` is the pattern: build a fresh object, never pass a + slice of parsed user input downstream, and surface every rejection as a + visible `deployRejections`-style entry rather than a dropped `logger.warn`. +3. **Tie into the existing protected-config signal.** Migration 039 added + `changes.touches_protected_config`. A change that edits the evaluator's + provider block is exactly what that flag is for, and it should require review + at a level above an ordinary change. + +Note also that the key itself is a redaction surface: `src/deploy/redact.ts` +already exists for provider output on `deployments.log_tail`, and provider error +bodies from an LLM call need the same treatment before they reach an +`EvalResult.reason`. + +### 5.5 Billing consequence + +BYOK usage is recorded with `source = 'byok'` and does not decrement the +`llm_tokens_month` allowance. It is not free to the operator, though — it still +costs a subrequest and Worker wall time — so BYOK moves a project from a token +quota onto the request-rate quota rather than off all quotas. State that plainly +in the docs; a user who discovers it at a 429 will assume they were misled. + +### 5.6 The pattern is already the house style + +Deploy credentials are BYO (`project_secrets` holds the Vercel and Cloudflare +tokens), and BYO-R2 for backups is a natural next one. "Stratum orchestrates, +you own the accounts it spends against" is a coherent product position, not a +concession, and it is materially easier to sell to a security team than a +platform that proxies everything. + +--- + +## 6. Abuse + +Signup is currently gated by invite codes (`src/beta/gate.ts`), which is doing +the anti-abuse work today and hides the problem. When that gate lifts, a free +LLM allowance on an open signup becomes a target. + +The design has an unusual advantage here: the evaluator is not a chat endpoint. +`SYSTEM_PROMPT` is fixed in the source, the output is parsed as a strict verdict +object, and anything else fails closed. It is close to useless as a general +model proxy, which removes the main motive. What remains is ordinary +free-tier farming, handled by keying quota on the **owner** rather than the +project — otherwise forking a project resets the counter — and by keeping the +existing verified-identity signup paths (GitHub, Google, magic link). + +--- + +## 7. Sequencing + +Nothing here should ship in the order it was written. The order that matters: + +**0. Trademark policy and registration.** Independent of all engineering, and +the only real lever against a hosted fork. Do it first because it has a lead +time. + +**1. Measure.** `cost_records` is already being written and +`getProjectCostSummary` in `src/storage/costs.ts` is currently **dead code** — +defined, exported, never called. Add `owner_id` and `source`, then build the +per-owner usage view on top of that existing function. Run it for a month +against real traffic on `app.usestratum.dev` before choosing a single number in +§3. Every allowance in this document is deliberately unspecified for that reason. + +**2. Entitlements interface**, with `UnlimitedEntitlements` wired in as the +default. Zero behavior change on merge, which makes it reviewable on its own. + +**3. `UsageMeter` Durable Object** plus reconciliation against `cost_records`. +Still no enforcement. + +**4. BYOK.** Ships as a feature to self-hosters — pick your own model for the +merge gate — and lands before any cap so the escape hatch predates the wall. + +**5. The usage UI.** Non-negotiable prerequisite for step 6. If the first thing +a user learns about a quota is a blocked merge, the quota was a bug report. + +**6. Cloud entitlements provider and Stripe**, in the private layer. + +**7. Turn on free-tier caps**, announced at least 30 days ahead, with the usage +page live throughout. + +--- + +## 8. What not to do + +Stated explicitly, because each of these is the default mistake: + +Do not add a license-key check to the AGPL server, or obfuscate any part of it. +It is removable in one commit, it reads as bad faith, and it buys nothing. + +Do not move a shipped feature behind the wall. Everything in +`docs/CURRENT_CAPABILITIES.md` shipped under MIT or AGPL. Taking any of it back +costs more in credibility than it can return in revenue. Only *new* hosted +capacity limits, and self-host stays unlimited by default. + +Do not degrade the open-source build to make the hosted one look better. The +self-hosted path is the top of the funnel for the commercial license, which is +the higher-margin product of the two. + +Do not gate collaborators or public projects. Free multiplayer on public work is +how the product gets seen at all. From 9f475871feee91a804ef8e60dcb62b91dbd795f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 00:33:12 +0000 Subject: [PATCH 02/26] Give the evaluator layer a project and a billing subject Task 1 of the metering/BYOK work, and a pure refactor: no behavior change beyond one log line, no new feature. buildEvaluators took a display name, and runEvaluation took no project at all, so nothing at the evaluator layer could name the project it was judging or the account that pays for what it spends. The LLM evaluator spends the operator's Workers AI budget on every call. Both now receive the ProjectEntry, and an optional BillingContext rides on EvaluationContext so evaluators that consume nothing metered can ignore it. billingContextFor yields nothing for an agent-owned project rather than guessing. ProjectEntry.ownerType admits "agent", which is not a payer but belongs to one; walking agents.owner_id to that user is the attribution work that follows, and attributing spend to an agent id would be worse than attributing none. The conflict-resolution path in sync-management gets the context too. It runs the same evaluator suite, LLM included, and is the one metered path that records no cost today. The only observable difference: the unknown-evaluator warning from the re-evaluation route now names a project as "@namespace/slug" rather than by bare name. projectName is log-only inside buildEvaluators. Also fixes package-lock.json, which still declared MIT after the relicense to AGPL-3.0-or-later. Typecheck clean; 3716 tests pass. Task 1's own tests are still outstanding. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- package-lock.json | 2 +- src/evaluation/types.ts | 31 +++++++++++++++++++++++++ src/routes/changes.ts | 4 +++- src/routes/sync-management.ts | 9 +++++--- src/services/change-flow.ts | 40 ++++++++++++++++++++++++++++++--- tests/sandbox-evaluator.test.ts | 18 +++++++++++---- 6 files changed, 92 insertions(+), 12 deletions(-) diff --git a/package-lock.json b/package-lock.json index 5cdf67f7..3dd28338 100644 --- a/package-lock.json +++ b/package-lock.json @@ -7,7 +7,7 @@ "": { "name": "stratum", "version": "0.2.0", - "license": "MIT", + "license": "AGPL-3.0-or-later", "dependencies": { "@isomorphic-git/lightning-fs": "^4.6.2", "diff": "^9.0.0", diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index d22f04b5..dac6d31d 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -11,6 +11,28 @@ export interface EvalResult { costs?: Array<{ kind: "llm_tokens" | "sandbox_ms"; quantity: number; estimated?: boolean }>; } +/** + * Who pays for the metered resources an evaluation consumes. + * + * The evaluator layer has no project today — `buildEvaluators` was given a + * display name and `runEvaluation` nothing at all — so the LLM evaluator's + * spend lands on the operator's Workers AI account with nothing recording + * whose change caused it. Carrying the subject on the evaluation context is + * what makes that spend attributable and, later, meterable, without every + * evaluator having to grow a project-shaped constructor argument. + * + * `ownerType` is narrower than `ProjectEntry.ownerType`, which also admits + * `"agent"`: an agent is not a billing subject, it resolves to the user that + * owns it. That resolution is deliberately not done here. + */ +export interface BillingContext { + /** The paying user or org. Never an agent id — see `ownerType`. */ + ownerId: string; + ownerType: "user" | "org"; + /** The project whose policy is being enforced. Keys per-project credentials. */ + projectId: string; +} + /** * What the diff is a diff *of*. A diff alone does not identify the tree it * applies to, so an evaluator that reproduces the change out-of-process (the @@ -28,6 +50,15 @@ export interface EvaluationContext { * commit would report a verdict for a combination the change never proposed. */ baseSha?: string; + /** + * Who pays for whatever this evaluation spends. + * + * Optional so evaluators that consume nothing metered — every evaluator + * today except `llm` — can ignore it entirely. Absent where no billing + * subject can be named, which is deliberately not the same as free: an + * agent-owned project has a payer, it just has to be resolved first. + */ + billing?: BillingContext; } export interface Evaluator { diff --git a/src/routes/changes.ts b/src/routes/changes.ts index caca75af..8b7a566a 100644 --- a/src/routes/changes.ts +++ b/src/routes/changes.ts @@ -9,6 +9,7 @@ import { enqueueMergeDeploy } from "../queue/deploy-queue"; import { emitEvent } from "../queue/events"; import type { MergeOutcome } from "../queue/merge-queue"; import { + billingContextFor, buildEvaluators, createChangeWithEvaluation, resolveProjectHead, @@ -1397,7 +1398,7 @@ app.post("/changes/:id/evaluate", async (c) => { baseOid, } = diffResult.data; - const evaluators = buildEvaluators(c.env, policy, change.project, logger, { + const evaluators = buildEvaluators(c.env, policy, project, logger, { remote: workspace.remote, token: workspaceReadToken.data, ref: evaluatedSha, @@ -1407,6 +1408,7 @@ app.post("/changes/:id/evaluate", async (c) => { // stale by the time a change is re-evaluated (#274). const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha: baseOid, + billing: billingContextFor(project), }); const recordResult = await recordEvalRuns(c.env.DB, logger, id, evalRuns); diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index de472a2a..beb7c5e8 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -3,7 +3,7 @@ import { loadPolicy } from "../evaluation/policy-loader"; import { scanContentForSecrets } from "../evaluation/secret-scanner"; import { checkResolutionMergeProtection } from "../merge/protection"; import { authMiddleware } from "../middleware/auth"; -import { buildEvaluators, runEvaluation } from "../services/change-flow"; +import { billingContextFor, buildEvaluators, runEvaluation } from "../services/change-flow"; import { recordAudit } from "../storage/audit"; import { getChange } from "../storage/changes"; import { @@ -656,17 +656,20 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { const { diff, baseSha } = diffResult.data; const policy = await loadPolicy(project.remote, projectToken.data, logger, branch); - const projectName = `${conflictCtx.namespace}/${conflictCtx.slug}`; // No workspace repo access is passed for the sandbox evaluator: the content // being judged here has no commit of its own yet — that's the point, it must // pass BEFORE resolveConflict creates one — so there is no ref a sandbox // could check out. A policy naming `sandbox` fails closed via // UnavailableEvaluator, same as any other missing prerequisite. - const evaluators = buildEvaluators(c.env, policy, projectName, logger); + const evaluators = buildEvaluators(c.env, policy, project, logger); // `buildManualResolutionDiff` resolved this base from the clone it built the // diff on, so it names the tree the resolution actually applies to (#274). + // This path runs the LLM evaluator too, so it needs a payer as much as + // change creation does — it is the one metered path that records nothing + // today. const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha, + billing: billingContextFor(project), }); if (!evalResult.passed) { diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index de199e6d..f674414c 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -9,7 +9,13 @@ import { loadPolicy, } from "../evaluation"; import type { SandboxRepoAccess } from "../evaluation/sandbox-evaluator"; -import type { EvalPolicy, EvalResult, EvaluationContext, Evaluator } from "../evaluation/types"; +import type { + BillingContext, + EvalPolicy, + EvalResult, + EvaluationContext, + Evaluator, +} from "../evaluation/types"; import { buildEvaluationReport, reportEvaluationToGitHub } from "../github/sync"; import { type EventActor, emitEvent } from "../queue/events"; import { getAgent } from "../storage/agents"; @@ -96,11 +102,34 @@ export class UnavailableEvaluator implements Evaluator { } } +/** + * The billing subject an evaluation of `project` should be charged to, or + * `undefined` when the project has none that can be named yet. + * + * `ProjectEntry.ownerType` admits `"agent"`, which `BillingContext` does not: + * an agent is not a payer, it belongs to one. Walking `agents.owner_id` to that + * user is `resolveBillingSubject`'s job in the attribution work that follows — + * so an agent-owned project yields no context here rather than a guessed one. + * Attributing spend to an agent id would be worse than attributing none. + */ +export function billingContextFor(project: ProjectEntry): BillingContext | undefined { + if (project.ownerType === "agent") return undefined; + return { + ownerId: project.ownerId, + ownerType: project.ownerType, + projectId: project.id, + }; +} + /** * Build the evaluator set for a policy: the always-on blocking secret scan plus * whatever the policy configures. Evaluators whose binding is missing become * UnavailableEvaluator (score 0, fail) rather than silently vanishing. * + * Takes the whole `ProjectEntry` rather than a display name so an evaluator can + * be told who owns the project it is judging; the display name every call site + * used to pass is derived here so no caller has to build it twice. + * * `workspaceRepo` is read access to the workspace being evaluated (remote + * read token + the pinned evaluated commit); the sandbox evaluator needs it to * materialize the full tree it runs the test command against. Without it — or @@ -111,10 +140,14 @@ export class UnavailableEvaluator implements Evaluator { export function buildEvaluators( env: Env, policy: EvalPolicy, - projectName: string, + project: ProjectEntry, logger: Logger, workspaceRepo?: SandboxRepoAccess, ): Array<{ type: string; evaluator: Evaluator }> { + // The whole entry, not a name: an evaluator that spends money has to be able + // to name the payer, and a "@ns/slug" string identifies no owner and no + // project id. Only the log line below needs the human-readable form. + const projectName = `${project.namespace}/${project.slug}`; const evaluators: Array<{ type: string; evaluator: Evaluator }> = [ { type: "secret_scan", evaluator: new SecretScanEvaluator() }, ]; @@ -393,7 +426,7 @@ export async function createChangeWithEvaluation( baseOid, } = diffResult.data; - const evaluators = buildEvaluators(env, policy, projectName, logger, { + const evaluators = buildEvaluators(env, policy, project, logger, { remote: workspaceRemote, token: workspaceReadToken.data, ref: evaluatedSha, @@ -403,6 +436,7 @@ export async function createChangeWithEvaluation( // base must be the one the diff was actually built against (#274). const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha: baseOid, + billing: billingContextFor(project), }); const newStatus: Change["status"] = evalResult.passed ? "accepted" : "needs_changes"; diff --git a/tests/sandbox-evaluator.test.ts b/tests/sandbox-evaluator.test.ts index 422825c7..c589f15f 100644 --- a/tests/sandbox-evaluator.test.ts +++ b/tests/sandbox-evaluator.test.ts @@ -21,7 +21,7 @@ import type { EvalPolicy } from "../src/evaluation/types"; import { buildEvaluators } from "../src/services/change-flow"; import { type NodeFS, readTreeAtCommit } from "../src/storage/git-ops"; import { MemoryFS } from "../src/storage/memory-fs"; -import type { Env, SandboxBinding, SandboxInstance } from "../src/types"; +import type { Env, ProjectEntry, SandboxBinding, SandboxInstance } from "../src/types"; import { AppError } from "../src/utils/errors"; import type { Logger } from "../src/utils/logger"; import { err, ok } from "../src/utils/result"; @@ -786,6 +786,16 @@ describe("SandboxEvaluator — reason field", () => { describe("buildEvaluators — sandbox wiring", () => { const sandboxPolicy: EvalPolicy = { evaluators: [{ type: "sandbox" }] }; + const project: ProjectEntry = { + id: "proj_sandbox", + name: "proj", + slug: "proj", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://artifacts.example.com/repos/proj", + createdAt: "2026-01-01T00:00:00.000Z", + }; function findSandbox(evaluators: ReturnType) { const entry = evaluators.find((e) => e.type === "sandbox"); @@ -795,7 +805,7 @@ describe("buildEvaluators — sandbox wiring", () => { } it("no SANDBOX binding → fails closed with an actionable wrangler.toml reason", async () => { - const evaluators = buildEvaluators({} as Env, sandboxPolicy, "proj", mockLogger, repo); + const evaluators = buildEvaluators({} as Env, sandboxPolicy, project, mockLogger, repo); const result = await findSandbox(evaluators).evaluate("", sandboxPolicy, mockLogger); expect(result.success).toBe(true); if (!result.success) return; @@ -810,7 +820,7 @@ describe("buildEvaluators — sandbox wiring", () => { const evaluators = buildEvaluators( { SANDBOX: binding } as Env, sandboxPolicy, - "proj", + project, mockLogger, ); const result = await findSandbox(evaluators).evaluate("", sandboxPolicy, mockLogger); @@ -826,7 +836,7 @@ describe("buildEvaluators — sandbox wiring", () => { const evaluators = buildEvaluators( { SANDBOX: binding } as Env, sandboxPolicy, - "proj", + project, mockLogger, repo, ); From b8ee983e39fe6122b4bc7006c710760f8d8c6e13 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 00:40:18 +0000 Subject: [PATCH 03/26] Test that every evaluation names its payer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 1.5. Seventeen tests covering the three call sites that build evaluators, asserting on the EvaluationContext each evaluator is actually handed rather than on any route's response body — the context is the only observable effect this refactor has. Recording stand-ins replace the secret scanner and the LLM evaluator. The secret scan is unconditional, so it is the one evaluator guaranteed to observe the context at all three sites. Covers billingContextFor across user, org and agent ownership; that buildEvaluators still constructs the same evaluator per policy type, including the fail-closed UnavailableEvaluator for absent AI and SANDBOX bindings; that runEvaluation forwards the context to every evaluator; and that each of the three call sites drives a real code path to a populated subject — including the org case, and the agent case where the subject is deliberately absent. 3733 tests pass, up 17 from the baseline. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- tests/evaluator-billing-context.test.ts | 524 ++++++++++++++++++++++++ 1 file changed, 524 insertions(+) create mode 100644 tests/evaluator-billing-context.test.ts diff --git a/tests/evaluator-billing-context.test.ts b/tests/evaluator-billing-context.test.ts new file mode 100644 index 00000000..5604ce25 --- /dev/null +++ b/tests/evaluator-billing-context.test.ts @@ -0,0 +1,524 @@ +import { Hono } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { EvalPolicy, EvaluationContext, Evaluator } from "../src/evaluation/types"; +import { authMiddleware } from "../src/middleware/auth"; +import type { Change, Env, ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; + +/** + * Task 1 is a pure refactor whose only observable effect is what reaches the + * evaluator layer: `buildEvaluators` now takes the whole `ProjectEntry`, and + * every `runEvaluation` call site carries a `BillingContext`. So these tests + * assert on the `EvaluationContext` each evaluator is actually handed, at each + * of the three call sites, rather than on any route's response body. + */ + +/** Contexts captured from the evaluator classes the call sites really build. */ +const captured = vi.hoisted(() => ({ + calls: [] as Array<{ evaluator: string; context?: EvaluationContext }>, +})); + +function recordingEvaluatorClass(name: string) { + return class { + async evaluate( + _diff: string, + _policy: EvalPolicy, + _logger: Logger, + context?: EvaluationContext, + ) { + captured.calls.push({ evaluator: name, context }); + return { success: true as const, data: { score: 1, passed: true, reason: `${name} ok` } }; + } + }; +} + +// The secret scan is on every evaluator set unconditionally, so it is the one +// evaluator guaranteed to observe the context at all three call sites. +vi.mock("../src/evaluation/secret-scanner", async (importActual) => ({ + ...(await importActual()), + SecretScanEvaluator: recordingEvaluatorClass("secret_scan"), +})); + +vi.mock("../src/evaluation/llm-evaluator", async (importActual) => ({ + ...(await importActual()), + LLMEvaluator: recordingEvaluatorClass("llm"), +})); + +vi.mock("../src/evaluation/policy-loader", async (importActual) => ({ + ...(await importActual()), + // The real loader clones the project repo to read .stratum/policy.yaml. + loadPolicy: vi.fn(), +})); + +vi.mock("../src/storage/git-ops", async (importActual) => ({ + ...(await importActual()), + freshRepoToken: vi.fn(async () => ({ success: true, data: "test-token" })), + getCommitLog: vi.fn(), + getDiffBetweenRepos: vi.fn(), + buildManualResolutionDiff: vi.fn(), + // The real one clones both repos and pushes; the conflict route's gate runs + // before it, which is the part under test here. + resolveConflict: vi.fn(async () => ({ success: true, data: { commitSha: "resolved_sha" } })), +})); + +vi.mock("../src/storage/sync", async (importActual) => ({ + ...(await importActual()), + recordSyncHistory: vi.fn(async () => undefined), +})); + +vi.mock("../src/storage/audit", async (importActual) => ({ + ...(await importActual()), + recordAudit: vi.fn(async () => ({ success: true, data: undefined })), +})); + +vi.mock("../src/storage/repo-snapshot", () => ({ + readRepoSnapshot: vi.fn(async () => ({ success: true, data: null })), +})); + +vi.mock("../src/storage/state", async (importActual) => ({ + ...(await importActual()), + getProject: vi.fn(), + getProjectByPath: vi.fn(), + getWorkspace: vi.fn(), +})); + +vi.mock("../src/storage/changes", async (importActual) => ({ + ...(await importActual()), + createChange: vi.fn(), + getChange: vi.fn(), + updateChangeStatus: vi.fn(async () => ({ success: true, data: undefined })), + dismissApprovalsAndUpdateStatus: vi.fn(async () => ({ + success: true, + data: { dismissedReviewerIds: [] }, + })), +})); + +vi.mock("../src/storage/eval-runs", async (importActual) => ({ + ...(await importActual()), + recordEvalRuns: vi.fn(async () => ({ success: true, data: [] })), +})); + +vi.mock("../src/storage/deletion", () => ({ + isTargetDeleting: vi.fn(async () => false), +})); + +vi.mock("../src/queue/events", () => ({ + emitEvent: vi.fn(async () => undefined), +})); + +vi.mock("../src/storage/users", () => ({ + getUserByToken: vi.fn(), + getUser: vi.fn(), +})); + +vi.mock("../src/storage/agents", () => ({ + getAgentByToken: vi.fn(async () => ({ success: false, error: { code: "NOT_FOUND" } })), + getAgent: vi.fn(), +})); + +import { DiffEvaluator } from "../src/evaluation/diff-evaluator"; +import { LLMEvaluator } from "../src/evaluation/llm-evaluator"; +import { loadPolicy } from "../src/evaluation/policy-loader"; +import { SecretScanEvaluator } from "../src/evaluation/secret-scanner"; +import { WebhookEvaluator } from "../src/evaluation/webhook-evaluator"; +import { changesRouter } from "../src/routes/changes"; +import { syncManagementRouter } from "../src/routes/sync-management"; +import { + UnavailableEvaluator, + billingContextFor, + buildEvaluators, + createChangeWithEvaluation, + runEvaluation, +} from "../src/services/change-flow"; +import { getChange } from "../src/storage/changes"; +import { + buildManualResolutionDiff, + getCommitLog, + getDiffBetweenRepos, +} from "../src/storage/git-ops"; +import { getProject, getProjectByPath, getWorkspace } from "../src/storage/state"; +import { getUserByToken } from "../src/storage/users"; +import { createLogger } from "../src/utils/logger"; + +const logger = createLogger({ component: "test" }); + +const USER_AUTH = { Authorization: "Bearer stratum_user_testtoken00000000000000000" }; + +function projectEntry(overrides: Partial = {}): ProjectEntry { + return { + id: "proj_abc", + name: "my-repo", + slug: "my-repo", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://artifacts.example.com/repos/my-repo", + createdAt: "2026-01-01T00:00:00.000Z", + ...overrides, + }; +} + +const WORKSPACE = { + name: "ws-1", + remote: "https://artifacts.example.com/repos/ws-1", + parent: "proj_abc", + createdAt: "2026-01-01T01:00:00.000Z", +}; + +const CHANGE: Change = { + id: "chg_abc123", + project: "@alice/my-repo", + projectId: "proj_abc", + workspace: "ws-1", + status: "open", + createdAt: "2026-01-01T02:00:00.000Z", +}; + +/** Only the bindings the paths under test touch; D1 is never really queried. */ +function makeEnv(): Env { + return { + ARTIFACTS: {} as Env["ARTIFACTS"], + STATE: {} as KVNamespace, + DB: {} as D1Database, + AI: {} as Env["AI"], + }; +} + +/** The context the secret scan — present in every evaluator set — was given. */ +function secretScanContext(): EvaluationContext | undefined { + const call = captured.calls.find((c) => c.evaluator === "secret_scan"); + expect(call).toBeDefined(); + return call?.context; +} + +function stubDiff() { + vi.mocked(getDiffBetweenRepos).mockResolvedValue({ + success: true, + data: { + diff: "diff --git a/a.ts b/a.ts\n+line", + workspaceOid: "ws_tip", + workspaceTreeOid: "ws_tree", + workspaceSha: "ws_tip", + baseOid: "base_from_diff_clone", + }, + }); +} + +beforeEach(() => { + vi.clearAllMocks(); + captured.calls.length = 0; + vi.mocked(loadPolicy).mockResolvedValue({ + evaluators: [{ type: "llm" }], + requireAll: true, + minScore: 0.7, + }); + vi.mocked(getCommitLog).mockResolvedValue({ + success: true, + data: [{ sha: "head_sha", message: "m", author: "a", timestamp: 0 }], + }); + stubDiff(); + vi.mocked(getUserByToken).mockResolvedValue({ + success: true, + data: { + id: "user_alice", + email: "alice@example.com", + username: "alice", + tokenHash: "hash", + createdAt: "2026-01-01T00:00:00.000Z", + }, + }); +}); + +// --------------------------------------------------------------------------- +// 1.1 — the billing subject derived from a ProjectEntry +// --------------------------------------------------------------------------- + +describe("billingContextFor", () => { + it("names a user-owned project's owner as the billing subject", () => { + expect(billingContextFor(projectEntry())).toEqual({ + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + }); + }); + + it("names an org-owned project's org as the billing subject", () => { + const project = projectEntry({ ownerId: "org_acme", ownerType: "org" }); + expect(billingContextFor(project)).toEqual({ + ownerId: "org_acme", + ownerType: "org", + projectId: "proj_abc", + }); + }); + + it("yields no subject for an agent-owned project rather than billing the agent", () => { + // An agent is not a payer — it resolves to its owning user, which is + // deliberately not done here. Attributing spend to an agent id would be + // worse than attributing none, so the field is absent, not guessed. + const project = projectEntry({ ownerId: "agent_bot", ownerType: "agent" }); + expect(billingContextFor(project)).toBeUndefined(); + }); +}); + +// --------------------------------------------------------------------------- +// 1.2 — buildEvaluators takes a ProjectEntry and behaves exactly as before +// --------------------------------------------------------------------------- + +describe("buildEvaluators — ProjectEntry replaces the display name", () => { + const project = projectEntry(); + + function evaluatorFor(policy: EvalPolicy, type: string, env: Env = makeEnv()) { + const entry = buildEvaluators(env, policy, project, logger).find((e) => e.type === type); + expect(entry).toBeDefined(); + return entry?.evaluator; + } + + it("always prepends the blocking secret scan", () => { + const built = buildEvaluators(makeEnv(), { evaluators: [] }, project, logger); + expect(built.map((e) => e.type)).toEqual(["secret_scan"]); + expect(built[0]?.evaluator).toBeInstanceOf(SecretScanEvaluator); + }); + + it("builds a real DiffEvaluator for a diff entry", () => { + expect(evaluatorFor({ evaluators: [{ type: "diff" }] }, "diff")).toBeInstanceOf(DiffEvaluator); + }); + + it("builds a real WebhookEvaluator for a webhook entry", () => { + const policy: EvalPolicy = { evaluators: [{ type: "webhook", url: "https://hook.test" }] }; + expect(evaluatorFor(policy, "webhook")).toBeInstanceOf(WebhookEvaluator); + }); + + it("builds a real LLMEvaluator when the AI binding is present", () => { + expect(evaluatorFor({ evaluators: [{ type: "llm" }] }, "llm")).toBeInstanceOf(LLMEvaluator); + }); + + it("still fails the llm evaluator closed when the AI binding is missing", async () => { + const policy: EvalPolicy = { evaluators: [{ type: "llm" }] }; + const evaluator = evaluatorFor(policy, "llm", {} as Env); + expect(evaluator).toBeInstanceOf(UnavailableEvaluator); + const result = await evaluator?.evaluate("", policy, logger); + expect(result?.success).toBe(true); + if (!result?.success) return; + expect(result.data.passed).toBe(false); + expect(result.data.score).toBe(0); + expect(result.data.reason).toContain("AI binding is not configured"); + }); + + it("still fails the sandbox evaluator closed when the binding is missing", async () => { + const policy: EvalPolicy = { evaluators: [{ type: "sandbox" }] }; + const evaluator = evaluatorFor(policy, "sandbox"); + expect(evaluator).toBeInstanceOf(UnavailableEvaluator); + const result = await evaluator?.evaluate("", policy, logger); + expect(result?.success).toBe(true); + if (!result?.success) return; + expect(result.data.reason).toContain("SANDBOX binding is not configured"); + }); + + it("drops an unknown evaluator type and names the project in the warning", () => { + const warn = vi.fn(); + const warningLogger = { ...logger, warn } as unknown as Logger; + const policy = { evaluators: [{ type: "quantum" }] } as unknown as EvalPolicy; + + const built = buildEvaluators(makeEnv(), policy, project, warningLogger); + + expect(built.map((e) => e.type)).toEqual(["secret_scan"]); + expect(warn).toHaveBeenCalledTimes(1); + // The display name is derived from the entry now; the log line that used to + // receive it as a parameter must still identify the project. + expect(warn.mock.calls[0]?.[0]).toContain("@alice/my-repo"); + expect(warn.mock.calls[0]?.[1]).toMatchObject({ projectName: "@alice/my-repo" }); + }); +}); + +// --------------------------------------------------------------------------- +// runEvaluation forwards the whole context, billing included +// --------------------------------------------------------------------------- + +describe("runEvaluation", () => { + it("forwards the billing context to every evaluator", async () => { + const seen: Array = []; + const spy: Evaluator = { + evaluate: async (_diff, _policy, _logger, context) => { + seen.push(context); + return { success: true, data: { score: 1, passed: true, reason: "ok" } }; + }, + }; + const policy: EvalPolicy = { evaluators: [], requireAll: true, minScore: 0.7 }; + + await runEvaluation( + [ + { type: "secret_scan", evaluator: spy }, + { type: "diff", evaluator: spy }, + ], + "diff", + policy, + logger, + { baseSha: "base", billing: billingContextFor(projectEntry()) }, + ); + + expect(seen).toHaveLength(2); + for (const context of seen) { + expect(context?.baseSha).toBe("base"); + expect(context?.billing).toEqual({ + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + }); + } + }); +}); + +// --------------------------------------------------------------------------- +// 1.3 — call site 1: createChangeWithEvaluation (src/services/change-flow.ts) +// --------------------------------------------------------------------------- + +describe("createChangeWithEvaluation passes the project through", () => { + beforeEach(() => { + vi.mocked(getProject).mockResolvedValue({ success: true, data: projectEntry() }); + }); + + async function create(project: ProjectEntry) { + const { createChange } = await import("../src/storage/changes"); + vi.mocked(createChange).mockResolvedValue({ success: true, data: CHANGE }); + return createChangeWithEvaluation(makeEnv(), logger, { + project, + projectName: "@alice/my-repo", + workspaceName: WORKSPACE.name, + workspaceRemote: WORKSPACE.remote, + actor: { userId: "user_alice" }, + }); + } + + it("bills a user-owned project's owner", async () => { + const outcome = await create(projectEntry()); + expect(outcome.success).toBe(true); + expect(secretScanContext()).toEqual({ + baseSha: "base_from_diff_clone", + billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + }); + }); + + it("bills an org-owned project's org", async () => { + await create(projectEntry({ ownerId: "org_acme", ownerType: "org" })); + expect(secretScanContext()?.billing).toEqual({ + ownerId: "org_acme", + ownerType: "org", + projectId: "proj_abc", + }); + }); + + it("carries no billing subject for an agent-owned project", async () => { + await create(projectEntry({ ownerId: "agent_bot", ownerType: "agent" })); + const context = secretScanContext(); + // The base is still pinned: dropping the billing subject must not drop the + // rest of the context the evaluators already relied on (#274). + expect(context?.baseSha).toBe("base_from_diff_clone"); + expect(context?.billing).toBeUndefined(); + }); + + it("still reaches the LLM evaluator, which is the metered one", async () => { + await create(projectEntry()); + const llm = captured.calls.find((c) => c.evaluator === "llm"); + expect(llm?.context?.billing).toEqual({ + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + }); + }); +}); + +// --------------------------------------------------------------------------- +// 1.3 — call site 2: POST /api/changes/:id/evaluate (src/routes/changes.ts) +// --------------------------------------------------------------------------- + +describe("POST /api/changes/:id/evaluate passes the project through", () => { + function makeApp() { + const app = new Hono<{ Bindings: Env }>(); + app.use("*", authMiddleware); + app.route("/api", changesRouter); + return app; + } + + it("bills the project's owner, not the change's project name string", async () => { + vi.mocked(getChange).mockResolvedValue({ success: true, data: CHANGE }); + vi.mocked(getProject).mockResolvedValue({ success: true, data: projectEntry() }); + vi.mocked(getWorkspace).mockResolvedValue({ success: true, data: WORKSPACE }); + + const res = await makeApp().fetch( + new Request("http://localhost/api/changes/chg_abc123/evaluate", { + method: "POST", + headers: { ...USER_AUTH, "Content-Type": "application/json" }, + body: "{}", + }), + makeEnv(), + ); + + expect(res.status).toBe(200); + // This route used to pass `change.project` — a string that names no owner + // and no project id. The evaluators now see the loaded entry. + expect(secretScanContext()).toEqual({ + baseSha: "base_from_diff_clone", + billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + }); + }); +}); + +// --------------------------------------------------------------------------- +// 1.4 — call site 3: POST /api/projects/conflicts/:id/resolve +// --------------------------------------------------------------------------- + +describe("POST /api/projects/conflicts/:id/resolve passes the project through", () => { + const CONFLICT_CTX = { + conflictId: "conflict-abc", + namespace: "@alice", + slug: "my-repo", + workspaceName: "ws-1", + conflictingFiles: ["src/a.ts"], + detectedAt: "2026-01-01T00:00:00.000Z", + }; + + function makeApp() { + const app = new Hono<{ Bindings: Env }>(); + app.use("*", authMiddleware); + app.route("/api", syncManagementRouter); + return app; + } + + function makeKv(): KVNamespace { + const store: Record = { + [`conflict:${CONFLICT_CTX.conflictId}`]: JSON.stringify(CONFLICT_CTX), + }; + return { + get: vi.fn(async (key: string) => store[key] ?? null), + put: vi.fn(async () => undefined), + delete: vi.fn(async () => undefined), + } as unknown as KVNamespace; + } + + it("bills the loaded project rather than the synthesized @ns/slug string", async () => { + vi.mocked(getProjectByPath).mockResolvedValue({ success: true, data: projectEntry() }); + vi.mocked(getWorkspace).mockResolvedValue({ success: true, data: WORKSPACE }); + vi.mocked(buildManualResolutionDiff).mockResolvedValue({ + success: true, + data: { diff: "diff --git a/a.ts b/a.ts\n+line", baseSha: "conflict_base_sha" }, + }); + + const res = await makeApp().fetch( + new Request("http://localhost/api/projects/conflicts/conflict-abc/resolve", { + method: "POST", + headers: { ...USER_AUTH, "Content-Type": "application/json" }, + body: JSON.stringify({ + strategy: "manual", + resolutions: [{ file: "src/a.ts", content: "resolved\n" }], + }), + }), + { ...makeEnv(), STATE: makeKv() }, + ); + + expect(res.status).toBe(200); + expect(secretScanContext()).toEqual({ + baseSha: "conflict_base_sha", + billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + }); + }); +}); From a853ff47c4ef1658a82ebe560d4c324707553194 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 00:49:28 +0000 Subject: [PATCH 04/26] Fix review findings on the evaluator billing context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three parallel reviews of the Task 1 refactor. No bugs in the threading itself, but one regression it introduced and one gap it left. buildEvaluators interpolated the display name unguarded, so a legacy entry with no namespace logged "undefined/old-project" where the re-evaluate route previously passed a real name. This file already guards that exact shape twice, at resolveProjectHead and ensureWorkspaceRepo; now so does this. It was a regression in the single observable behavior the refactor claimed to change. billingContextFor checked only for agent ownership and passed every other field straight through. KV entries are cast without shape validation, legacy rows genuinely lack fields, and a restored backup re-writes whatever JSON it held — so the type's promise is not a runtime guarantee, and the function could return a context keyed on "". Spend that looks attributed but aggregates under an empty key nobody will reconcile is worse than spend visibly attributed to nobody, so both cases now yield undefined. Two properties the reviews found true but unenforced are now pinned by tests. The webhook evaluator POSTs to a URL taken from the policy file — repository content — and billing.ownerId stays out of that request only because the body names its fields instead of spreading the context; a test asserts the subject never appears in it, and the type says why. The missing-field cases get their own case too. Doc comments corrected where they described the world the refactor abolished, claimed callers no longer build a name they still build, or credited buildEvaluators with handing evaluators a payer it never hands them — the subject travels on the context. runEvaluation now records that one context object is shared across evaluators running under Promise.all, so it must stay immutable data. Two test names claimed to build "a real" evaluator while asserting against this file's own recording stand-ins. Typecheck clean; lint clean; 3736 tests pass, up 3. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- src/evaluation/types.ts | 32 ++++++++----- src/services/change-flow.ts | 47 +++++++++++++------ tests/evaluator-billing-context.test.ts | 61 +++++++++++++++++++++---- 3 files changed, 105 insertions(+), 35 deletions(-) diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index dac6d31d..f85a878d 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -14,30 +14,38 @@ export interface EvalResult { /** * Who pays for the metered resources an evaluation consumes. * - * The evaluator layer has no project today — `buildEvaluators` was given a - * display name and `runEvaluation` nothing at all — so the LLM evaluator's - * spend lands on the operator's Workers AI account with nothing recording - * whose change caused it. Carrying the subject on the evaluation context is - * what makes that spend attributable and, later, meterable, without every - * evaluator having to grow a project-shaped constructor argument. + * The LLM evaluator spends the operator's Workers AI budget on every call, and + * before this existed nothing at the evaluator layer could say whose change + * caused it. Carrying the subject on the evaluation context is what makes that + * spend attributable, and later meterable, without every evaluator having to + * grow a project-shaped constructor argument. * * `ownerType` is narrower than `ProjectEntry.ownerType`, which also admits * `"agent"`: an agent is not a billing subject, it resolves to the user that - * owns it. That resolution is deliberately not done here. + * owns it. `billingContextFor` yields nothing rather than performing that + * resolution. */ export interface BillingContext { /** The paying user or org. Never an agent id — see `ownerType`. */ ownerId: string; ownerType: "user" | "org"; - /** The project whose policy is being enforced. Keys per-project credentials. */ + /** The project whose policy is being enforced. */ projectId: string; } /** - * What the diff is a diff *of*. A diff alone does not identify the tree it - * applies to, so an evaluator that reproduces the change out-of-process (the - * webhook evaluator) cannot tell which base it should apply the hunks to - * (#274). + * What an evaluation is *of*, and who pays for it: the tree the diff applies + * to, plus the billing subject. A diff alone does not identify the tree, so an + * evaluator that reproduces the change out-of-process (the webhook evaluator) + * cannot tell which base it should apply the hunks to (#274). + * + * **Nothing on this type may be forwarded off-box wholesale.** The webhook + * evaluator POSTs to a URL taken from `.stratum/policy.yaml` — repository + * content — and it stays safe only because it names the fields it sends + * (`diff`, sanitized policy, `baseSha`) instead of spreading this object. A + * `...context` there would ship `billing.ownerId` to an arbitrary endpoint the + * policy file chose. Add a field here and that exclusion is one careless + * refactor away; `tests/evaluator-billing-context.test.ts` pins it. */ export interface EvaluationContext { /** diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index f674414c..f3231813 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -104,16 +104,28 @@ export class UnavailableEvaluator implements Evaluator { /** * The billing subject an evaluation of `project` should be charged to, or - * `undefined` when the project has none that can be named yet. + * `undefined` when the project has none that can be named. + * + * Two cases yield nothing, and both are deliberate. * * `ProjectEntry.ownerType` admits `"agent"`, which `BillingContext` does not: * an agent is not a payer, it belongs to one. Walking `agents.owner_id` to that * user is `resolveBillingSubject`'s job in the attribution work that follows — * so an agent-owned project yields no context here rather than a guessed one. * Attributing spend to an agent id would be worse than attributing none. + * + * The field check is the same principle applied to a shape the type system + * promises but KV cannot: entries are cast without validation + * (`src/storage/state.ts`), legacy rows genuinely lack `namespace` and can lack + * `ownerId` — `canWriteProject` guards that exact case — and a restored backup + * re-writes whatever JSON it held. Returning a context keyed on `""` or + * `undefined` would be worse than returning none: spend that *looks* attributed + * aggregates under an empty key that no account will ever reconcile, where an + * absent subject is at least visibly absent. */ export function billingContextFor(project: ProjectEntry): BillingContext | undefined { - if (project.ownerType === "agent") return undefined; + if (project.ownerType !== "user" && project.ownerType !== "org") return undefined; + if (!project.ownerId || !project.id) return undefined; return { ownerId: project.ownerId, ownerType: project.ownerType, @@ -126,9 +138,11 @@ export function billingContextFor(project: ProjectEntry): BillingContext | undef * whatever the policy configures. Evaluators whose binding is missing become * UnavailableEvaluator (score 0, fail) rather than silently vanishing. * - * Takes the whole `ProjectEntry` rather than a display name so an evaluator can - * be told who owns the project it is judging; the display name every call site - * used to pass is derived here so no caller has to build it twice. + * Takes the whole `ProjectEntry` rather than a display name because the payer + * has to be nameable from it. Nothing built here receives the project: the + * billing subject reaches evaluators on the `EvaluationContext` that + * `runEvaluation` forwards, and inside this function the entry is used for the + * log line below and nothing else. * * `workspaceRepo` is read access to the workspace being evaluated (remote + * read token + the pinned evaluated commit); the sandbox evaluator needs it to @@ -144,10 +158,12 @@ export function buildEvaluators( logger: Logger, workspaceRepo?: SandboxRepoAccess, ): Array<{ type: string; evaluator: Evaluator }> { - // The whole entry, not a name: an evaluator that spends money has to be able - // to name the payer, and a "@ns/slug" string identifies no owner and no - // project id. Only the log line below needs the human-readable form. - const projectName = `${project.namespace}/${project.slug}`; + // Guarded because a legacy entry can carry no namespace, as this file already + // assumes at `resolveProjectHead` and `ensureWorkspaceRepo`. Interpolating + // unguarded would log "undefined/old-project" for exactly the projects whose + // identity is already hardest to trace. + const projectName = + project.namespace && project.slug ? `${project.namespace}/${project.slug}` : project.name; const evaluators: Array<{ type: string; evaluator: Evaluator }> = [ { type: "secret_scan", evaluator: new SecretScanEvaluator() }, ]; @@ -224,10 +240,15 @@ export async function runEvaluation( diff: string, policy: EvalPolicy, logger: Logger, - /** What the diff is a diff of. Forwarded to every evaluator so one that - * reproduces the change out-of-process can pin the base it applies to (#274). - * Callers pass the base resolved from the clone the diff came from, never a - * fresh read of the project head. */ + /** What the diff is a diff of, and who pays for evaluating it. Forwarded to + * every evaluator so one that reproduces the change out-of-process can pin + * the base it applies to (#274). Callers pass the base resolved from the + * clone the diff came from, never a fresh read of the project head, and the + * billing subject from `billingContextFor`. + * + * One object is shared by every evaluator in the `Promise.all` below, so it + * must stay immutable data. Hanging a mutable balance or counter off it would + * be a race across concurrently running evaluators, not a shared budget. */ context?: EvaluationContext, ): Promise<{ evalRuns: EvaluationRun[]; evalResult: EvalResult }> { const evalRuns = await Promise.all( diff --git a/tests/evaluator-billing-context.test.ts b/tests/evaluator-billing-context.test.ts index 5604ce25..7f5c1b48 100644 --- a/tests/evaluator-billing-context.test.ts +++ b/tests/evaluator-billing-context.test.ts @@ -4,6 +4,7 @@ import type { EvalPolicy, EvaluationContext, Evaluator } from "../src/evaluation import { authMiddleware } from "../src/middleware/auth"; import type { Change, Env, ProjectEntry } from "../src/types"; import type { Logger } from "../src/utils/logger"; +import { makeFakeKV } from "./helpers/fake-kv"; /** * Task 1 is a pure refactor whose only observable effect is what reaches the @@ -251,6 +252,17 @@ describe("billingContextFor", () => { }); }); + it.each([ + ["a missing ownerId", { ownerId: "" }], + ["a missing project id", { id: "" }], + ])("yields no subject for %s rather than one keyed on an empty string", (_label, overrides) => { + // KV entries are cast without shape validation and legacy rows genuinely + // lack fields, so the type's promise is not a runtime guarantee. Spend that + // *looks* attributed but aggregates under "" is worse than spend visibly + // attributed to nobody. + expect(billingContextFor(projectEntry(overrides))).toBeUndefined(); + }); + it("yields no subject for an agent-owned project rather than billing the agent", () => { // An agent is not a payer — it resolves to its owning user, which is // deliberately not done here. Attributing spend to an agent id would be @@ -273,6 +285,10 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { return entry?.evaluator; } + // `SecretScanEvaluator` and `LLMEvaluator` are the recording stand-ins this + // file substitutes above, not the shipped classes. The assertions still + // separate the wired branch from `UnavailableEvaluator`, which is what they + // are for — they just don't prove the real class was constructed. it("always prepends the blocking secret scan", () => { const built = buildEvaluators(makeEnv(), { evaluators: [] }, project, logger); expect(built.map((e) => e.type)).toEqual(["secret_scan"]); @@ -288,8 +304,10 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { expect(evaluatorFor(policy, "webhook")).toBeInstanceOf(WebhookEvaluator); }); - it("builds a real LLMEvaluator when the AI binding is present", () => { - expect(evaluatorFor({ evaluators: [{ type: "llm" }] }, "llm")).toBeInstanceOf(LLMEvaluator); + it("wires the llm evaluator, not UnavailableEvaluator, when the AI binding is present", () => { + const evaluator = evaluatorFor({ evaluators: [{ type: "llm" }] }, "llm"); + expect(evaluator).toBeInstanceOf(LLMEvaluator); + expect(evaluator).not.toBeInstanceOf(UnavailableEvaluator); }); it("still fails the llm evaluator closed when the AI binding is missing", async () => { @@ -334,6 +352,34 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { // runEvaluation forwards the whole context, billing included // --------------------------------------------------------------------------- +describe("the billing subject stays inside the Worker", () => { + it("is not forwarded to the policy-supplied webhook URL", async () => { + // The webhook evaluator POSTs to a URL taken from .stratum/policy.yaml — + // repository content. It is safe only because it names the fields it sends + // rather than spreading the context. This pins that: a `...context` there + // would ship ownerId and projectId to an endpoint the policy file chose. + const fetchMock = vi.fn(async () => new Response(JSON.stringify({ score: 1, passed: true }))); + const original = globalThis.fetch; + globalThis.fetch = fetchMock as unknown as typeof fetch; + try { + const policy: EvalPolicy = { evaluators: [{ type: "webhook", url: "https://hook.test" }] }; + await new WebhookEvaluator().evaluate("a diff", policy, logger, { + baseSha: "base", + billing: billingContextFor(projectEntry()), + }); + + expect(fetchMock).toHaveBeenCalledTimes(1); + const body = String((fetchMock.mock.calls[0] as unknown as [string, RequestInit])[1].body); + expect(body).toContain("baseSha"); + expect(body).not.toContain("billing"); + expect(body).not.toContain("user_alice"); + expect(body).not.toContain("proj_abc"); + } finally { + globalThis.fetch = original; + } + }); +}); + describe("runEvaluation", () => { it("forwards the billing context to every evaluator", async () => { const seen: Array = []; @@ -485,14 +531,9 @@ describe("POST /api/projects/conflicts/:id/resolve passes the project through", } function makeKv(): KVNamespace { - const store: Record = { - [`conflict:${CONFLICT_CTX.conflictId}`]: JSON.stringify(CONFLICT_CTX), - }; - return { - get: vi.fn(async (key: string) => store[key] ?? null), - put: vi.fn(async () => undefined), - delete: vi.fn(async () => undefined), - } as unknown as KVNamespace; + const kv = makeFakeKV(); + kv.store.set(`conflict:${CONFLICT_CTX.conflictId}`, JSON.stringify(CONFLICT_CTX)); + return kv; } it("bills the loaded project rather than the synthesized @ns/slug string", async () => { From 359f6946313536a9c7246b5800e39b6655dbbca6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 01:08:14 +0000 Subject: [PATCH 05/26] Name who pays for each cost sample MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 2 of the metering work. cost_records has recorded what a change consumed since it existed but never who owes for it, so migration 048 adds owner_id, owner_type and source, and an owner-scoped index. Owner columns are nullable because the mapping needed to backfill them lives in KV, not D1 — the same reason 025 backfilled no project_id — and discarding rows that cannot be attributed would delete the only record of spend that happened. A NULL owner is a row nobody can be billed for, which is honest; an owner guessed from a project name would be a row billed to possibly the wrong account. source is NOT NULL DEFAULT 'platform' because the backfill is correct by construction: BYOK does not exist yet, so every existing row was spend on the operator's own provider account. resolveBillingSubject is the single home for the owner-to-payer mapping. An agent is not a payer, so it walks agents.owner_id to the owning user. It never throws: recordCosts is best-effort by contract and runs inside change creation, merge and deploy, so an unresolvable owner costs the attribution, never the change. Its owner type is spelled out rather than reused from ProjectEntry so a fourth owner kind fails to compile and someone has to decide who pays. The conflict-resolution path now records its spend, closing the one metered path that recorded nothing. It records before the verdict is acted on, because a resolution the suite rejects burned the same model tokens as one it accepts and the early return would drop them. Adds projectDisplayName, and routes three call sites through it. The new recording site interpolated namespace and slug unguarded, which for a legacy entry carrying no namespace writes the literal "undefined/slug" — tolerable in the log line where this pattern was fixed before, wrong in cost_records.project, where it is a persisted identifier that matches no project during the deletion cascade's name-form lookup. Typecheck clean; 3764 tests pass, up 28. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- migrations/048_cost_attribution.sql | 61 +++ src/deploy/runner.ts | 4 +- src/evaluation/types.ts | 18 +- src/merge/post-merge.ts | 10 +- src/routes/changes.ts | 13 +- src/routes/sync-management.ts | 27 +- src/services/change-flow.ts | 24 +- src/storage/costs.ts | 98 ++++- src/types.ts | 19 + tests/cost-attribution-paths.test.ts | 470 ++++++++++++++++++++++++ tests/cost-attribution.test.ts | 284 ++++++++++++++ tests/deploy-runner.test.ts | 32 ++ tests/evaluator-billing-context.test.ts | 9 +- tests/post-merge.test.ts | 79 ++++ 14 files changed, 1128 insertions(+), 20 deletions(-) create mode 100644 migrations/048_cost_attribution.sql create mode 100644 tests/cost-attribution-paths.test.ts create mode 100644 tests/cost-attribution.test.ts diff --git a/migrations/048_cost_attribution.sql b/migrations/048_cost_attribution.sql new file mode 100644 index 00000000..91b533b6 --- /dev/null +++ b/migrations/048_cost_attribution.sql @@ -0,0 +1,61 @@ +-- Migration 048: name who pays for each cost sample. +-- +-- `cost_records` (021, given `project_id` by 025) has recorded what a change +-- consumed since it existed, but never who owes for it. A project name is not a +-- billing subject: names are unique only per-namespace, projects are transferred +-- and deleted, and an aggregate keyed on one cannot be reconciled against an +-- account. These three columns make every new row attributable at the moment +-- the spend is incurred, which is the only moment the payer is known for +-- certain. +-- +-- `owner_id`/`owner_type` are NULLABLE, deliberately. Every historical row +-- predates attribution and the mapping needed to backfill it lives in KV (the +-- ProjectEntry's owner), not in D1 — the same reason 025 backfilled no +-- `project_id`. SQL cannot resolve it here, and discarding rows that cannot be +-- attributed would delete the only record of spend that actually happened. So +-- unattributed rows stay, visibly unattributed: a NULL owner is a row nobody +-- can be billed for, which is honest, where an owner guessed from a project +-- name would be a row billed to possibly the wrong account. Live code can also +-- write NULL — `resolveBillingSubject` (src/storage/costs.ts) returns null +-- rather than throwing when a project's owner cannot be named — so this is not +-- a transitional state that later becomes NOT NULL. +-- +-- `source` is NOT NULL DEFAULT 'platform' because the backfill is correct by +-- construction, not merely convenient: BYOK does not exist yet, so every row +-- already in this table was spend on the operator's own provider account. The +-- default backfills them in place with the truth. Precedent for both shapes of +-- ALTER: a CHECK on an added column in 037, NOT NULL DEFAULT in 032/034/041. +-- No table rebuild — unlike `kind`, whose closed CHECK would need one to widen. +-- +-- The new index sorts by `created_at`, which is safe HERE specifically. +-- `cost_records.created_at` is written exclusively by `recordCosts` +-- (src/storage/costs.ts) as `new Date().toISOString()` — that function holds the +-- only INSERT into this table anywhere in the tree — so every value is ISO 8601 +-- with a 'T' separator. The hazard 042/044/047 warn about is a column mixing +-- that format with `datetime('now')`'s space separator, which sorts differently +-- as TEXT (' ' 0x20 < 'T' 0x54) and silently breaks range and ORDER BY +-- comparisons. A single-format column has no such problem, so an owner's usage +-- history can be paged newest-first in SQL rather than in JavaScript. +-- +-- ROLLBACK: drop `idx_costs_owner`; leave the three columns in place. SQLite's +-- DROP COLUMN would rewrite the table for no gain, and the columns are inert +-- when unwritten — reads name them explicitly, an absent owner is already the +-- normal historical case, and `source` defaults itself. Code must therefore +-- tolerate the columns being present but never populated. No migration in this +-- repo has been reverted; this records the intent rather than building +-- machinery for it. + +-- CHECK'd rather than free text: an owner_type outside this set is unreachable +-- from `resolveBillingSubject`, which maps ProjectEntry's third kind ('agent') +-- onto the user that owns the agent instead of passing it through. A row +-- claiming a type no billing subject can have would aggregate under an account +-- that will never be invoiced. +ALTER TABLE cost_records ADD COLUMN owner_id TEXT; +ALTER TABLE cost_records ADD COLUMN owner_type TEXT CHECK(owner_type IN ('user','org')); +ALTER TABLE cost_records ADD COLUMN source TEXT NOT NULL DEFAULT 'platform' + CHECK(source IN ('platform','byok')); + +-- Covers the owner-scoped read this exists for: `WHERE owner_id = ? ORDER BY +-- created_at DESC`. Rows with a NULL owner are still indexed (SQLite indexes +-- NULLs) but are never selected by that predicate, which is the intent. +CREATE INDEX IF NOT EXISTS idx_costs_owner ON cost_records(owner_id, created_at DESC); diff --git a/src/deploy/runner.ts b/src/deploy/runner.ts index df4dcf64..a4d136e1 100644 --- a/src/deploy/runner.ts +++ b/src/deploy/runner.ts @@ -33,7 +33,7 @@ import { parsePolicyContent } from "../evaluation/policy-loader"; import type { DeployConfig, DeployRejection } from "../evaluation/types"; import { type StratumEvent, emitEvent } from "../queue/events"; import { getChange } from "../storage/changes"; -import { recordCosts } from "../storage/costs"; +import { recordCosts, resolveBillingSubject } from "../storage/costs"; import { isTargetDeleting } from "../storage/deletion"; import { type Deployment, @@ -854,6 +854,7 @@ async function readTree( // Recorded whether or not the read succeeded: a clone that failed still cost // the round trip. + const subject = await resolveBillingSubject(env.DB, logger, project); await recordCosts( env.DB, logger, @@ -861,6 +862,7 @@ async function readTree( project: project.name, projectId: project.id, ...(changeId !== undefined ? { changeId } : {}), + ...(subject ?? {}), }, [{ kind: "git_ops", quantity: 1 }], ); diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index f85a878d..ad4c733b 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -7,8 +7,22 @@ export interface EvalResult { passed: boolean; reason: string; issues?: string[]; - /** Resource usage the evaluator incurred, recorded for cost tracking. */ - costs?: Array<{ kind: "llm_tokens" | "sandbox_ms"; quantity: number; estimated?: boolean }>; + /** + * Resource usage the evaluator incurred, recorded for cost tracking. + * + * Structurally a subset of `CostSample` (src/storage/costs.ts) rather than an + * import of it, so the evaluation layer does not depend on storage. `source` + * is what lets an evaluator running on the project's own provider credential + * say so: both recording sites flatten this array straight into `recordCosts`, + * so a field missing here is a distinction that cannot be made at all. + * Omitted means `"platform"` — the operator paid. + */ + costs?: Array<{ + kind: "llm_tokens" | "sandbox_ms"; + quantity: number; + estimated?: boolean; + source?: "platform" | "byok"; + }>; } /** diff --git a/src/merge/post-merge.ts b/src/merge/post-merge.ts index 21adc24f..82d5182f 100644 --- a/src/merge/post-merge.ts +++ b/src/merge/post-merge.ts @@ -2,7 +2,7 @@ import { materializeTree } from "../evaluation/sandbox-evaluator"; import type { EvalPolicy } from "../evaluation/types"; import { emitEvent } from "../queue/events"; import { updateChangeStatus } from "../storage/changes"; -import { recordCosts } from "../storage/costs"; +import { recordCosts, resolveBillingSubject } from "../storage/costs"; import { freshRepoToken, getCommitParent, readRepoFiles, revertToCommit } from "../storage/git-ops"; import type { Env, ProjectEntry } from "../types"; import { projectDefaultBranch } from "../types"; @@ -85,10 +85,16 @@ export async function runPostMergeCheck( const run = await sandbox.run(command, { timeout: merge?.postMergeTimeoutMs ?? DEFAULT_POST_MERGE_TIMEOUT_MS, }); + const subject = await resolveBillingSubject(env.DB, logger, project); await recordCosts( env.DB, logger, - { project: project.name, projectId: project.id, changeId: opts.changeId }, + { + project: project.name, + projectId: project.id, + changeId: opts.changeId, + ...(subject ?? {}), + }, [ { kind: "sandbox_ms", quantity: Date.now() - runStartedAt }, { kind: "git_ops", quantity: 1 }, diff --git a/src/routes/changes.ts b/src/routes/changes.ts index 8b7a566a..33e80157 100644 --- a/src/routes/changes.ts +++ b/src/routes/changes.ts @@ -25,7 +25,12 @@ import { mergeTransitionOpts, updateChangeStatus, } from "../storage/changes"; -import { type CostSample, getChangeCostSummary, recordCosts } from "../storage/costs"; +import { + type CostSample, + getChangeCostSummary, + recordCosts, + resolveBillingSubject, +} from "../storage/costs"; import { isTargetDeleting } from "../storage/deletion"; import { listEvalRuns, recordEvalRuns } from "../storage/eval-runs"; import { @@ -801,6 +806,7 @@ app.post("/changes/:id/merge", async (c) => { }); } + const mergeSubject = await resolveBillingSubject(c.env.DB, logger, project); await recordCosts( c.env.DB, logger, @@ -810,6 +816,9 @@ app.post("/changes/:id/merge", async (c) => { projectId: change.projectId ?? project.id, changeId: id, workspace: change.workspace, + // The project's owner now, not the author of the change: this is the + // account the clone and push were spent against. + ...(mergeSubject ?? {}), }, [{ kind: "git_ops", quantity: 2 }], ); @@ -1421,6 +1430,7 @@ app.post("/changes/:id/evaluate", async (c) => { { kind: "git_ops", quantity: 2 }, ...evalRuns.flatMap(({ result }) => result.costs ?? []), ]; + const evaluateSubject = await resolveBillingSubject(c.env.DB, logger, project); await recordCosts( c.env.DB, logger, @@ -1429,6 +1439,7 @@ app.post("/changes/:id/evaluate", async (c) => { projectId: change.projectId ?? project.id, changeId: id, workspace: change.workspace, + ...(evaluateSubject ?? {}), }, evaluateCostSamples, ); diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index beb7c5e8..89d627ef 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -6,6 +6,7 @@ import { authMiddleware } from "../middleware/auth"; import { billingContextFor, buildEvaluators, runEvaluation } from "../services/change-flow"; import { recordAudit } from "../storage/audit"; import { getChange } from "../storage/changes"; +import { recordCosts, resolveBillingSubject } from "../storage/costs"; import { MAX_FILE_BYTES, buildManualResolutionDiff, @@ -22,7 +23,7 @@ import { updateProjectAfterSync, } from "../storage/sync"; import type { Env } from "../types"; -import { projectDefaultBranch } from "../types"; +import { projectDefaultBranch, projectDisplayName } from "../types"; import { canWriteProject } from "../utils/authz"; import { createLogger } from "../utils/logger"; import { readJsonWithLimit } from "../utils/request-body"; @@ -665,13 +666,33 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { // `buildManualResolutionDiff` resolved this base from the clone it built the // diff on, so it names the tree the resolution actually applies to (#274). // This path runs the LLM evaluator too, so it needs a payer as much as - // change creation does — it is the one metered path that records nothing - // today. + // change creation does. const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha, billing: billingContextFor(project), }); + // Recorded here, BEFORE the verdict is acted on, because the spend is + // already incurred: a resolution the suite rejects burned exactly the same + // model tokens as one it accepts, and the early return below would drop + // them. Mirrors the recording POST /changes/:id/evaluate does, minus the + // git_ops pair — `buildManualResolutionDiff` cloned once, above. + const resolveSubject = await resolveBillingSubject(c.env.DB, logger, project); + await recordCosts( + c.env.DB, + logger, + { + project: projectDisplayName(project), + projectId: project.id, + // The Change whose merge attempt produced this conflict, when the + // conflict context named one — a resolution is part of landing it. + ...(conflictCtx.changeId ? { changeId: conflictCtx.changeId } : {}), + workspace: conflictCtx.workspaceName, + ...(resolveSubject ?? {}), + }, + [{ kind: "git_ops", quantity: 1 }, ...evalRuns.flatMap(({ result }) => result.costs ?? [])], + ); + if (!evalResult.passed) { logger.warn("Manual conflict resolution blocked by evaluator suite", { conflictId, diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index f3231813..89cd71d9 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -20,7 +20,7 @@ import { buildEvaluationReport, reportEvaluationToGitHub } from "../github/sync" import { type EventActor, emitEvent } from "../queue/events"; import { getAgent } from "../storage/agents"; import { createChange, updateChangeStatus } from "../storage/changes"; -import { type CostSample, recordCosts } from "../storage/costs"; +import { type CostSample, recordCosts, resolveBillingSubject } from "../storage/costs"; import { recordEvalRuns } from "../storage/eval-runs"; import { artifactsRepoNameFromRemote, @@ -36,6 +36,7 @@ import { type ProjectEntry, getArtifactsRepoName, projectDefaultBranch, + projectDisplayName, } from "../types"; import { AppError } from "../utils/errors"; import type { Logger } from "../utils/logger"; @@ -158,12 +159,9 @@ export function buildEvaluators( logger: Logger, workspaceRepo?: SandboxRepoAccess, ): Array<{ type: string; evaluator: Evaluator }> { - // Guarded because a legacy entry can carry no namespace, as this file already - // assumes at `resolveProjectHead` and `ensureWorkspaceRepo`. Interpolating - // unguarded would log "undefined/old-project" for exactly the projects whose - // identity is already hardest to trace. - const projectName = - project.namespace && project.slug ? `${project.namespace}/${project.slug}` : project.name; + // Guarded, because a legacy entry can carry no namespace — see + // `projectDisplayName`, and this file's own `resolveProjectHead`. + const projectName = projectDisplayName(project); const evaluators: Array<{ type: string; evaluator: Evaluator }> = [ { type: "secret_scan", evaluator: new SecretScanEvaluator() }, ]; @@ -475,10 +473,20 @@ export async function createChangeWithEvaluation( { kind: "git_ops", quantity: 2 }, ...evalRuns.flatMap(({ result }) => result.costs ?? []), ]; + // Resolved rather than reused from `billingContextFor` above: that one yields + // nothing for an agent-owned project, while the ledger can still name the + // agent's owning user as the payer. + const createSubject = await resolveBillingSubject(env.DB, logger, project); await recordCosts( env.DB, logger, - { project: projectName, projectId: project.id, changeId: change.id, workspace: workspaceName }, + { + project: projectName, + projectId: project.id, + changeId: change.id, + workspace: workspaceName, + ...(createSubject ?? {}), + }, createCostSamples, ); diff --git a/src/storage/costs.ts b/src/storage/costs.ts index fc3f7ca1..1d33c97d 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -2,14 +2,36 @@ import { AppError } from "../utils/errors"; import { newId } from "../utils/ids"; import type { Logger } from "../utils/logger"; import { type Result, err, ok } from "../utils/result"; +import { getAgent } from "./agents"; export type CostKind = "llm_tokens" | "sandbox_ms" | "git_ops"; +/** + * Whose provider account paid for a sample: the operator's (`platform`) or the + * project's own credential (`byok`). Recorded per sample rather than per change + * because one evaluation can mix the two — a BYOK `llm` evaluator alongside + * git operations that are always the operator's cost. + */ +export type CostSource = "platform" | "byok"; + export interface CostSample { kind: CostKind; quantity: number; /** True when the quantity is an estimate (e.g. character-based token counts). */ estimated?: boolean; + /** Defaults to `"platform"`: an evaluator that says nothing spent our money. */ + source?: CostSource; +} + +/** + * An account a cost row can actually be billed to. + * + * Narrower than `ProjectEntry.ownerType`, which also admits `"agent"` — see + * `resolveBillingSubject`, which is the only thing allowed to make that mapping. + */ +export interface BillingSubject { + ownerId: string; + ownerType: "user" | "org"; } export interface CostSummaryEntry { @@ -24,14 +46,81 @@ interface SummaryRow { any_estimated: number; } +/** + * The account that owes for spend incurred on behalf of `owner`, or `null` when + * no account can be named. + * + * The single home for the `ownerType -> billing subject` mapping: every + * recording site resolves through here so the agent walk and the + * cannot-be-attributed rule exist once rather than five times. Accepts a whole + * `ProjectEntry` (structurally) or a bare owner pair. + * + * An **agent is not a payer** — it belongs to one — so `"agent"` walks + * `agents.owner_id` (migrations/001_core.sql) to the user that owns it. + * `billingContextFor` (src/services/change-flow.ts) deliberately does not do + * this: it is synchronous and KV-only, while this needs D1. No creation path + * writes `ownerType: "agent"` today, so that branch is defensive against + * restored backups and future ownership transfer, not a live code path. + * + * **Never throws.** `recordCosts` is best-effort by contract and runs inside + * change creation, merge and deploy; an unresolvable owner must cost the + * attribution, never the change. Every failure returns `null` and logs — a row + * with a NULL owner is honest about being unattributed, which is why migration + * 048 leaves those columns nullable. + * + * @param owner - A `ProjectEntry`, or any object carrying its owner fields + * @returns The billing subject, or `null` when none can be named + */ +export async function resolveBillingSubject( + db: D1Database, + logger: Logger, + // Spelled out rather than `ProjectEntry["ownerType"]` on purpose: if a fourth + // owner kind is ever added, this must fail to compile so someone decides who + // pays, instead of silently falling through to the agent walk. + owner: { ownerId?: string; ownerType?: "user" | "org" | "agent" }, +): Promise { + const { ownerId, ownerType } = owner; + // KV entries are cast without shape validation and legacy rows genuinely lack + // fields, so the type's promise is not a runtime guarantee. + if (!ownerId || !ownerType) { + logger.warn("Cost attribution skipped: project names no owner", { ownerId, ownerType }); + return null; + } + if (ownerType === "user" || ownerType === "org") return { ownerId, ownerType }; + + const agentResult = await getAgent(db, ownerId, logger); + if (!agentResult.success) { + logger.warn("Cost attribution skipped: agent owner could not be resolved to a user", { + agentId: ownerId, + error: agentResult.error.message, + }); + return null; + } + // No emptiness check on the resolved id, unlike the owner fields above: + // `agents.owner_id` is NOT NULL and REFERENCES users(id), and D1 enforces + // foreign keys, so a row that exists names a user that exists. + return { ownerId: agentResult.data.ownerId, ownerType: "user" }; +} + /** * Record cost samples for a change. Best-effort: failures are logged and * reported, but callers treat cost recording as non-blocking. + * + * `ownerId`/`ownerType` name who pays, and come from `resolveBillingSubject` — + * omitting them records the spend unattributed rather than refusing to record + * it, because a sample nobody can be billed for is still a sample that happened. */ export async function recordCosts( db: D1Database, logger: Logger, - opts: { project: string; projectId?: string; changeId?: string; workspace?: string }, + opts: { + project: string; + projectId?: string; + changeId?: string; + workspace?: string; + ownerId?: string; + ownerType?: BillingSubject["ownerType"]; + }, samples: CostSample[], ): Promise> { if (samples.length === 0) return ok(undefined); @@ -39,7 +128,7 @@ export async function recordCosts( try { const stmt = db.prepare( - "INSERT INTO cost_records (id, project, project_id, change_id, workspace, kind, quantity, estimated, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "INSERT INTO cost_records (id, project, project_id, change_id, workspace, kind, quantity, estimated, created_at, owner_id, owner_type, source) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ); await db.batch( samples.map((sample) => @@ -53,6 +142,11 @@ export async function recordCosts( sample.quantity, sample.estimated ? 1 : 0, createdAt, + opts.ownerId ?? null, + opts.ownerType ?? null, + // Bound explicitly rather than left to the column DEFAULT: the column + // is NOT NULL, so an unset source must resolve here, not in SQLite. + sample.source ?? "platform", ), ), ); diff --git a/src/types.ts b/src/types.ts index e3f4c102..789986fe 100644 --- a/src/types.ts +++ b/src/types.ts @@ -329,6 +329,25 @@ export function projectDefaultBranch(project: { return project.sourceDefaultBranch || project.githubDefaultBranch || "main"; } +/** + * A project's `@namespace/slug` label, falling back to its bare name. + * + * The fallback is the point. A legacy entry can carry no namespace — KV records + * are cast without shape validation, and this file's own `projectDefaultBranch` + * exists for the same class of drift — so interpolating the pair unguarded + * yields the string `"undefined/my-repo"`. That is merely ugly in a log line + * and actively wrong in `cost_records.project`, where it becomes a persisted + * identifier that matches no project during the deletion cascade's name-form + * lookup. + */ +export function projectDisplayName(project: { + namespace?: string; + slug?: string; + name: string; +}): string { + return project.namespace && project.slug ? `${project.namespace}/${project.slug}` : project.name; +} + export interface WorkspaceEntry { name: string; remote: string; diff --git a/tests/cost-attribution-paths.test.ts b/tests/cost-attribution-paths.test.ts new file mode 100644 index 00000000..718b4e25 --- /dev/null +++ b/tests/cost-attribution-paths.test.ts @@ -0,0 +1,470 @@ +/** + * Every path that records a cost row names who pays. + * + * Attribution is only worth having if it is uniform: one unattributed recording + * site is a hole an account can be steered through. So these tests drive the + * real routes and services and read the rows out of a real SQLite D1 with the + * production migrations applied, rather than asserting on a spy's arguments — + * a `recordCosts` mock would pass whether or not migration 048's columns + * accept what the call site sends. + * + * The conflict-resolution path is the reason this file exists at all: it ran + * the full evaluator suite, LLM included, and recorded nothing. + */ +import { Hono } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { EvalPolicy, EvaluationContext } from "../src/evaluation/types"; +import { authMiddleware } from "../src/middleware/auth"; +import type { Change, Env, ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { makeFakeKV } from "./helpers/fake-kv"; +import { makeSqliteD1 } from "./helpers/sqlite-d1"; + +/** Lets a test fail the suite without stopping the evaluators reporting costs. */ +const control = vi.hoisted(() => ({ llmPassed: true })); + +/** + * Stand-in evaluators that report spend, so the flattening of + * `EvalResult.costs` into `recordCosts` is exercised rather than assumed. The + * LLM one claims BYOK: `source` has to survive the flatten, or a project + * paying its own provider bill is recorded as platform spend. + */ +function reportingEvaluatorClass(name: string, costs: unknown, passes = () => true) { + return class { + async evaluate( + _diff: string, + _policy: EvalPolicy, + _logger: Logger, + _context?: EvaluationContext, + ) { + const passed = passes(); + const reason = `${name} ${passed ? "ok" : "blocked"}`; + return { + success: true as const, + data: { score: passed ? 1 : 0, passed, reason, costs }, + }; + } + }; +} + +vi.mock("../src/evaluation/secret-scanner", async (importActual) => ({ + ...(await importActual()), + SecretScanEvaluator: reportingEvaluatorClass("secret_scan", undefined), +})); + +vi.mock("../src/evaluation/llm-evaluator", async (importActual) => ({ + ...(await importActual()), + LLMEvaluator: reportingEvaluatorClass( + "llm", + [{ kind: "llm_tokens", quantity: 1500, estimated: true, source: "byok" }], + () => control.llmPassed, + ), +})); + +vi.mock("../src/evaluation/policy-loader", async (importActual) => ({ + ...(await importActual()), + loadPolicy: vi.fn(), +})); + +vi.mock("../src/storage/git-ops", async (importActual) => ({ + ...(await importActual()), + freshRepoToken: vi.fn(async () => ({ success: true, data: "test-token" })), + getCommitLog: vi.fn(), + getDiffBetweenRepos: vi.fn(), + buildManualResolutionDiff: vi.fn(), + mergeWorkspaceIntoProject: vi.fn(async () => ({ success: true, data: "sha_merged" })), + resolveConflict: vi.fn(async () => ({ success: true, data: { commitSha: "resolved_sha" } })), +})); + +vi.mock("../src/storage/sync", async (importActual) => ({ + ...(await importActual()), + recordSyncHistory: vi.fn(async () => undefined), +})); + +vi.mock("../src/storage/repo-snapshot", () => ({ + readRepoSnapshot: vi.fn(async () => ({ success: true, data: null })), +})); + +vi.mock("../src/storage/state", async (importActual) => ({ + ...(await importActual()), + getProject: vi.fn(), + getProjectByPath: vi.fn(), + getWorkspace: vi.fn(), +})); + +vi.mock("../src/storage/changes", async (importActual) => ({ + ...(await importActual()), + createChange: vi.fn(), + getChange: vi.fn(), + updateChangeStatus: vi.fn(async () => ({ success: true, data: undefined })), + markChangeMerged: vi.fn(async () => ({ success: true, data: { transitioned: true } })), + dismissApprovalsAndUpdateStatus: vi.fn(async () => ({ + success: true, + data: { dismissedReviewerIds: [] }, + })), +})); + +vi.mock("../src/storage/eval-runs", async (importActual) => ({ + ...(await importActual()), + recordEvalRuns: vi.fn(async () => ({ success: true, data: [] })), +})); + +vi.mock("../src/storage/provenance", () => ({ + recordProvenance: vi.fn(async () => ({ success: true, data: undefined })), +})); + +vi.mock("../src/merge/post-merge", () => ({ + runPostMergeCheck: vi.fn(async () => ({ status: "skipped" })), +})); + +vi.mock("../src/storage/deletion", () => ({ + isTargetDeleting: vi.fn(async () => false), +})); + +vi.mock("../src/queue/events", () => ({ + emitEvent: vi.fn(async () => undefined), +})); + +vi.mock("../src/storage/users", () => ({ + getUserByToken: vi.fn(), + getUser: vi.fn(), +})); + +import { loadPolicy } from "../src/evaluation/policy-loader"; +import { changesRouter } from "../src/routes/changes"; +import { syncManagementRouter } from "../src/routes/sync-management"; +import { createChangeWithEvaluation } from "../src/services/change-flow"; +import { createChange, getChange } from "../src/storage/changes"; +import { + buildManualResolutionDiff, + getCommitLog, + getDiffBetweenRepos, +} from "../src/storage/git-ops"; +import { getProject, getProjectByPath, getWorkspace } from "../src/storage/state"; +import { getUserByToken } from "../src/storage/users"; +import { createLogger } from "../src/utils/logger"; + +const logger = createLogger({ component: "test" }); +const USER_AUTH = { Authorization: "Bearer stratum_user_testtoken00000000000000000" }; + +function projectEntry(overrides: Partial = {}): ProjectEntry { + return { + id: "proj_abc", + name: "my-repo", + slug: "my-repo", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://artifacts.example.com/repos/my-repo", + createdAt: "2026-01-01T00:00:00.000Z", + ...overrides, + }; +} + +const WORKSPACE = { + name: "ws-1", + remote: "https://artifacts.example.com/repos/ws-1", + parent: "proj_abc", + createdAt: "2026-01-01T01:00:00.000Z", +}; + +const CHANGE: Change = { + id: "chg_abc123", + project: "@alice/my-repo", + projectId: "proj_abc", + workspace: "ws-1", + status: "accepted", + evalPassed: true, + evalScore: 1, + createdAt: "2026-01-01T02:00:00.000Z", +}; + +interface CostRow { + project: string; + project_id: string | null; + change_id: string | null; + workspace: string | null; + kind: string; + quantity: number; + owner_id: string | null; + owner_type: string | null; + source: string; +} + +let db: D1Database; +let raw: ReturnType["raw"]; + +function costRows(): CostRow[] { + return raw.prepare("SELECT * FROM cost_records ORDER BY kind").all() as unknown as CostRow[]; +} + +function makeEnv(overrides: Partial = {}): Env { + return { + ARTIFACTS: {} as Env["ARTIFACTS"], + STATE: makeFakeKV(), + DB: db, + AI: {} as Env["AI"], + ...overrides, + } as Env; +} + +beforeEach(() => { + const made = makeSqliteD1(); + db = made.db; + raw = made.raw; + vi.clearAllMocks(); + control.llmPassed = true; + + vi.mocked(loadPolicy).mockResolvedValue({ + evaluators: [{ type: "llm" }], + requireAll: true, + minScore: 0.7, + }); + vi.mocked(getCommitLog).mockResolvedValue({ + success: true, + data: [{ sha: "head_sha", message: "m", author: "a", timestamp: 0 }], + }); + vi.mocked(getDiffBetweenRepos).mockResolvedValue({ + success: true, + data: { + diff: "diff --git a/a.ts b/a.ts\n+line", + workspaceOid: "ws_tip", + workspaceTreeOid: "ws_tree", + workspaceSha: "ws_tip", + baseOid: "base_from_diff_clone", + }, + }); + vi.mocked(createChange).mockResolvedValue({ success: true, data: CHANGE }); + vi.mocked(getChange).mockResolvedValue({ success: true, data: CHANGE }); + vi.mocked(getProject).mockResolvedValue({ success: true, data: projectEntry() }); + vi.mocked(getProjectByPath).mockResolvedValue({ success: true, data: projectEntry() }); + vi.mocked(getWorkspace).mockResolvedValue({ success: true, data: WORKSPACE }); + vi.mocked(getUserByToken).mockResolvedValue({ + success: true, + data: { + id: "user_alice", + email: "alice@example.com", + username: "alice", + tokenHash: "hash", + createdAt: "2026-01-01T00:00:00.000Z", + }, + }); +}); + +function makeApp(router: Hono<{ Bindings: Env }>) { + const app = new Hono<{ Bindings: Env }>(); + app.use("*", authMiddleware); + app.route("/api", router); + return app; +} + +// --------------------------------------------------------------------------- +// Change creation (src/services/change-flow.ts) +// --------------------------------------------------------------------------- + +describe("createChangeWithEvaluation records attributed costs", () => { + async function create(project: ProjectEntry) { + return createChangeWithEvaluation(makeEnv(), logger, { + project, + projectName: "@alice/my-repo", + workspaceName: WORKSPACE.name, + workspaceRemote: WORKSPACE.remote, + actor: { userId: "user_alice" }, + }); + } + + it("bills the owning user, and carries the evaluator's own source through", async () => { + const outcome = await create(projectEntry()); + expect(outcome.success).toBe(true); + + const byKind = Object.fromEntries(costRows().map((r) => [r.kind, r])); + expect(byKind.git_ops).toMatchObject({ + owner_id: "user_alice", + owner_type: "user", + source: "platform", + change_id: "chg_abc123", + project_id: "proj_abc", + workspace: "ws-1", + }); + // The evaluator paid its own provider; the clone of both repos was ours. + expect(byKind.llm_tokens).toMatchObject({ + owner_id: "user_alice", + owner_type: "user", + source: "byok", + quantity: 1500, + }); + }); + + it("bills an org-owned project to the org", async () => { + await create(projectEntry({ ownerId: "org_acme", ownerType: "org" })); + for (const row of costRows()) { + expect(row.owner_id).toBe("org_acme"); + expect(row.owner_type).toBe("org"); + } + }); + + it("walks an agent-owned project to its user, where the evaluation context could not", async () => { + // `billingContextFor` yields no context for an agent-owned project, so the + // evaluators see none. The ledger can still name a payer, because + // `resolveBillingSubject` is allowed the D1 read that walk needs. + raw + .prepare("INSERT INTO users (id, email, token_hash, created_at) VALUES (?, ?, ?, ?)") + .run("user_alice", "alice@example.com", "hash", "2026-01-01T00:00:00.000Z"); + raw + .prepare("INSERT INTO agents (id, name, owner_id, token_hash) VALUES (?, ?, ?, ?)") + .run("agt_bot", "bot", "user_alice", "agent-hash"); + + await create(projectEntry({ ownerId: "agt_bot", ownerType: "agent" })); + + expect(costRows().length).toBeGreaterThan(0); + for (const row of costRows()) { + expect(row.owner_id).toBe("user_alice"); + expect(row.owner_type).toBe("user"); + } + }); + + it("still records the change's costs when no owner can be resolved", async () => { + // A failed resolution must cost the attribution, never the change: the + // change is created and the row lands with a NULL owner. + const outcome = await create(projectEntry({ ownerId: "" })); + expect(outcome.success).toBe(true); + expect(costRows().length).toBeGreaterThan(0); + for (const row of costRows()) expect(row.owner_id).toBeNull(); + }); +}); + +// --------------------------------------------------------------------------- +// Re-evaluation and merge (src/routes/changes.ts) +// --------------------------------------------------------------------------- + +describe("POST /api/changes/:id/evaluate records attributed costs", () => { + it("attributes the re-evaluation's spend to the project's owner", async () => { + const res = await makeApp(changesRouter).fetch( + new Request("http://localhost/api/changes/chg_abc123/evaluate", { + method: "POST", + headers: { ...USER_AUTH, "Content-Type": "application/json" }, + body: "{}", + }), + makeEnv(), + ); + + expect(res.status).toBe(200); + const byKind = Object.fromEntries(costRows().map((r) => [r.kind, r])); + expect(byKind.llm_tokens).toMatchObject({ + owner_id: "user_alice", + owner_type: "user", + source: "byok", + change_id: "chg_abc123", + }); + expect(byKind.git_ops?.owner_id).toBe("user_alice"); + }); +}); + +describe("POST /api/changes/:id/merge records attributed costs", () => { + it("attributes the merge's git operations to the project's owner", async () => { + const res = await makeApp(changesRouter).fetch( + new Request("http://localhost/api/changes/chg_abc123/merge", { + method: "POST", + headers: USER_AUTH, + }), + makeEnv(), + ); + + expect(res.status).toBe(200); + expect(costRows()).toHaveLength(1); + expect(costRows()[0]).toMatchObject({ + kind: "git_ops", + owner_id: "user_alice", + owner_type: "user", + source: "platform", + change_id: "chg_abc123", + }); + }); +}); + +// --------------------------------------------------------------------------- +// 2.5 — the path that recorded nothing +// --------------------------------------------------------------------------- + +describe("POST /api/projects/conflicts/:id/resolve records attributed costs", () => { + const CONFLICT_CTX = { + conflictId: "conflict-abc", + namespace: "@alice", + slug: "my-repo", + workspaceName: "ws-1", + conflictingFiles: ["src/a.ts"], + detectedAt: "2026-01-01T00:00:00.000Z", + changeId: "chg_abc123", + }; + + function envWithConflict(ctx: Record = CONFLICT_CTX): Env { + const kv = makeFakeKV(); + kv.store.set(`conflict:${CONFLICT_CTX.conflictId}`, JSON.stringify(ctx)); + return makeEnv({ STATE: kv }); + } + + function resolveRequest() { + return new Request("http://localhost/api/projects/conflicts/conflict-abc/resolve", { + method: "POST", + headers: { ...USER_AUTH, "Content-Type": "application/json" }, + body: JSON.stringify({ + strategy: "manual", + resolutions: [{ file: "src/a.ts", content: "resolved\n" }], + }), + }); + } + + beforeEach(() => { + vi.mocked(buildManualResolutionDiff).mockResolvedValue({ + success: true, + data: { diff: "diff --git a/a.ts b/a.ts\n+line", baseSha: "conflict_base_sha" }, + }); + }); + + it("records the suite's spend against the project's owner", async () => { + const res = await makeApp(syncManagementRouter).fetch(resolveRequest(), envWithConflict()); + + expect(res.status).toBe(200); + const byKind = Object.fromEntries(costRows().map((r) => [r.kind, r])); + expect(byKind.llm_tokens).toMatchObject({ + owner_id: "user_alice", + owner_type: "user", + source: "byok", + quantity: 1500, + // A resolution is part of landing the change whose merge hit the conflict. + change_id: "chg_abc123", + workspace: "ws-1", + project: "@alice/my-repo", + }); + expect(byKind.git_ops?.owner_id).toBe("user_alice"); + }); + + it("records the spend even when the suite rejects the resolution", async () => { + // The tokens were spent producing the rejection. Recording after the + // verdict would make a blocked resolution free, which is exactly the shape + // of spend an unmetered account would learn to produce. The evaluator + // still runs and still reports what it cost — it just says no. + control.llmPassed = false; + + const res = await makeApp(syncManagementRouter).fetch(resolveRequest(), envWithConflict()); + + expect(res.status).toBe(422); + const byKind = Object.fromEntries(costRows().map((r) => [r.kind, r])); + expect(byKind.llm_tokens).toMatchObject({ owner_id: "user_alice", quantity: 1500 }); + }); + + it("leaves change_id null when the conflict predates that field", async () => { + const { changeId: _dropped, ...legacy } = CONFLICT_CTX; + const res = await makeApp(syncManagementRouter).fetch( + resolveRequest(), + envWithConflict(legacy), + ); + + expect(res.status).toBe(200); + for (const row of costRows()) { + expect(row.change_id).toBeNull(); + // Still attributed: the payer comes from the project, not the change. + expect(row.owner_id).toBe("user_alice"); + } + }); +}); diff --git a/tests/cost-attribution.test.ts b/tests/cost-attribution.test.ts new file mode 100644 index 00000000..b1e3ecb8 --- /dev/null +++ b/tests/cost-attribution.test.ts @@ -0,0 +1,284 @@ +/** + * Cost attribution: migration 048's columns, and the one function allowed to + * decide who a cost row is billed to. + * + * D1 is real SQLite with the production migrations applied, so these assertions + * run against the shipped schema — the CHECK constraints, the NOT NULL DEFAULT + * that backfills `source`, and the owner index — rather than against a stub's + * idea of it. + */ +import { describe, expect, it, vi } from "vitest"; +import { + type CostSample, + getChangeCostSummary, + recordCosts, + resolveBillingSubject, +} from "../src/storage/costs"; +import type { ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { makeSqliteD1, makeThrowingD1 } from "./helpers/sqlite-d1"; + +const logger: Logger = { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + child: vi.fn(() => logger), +}; + +interface CostRow { + project: string; + project_id: string | null; + change_id: string | null; + kind: string; + quantity: number; + owner_id: string | null; + owner_type: string | null; + source: string; + created_at: string; +} + +function projectEntry(overrides: Partial = {}): ProjectEntry { + return { + id: "proj_abc", + name: "my-repo", + slug: "my-repo", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://artifacts.example.com/repos/my-repo", + createdAt: "2026-01-01T00:00:00.000Z", + ...overrides, + }; +} + +// --------------------------------------------------------------------------- +// 2.2 — resolveBillingSubject +// --------------------------------------------------------------------------- + +describe("resolveBillingSubject", () => { + function withAgent(ownerId: string, agentId = "agt_bot") { + const { db, raw } = makeSqliteD1(); + raw + .prepare("INSERT INTO users (id, email, token_hash, created_at) VALUES (?, ?, ?, ?)") + .run("user_alice", "alice@example.com", "hash", "2026-01-01T00:00:00.000Z"); + raw + .prepare( + "INSERT INTO agents (id, name, owner_id, token_hash, created_at) VALUES (?, ?, ?, ?, ?)", + ) + .run(agentId, "bot", ownerId, "agent-hash", "2026-01-01T00:00:00.000Z"); + return db; + } + + it("bills a user-owned project to its owner", async () => { + const { db } = makeSqliteD1(); + await expect(resolveBillingSubject(db, logger, projectEntry())).resolves.toEqual({ + ownerId: "user_alice", + ownerType: "user", + }); + }); + + it("bills an org-owned project to the org", async () => { + const { db } = makeSqliteD1(); + const project = projectEntry({ ownerId: "org_acme", ownerType: "org" }); + await expect(resolveBillingSubject(db, logger, project)).resolves.toEqual({ + ownerId: "org_acme", + ownerType: "org", + }); + }); + + it("takes a bare owner pair as readily as a ProjectEntry", async () => { + // The recording sites hold a whole entry; nothing should have to synthesize + // one to ask this question. + const { db } = makeSqliteD1(); + await expect( + resolveBillingSubject(db, logger, { ownerId: "org_acme", ownerType: "org" }), + ).resolves.toEqual({ ownerId: "org_acme", ownerType: "org" }); + }); + + it("walks an agent-owned project to the user that owns the agent", async () => { + // Defensive: no creation path writes ownerType "agent" today. It exists so + // a restored backup or a future ownership transfer bills the human behind + // the agent rather than an id no account can be reconciled against. + const db = withAgent("user_alice"); + const project = projectEntry({ ownerId: "agt_bot", ownerType: "agent" }); + await expect(resolveBillingSubject(db, logger, project)).resolves.toEqual({ + ownerId: "user_alice", + ownerType: "user", + }); + }); + + it("yields nothing, and logs, when the agent row is gone", async () => { + const db = withAgent("user_alice", "agt_other"); + const project = projectEntry({ ownerId: "agt_bot", ownerType: "agent" }); + await expect(resolveBillingSubject(db, logger, project)).resolves.toBeNull(); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining("Cost attribution skipped"), + expect.objectContaining({ agentId: "agt_bot" }), + ); + }); + + it.each([ + ["no ownerId", { ownerId: "" }], + ["no ownerType", { ownerType: undefined }], + ])("yields nothing for a project with %s", async (_label, overrides) => { + // KV entries are cast without shape validation, so the type's promise is + // not a runtime guarantee. + const { db } = makeSqliteD1(); + const project = { ...projectEntry(), ...overrides } as ProjectEntry; + await expect(resolveBillingSubject(db, logger, project)).resolves.toBeNull(); + }); + + it("never throws when the database is unreachable", async () => { + // The contract that matters: this runs inside change creation, merge and + // deploy. A dead D1 must cost the attribution, never the change. + const project = projectEntry({ ownerId: "agt_bot", ownerType: "agent" }); + await expect( + resolveBillingSubject(makeThrowingD1("D1 down"), logger, project), + ).resolves.toBeNull(); + }); +}); + +// --------------------------------------------------------------------------- +// 2.1 / 2.3 — the columns, and what recordCosts writes into them +// --------------------------------------------------------------------------- + +describe("recordCosts attribution columns", () => { + function rows(raw: ReturnType["raw"]): CostRow[] { + return raw.prepare("SELECT * FROM cost_records ORDER BY kind").all() as unknown as CostRow[]; + } + + it("stamps the billing subject on every sample of a batch", async () => { + const { db, raw } = makeSqliteD1(); + const subject = await resolveBillingSubject(db, logger, projectEntry()); + expect(subject).not.toBeNull(); + + const result = await recordCosts( + db, + logger, + { + project: "@alice/my-repo", + projectId: "proj_abc", + changeId: "chg_1", + workspace: "ws-1", + ...(subject ?? {}), + }, + [ + { kind: "git_ops", quantity: 2 }, + { kind: "llm_tokens", quantity: 1200, estimated: true }, + ], + ); + + expect(result.success).toBe(true); + expect(rows(raw)).toHaveLength(2); + for (const row of rows(raw)) { + expect(row.owner_id).toBe("user_alice"); + expect(row.owner_type).toBe("user"); + } + }); + + it("records the row unattributed rather than refusing it when no owner resolves", async () => { + // The spend happened. A NULL owner is honest about being unbillable; losing + // the row would lose the only evidence the resource was consumed. + const { db, raw } = makeSqliteD1(); + const subject = await resolveBillingSubject(db, logger, projectEntry({ ownerId: "" })); + expect(subject).toBeNull(); + + const result = await recordCosts( + db, + logger, + { project: "@alice/my-repo", changeId: "chg_1", ...(subject ?? {}) }, + [{ kind: "llm_tokens", quantity: 900, estimated: true }], + ); + + expect(result.success).toBe(true); + expect(rows(raw)[0]?.owner_id).toBeNull(); + expect(rows(raw)[0]?.owner_type).toBeNull(); + // Still summarized for the change page, exactly as before attribution. + const summary = await getChangeCostSummary(db, logger, "chg_1"); + expect(summary.success && summary.data[0]?.total).toBe(900); + }); + + it("defaults source to platform for a sample that does not name one", async () => { + const { db, raw } = makeSqliteD1(); + await recordCosts(db, logger, { project: "p" }, [{ kind: "git_ops", quantity: 1 }]); + expect(rows(raw)[0]?.source).toBe("platform"); + }); + + it("carries a byok sample through unchanged", async () => { + // An evaluator running on the project's own credential is the only thing + // that can know this; both recording sites just flatten what it reports. + const { db, raw } = makeSqliteD1(); + const samples: CostSample[] = [ + { kind: "llm_tokens", quantity: 800, estimated: false, source: "byok" }, + { kind: "git_ops", quantity: 1 }, + ]; + await recordCosts(db, logger, { project: "p", ownerId: "org_acme", ownerType: "org" }, samples); + + const byKind = Object.fromEntries(rows(raw).map((r) => [r.kind, r])); + expect(byKind.llm_tokens?.source).toBe("byok"); + expect(byKind.git_ops?.source).toBe("platform"); + // Attribution is independent of who paid the provider: an org running BYOK + // is still the subject the usage is counted against. + expect(byKind.llm_tokens?.owner_type).toBe("org"); + }); + + it("backfills pre-048 rows as platform spend", async () => { + // The NOT NULL DEFAULT is correct by construction: BYOK did not exist when + // these rows were written, so every one of them was the operator's spend. + const { db, raw } = makeSqliteD1(); + raw + .prepare( + "INSERT INTO cost_records (id, project, kind, quantity, estimated, created_at) VALUES (?, ?, ?, ?, ?, ?)", + ) + .run("cost_legacy", "p", "llm_tokens", 500, 0, "2026-01-01T00:00:00.000Z"); + + const legacy = rows(raw)[0]; + expect(legacy?.source).toBe("platform"); + expect(legacy?.owner_id).toBeNull(); + // And the historical row still aggregates alongside attributed ones. + const summary = await getChangeCostSummary(db, logger, "chg_1"); + expect(summary.success).toBe(true); + }); + + it("serves an owner-scoped, newest-first history from the new index", async () => { + // What idx_costs_owner exists for. `created_at` is written only by + // recordCosts as an ISO 8601 string, so ordering it in SQL is safe here — + // the mixed-format hazard 042/044/047 warn about needs two writers. + const { db, raw } = makeSqliteD1(); + await recordCosts(db, logger, { project: "p", ownerId: "user_alice", ownerType: "user" }, [ + { kind: "git_ops", quantity: 1 }, + ]); + await recordCosts(db, logger, { project: "p", ownerId: "org_acme", ownerType: "org" }, [ + { kind: "git_ops", quantity: 5 }, + ]); + + const mine = raw + .prepare("SELECT quantity FROM cost_records WHERE owner_id = ? ORDER BY created_at DESC") + .all("user_alice") as unknown as Array<{ quantity: number }>; + expect(mine.map((r) => r.quantity)).toEqual([1]); + + const indexes = raw + .prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'cost_records'") + .all() as unknown as Array<{ name: string }>; + expect(indexes.map((i) => i.name)).toContain("idx_costs_owner"); + }); + + it("reports, and does not throw, when the owner_type CHECK rejects a row", async () => { + // The CHECK is the last line of defense against an owner_type no billing + // subject can have. It must surface as a Result, not an exception thrown + // into a merge. + const { db } = makeSqliteD1(); + const result = await recordCosts( + db, + logger, + // Only reachable by lying to the compiler — which is the point: the + // database refuses it even if a future caller bypasses the type. + { project: "p", ownerId: "agt_bot", ownerType: "agent" as "user" }, + [{ kind: "git_ops", quantity: 1 }], + ); + expect(result.success).toBe(false); + }); +}); diff --git a/tests/deploy-runner.test.ts b/tests/deploy-runner.test.ts index d09e3d54..f1d4c8dd 100644 --- a/tests/deploy-runner.test.ts +++ b/tests/deploy-runner.test.ts @@ -413,6 +413,38 @@ describe("'skipped' means only that nothing was configured", () => { }); }); +describe("cost attribution", () => { + it("bills the tree read to the project's owner", async () => { + // `readTree` records a git op whether or not the read succeeds, so the + // attribution has to be on it either way — this is a metered path that + // never goes near an evaluator. + await run(MERGE_MESSAGE, { readFiles: readsTree(tree(null)) }); + + const costs = raw + .prepare("SELECT owner_id, owner_type, source, kind, change_id FROM cost_records") + .all() as unknown as Array>; + expect(costs).toHaveLength(1); + expect(costs[0]).toMatchObject({ + kind: "git_ops", + // PROJECT is org-owned: the org is the payer, not the user who merged. + owner_id: "usr_alice", + owner_type: "org", + source: "platform", + change_id: CHANGE_ID, + }); + }); + + it("still records the read when the tree could not be read", async () => { + await run(MERGE_MESSAGE, { readFiles: unreadableTree }); + + const costs = raw + .prepare("SELECT owner_id FROM cost_records") + .all() as unknown as Array<{ owner_id: string | null }>; + expect(costs).toHaveLength(1); + expect(costs[0]?.owner_id).toBe("usr_alice"); + }); +}); + describe("tenant scoping", () => { it("fails closed when the message carries no projectId", async () => { const readFiles = readsTree(tree()); diff --git a/tests/evaluator-billing-context.test.ts b/tests/evaluator-billing-context.test.ts index 7f5c1b48..3b4a5a8b 100644 --- a/tests/evaluator-billing-context.test.ts +++ b/tests/evaluator-billing-context.test.ts @@ -114,7 +114,14 @@ vi.mock("../src/storage/users", () => ({ vi.mock("../src/storage/agents", () => ({ getAgentByToken: vi.fn(async () => ({ success: false, error: { code: "NOT_FOUND" } })), - getAgent: vi.fn(), + // Cost attribution walks this for an agent-owned project (`resolveBillingSubject`). + // These tests assert on the evaluation context, not the ledger, so the walk + // resolving to nothing is the right default here — it costs the attribution + // and never the change. + getAgent: vi.fn(async () => ({ + success: false, + error: { code: "NOT_FOUND", message: "Agent not found" }, + })), })); import { DiffEvaluator } from "../src/evaluation/diff-evaluator"; diff --git a/tests/post-merge.test.ts b/tests/post-merge.test.ts index b2f7eff3..93a54d9c 100644 --- a/tests/post-merge.test.ts +++ b/tests/post-merge.test.ts @@ -7,6 +7,7 @@ import { getCommitParent, readRepoFiles, revertToCommit } from "../src/storage/g import type { Env, ProjectEntry, SandboxInstance } from "../src/types"; import type { Logger } from "../src/utils/logger"; import { makeExecutingSandbox } from "./helpers/fake-sandbox"; +import { makeSqliteD1 } from "./helpers/sqlite-d1"; vi.mock("../src/storage/git-ops", () => ({ readRepoFiles: vi.fn(), @@ -322,3 +323,81 @@ describe("runPostMergeCheck", () => { expect(run).toHaveBeenCalledWith("make check", { timeout: 120_000 }); }); }); + +describe("runPostMergeCheck cost attribution", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(readRepoFiles).mockResolvedValue({ + success: true, + data: new Map([["src/index.ts", new TextEncoder().encode("export {};")]]), + }); + }); + + /** Real SQLite with the production migrations, so the row is really written. */ + function sandboxEnvWithDb(): { env: Env; raw: ReturnType["raw"] } { + const { db, raw } = makeSqliteD1(); + const sandbox: SandboxInstance = { + writeFile: vi.fn().mockResolvedValue(undefined), + run: vi.fn().mockResolvedValue({ exitCode: 0, stdout: "ok", stderr: "" }), + destroy: vi.fn().mockResolvedValue(undefined), + }; + const env = { + DB: db, + SANDBOX: { create: vi.fn().mockResolvedValue(sandbox) }, + } as unknown as Env; + return { env, raw }; + } + + it("bills the sandbox run and the tree read to the project's owner", async () => { + // This is the one metered path outside the evaluators: a post-merge command + // is sandbox time somebody pays for. + const { env, raw } = sandboxEnvWithDb(); + const result = await runPostMergeCheck( + env, + project, + { + changeId: "chg_1", + mergeCommit: "sha_merge", + policy: policyWith({ postMergeCommand: "npm test" }), + }, + mockLogger, + ); + + expect(result.status).toBe("passed"); + const rows = raw + .prepare( + "SELECT kind, owner_id, owner_type, source, change_id FROM cost_records ORDER BY kind", + ) + .all() as unknown as Array>; + expect(rows.map((r) => r.kind)).toEqual(["git_ops", "sandbox_ms"]); + for (const row of rows) { + expect(row.owner_id).toBe("user_1"); + expect(row.owner_type).toBe("user"); + // Sandbox time is always the operator's compute; BYOK is an LLM concept. + expect(row.source).toBe("platform"); + expect(row.change_id).toBe("chg_1"); + } + }); + + it("records the run unattributed rather than skipping it when the owner is unnameable", async () => { + const { env, raw } = sandboxEnvWithDb(); + const ownerless = { ...project, ownerId: "" } as ProjectEntry; + const result = await runPostMergeCheck( + env, + ownerless, + { + changeId: "chg_1", + mergeCommit: "sha_merge", + policy: policyWith({ postMergeCommand: "npm test" }), + }, + mockLogger, + ); + + expect(result.status).toBe("passed"); + const rows = raw + .prepare("SELECT owner_id FROM cost_records") + .all() as unknown as Array<{ owner_id: string | null }>; + expect(rows).toHaveLength(2); + for (const row of rows) expect(row.owner_id).toBeNull(); + }); +}); From 0be3c7135772e01fac2db457ddb1dd0dd21e7045 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 01:12:29 +0000 Subject: [PATCH 06/26] Correct claims the attribution comments got wrong MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Conventions review of Task 2. Two Biome formatting errors that would have failed CI, and four comments asserting things the code does not do. The miscounts were stale on arrival: two comments said "both recording sites" when this same diff added a third, and one said the resolution rule exists "once rather than five times" against six call sites. All now say it without counting, since the count is what keeps going stale. Migration 048's header claimed the new index serves a read that exists — no caller issues it yet; it is built for the owner-scoped usage read this work is heading toward, and now says so. It also claimed recordCosts holds the only INSERT into cost_records "anywhere in the tree", which the test added alongside it falsified by inserting a legacy row directly. The sortability argument survives narrowed to application code, which is the only part it needed. The header was also 46 lines of preamble where 042/047 run 9 to 11 and put reasoning inline against the column it constrains. Restructured to match: each ALTER now carries its own argument, and the speculative ROLLBACK section — a shape no other migration here uses, which conceded in its own text that nothing has ever been reverted — is gone. Typecheck clean; lint clean; 3764 tests pass. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- migrations/048_cost_attribution.sql | 77 ++++++++++++----------------- src/deploy/runner.ts | 2 +- src/evaluation/types.ts | 5 +- src/routes/sync-management.ts | 4 +- src/storage/costs.ts | 2 +- tests/cost-attribution.test.ts | 2 +- tests/deploy-runner.test.ts | 6 +-- tests/post-merge.test.ts | 6 +-- 8 files changed, 46 insertions(+), 58 deletions(-) diff --git a/migrations/048_cost_attribution.sql b/migrations/048_cost_attribution.sql index 91b533b6..3c344658 100644 --- a/migrations/048_cost_attribution.sql +++ b/migrations/048_cost_attribution.sql @@ -2,60 +2,47 @@ -- -- `cost_records` (021, given `project_id` by 025) has recorded what a change -- consumed since it existed, but never who owes for it. A project name is not a --- billing subject: names are unique only per-namespace, projects are transferred --- and deleted, and an aggregate keyed on one cannot be reconciled against an --- account. These three columns make every new row attributable at the moment --- the spend is incurred, which is the only moment the payer is known for +-- billing subject: names are unique only per-namespace, and projects are +-- transferred and deleted. These columns make every new row attributable at the +-- moment the spend is incurred, which is the only moment the payer is known for -- certain. --- --- `owner_id`/`owner_type` are NULLABLE, deliberately. Every historical row --- predates attribution and the mapping needed to backfill it lives in KV (the --- ProjectEntry's owner), not in D1 — the same reason 025 backfilled no --- `project_id`. SQL cannot resolve it here, and discarding rows that cannot be --- attributed would delete the only record of spend that actually happened. So --- unattributed rows stay, visibly unattributed: a NULL owner is a row nobody --- can be billed for, which is honest, where an owner guessed from a project --- name would be a row billed to possibly the wrong account. Live code can also --- write NULL — `resolveBillingSubject` (src/storage/costs.ts) returns null --- rather than throwing when a project's owner cannot be named — so this is not --- a transitional state that later becomes NOT NULL. --- --- `source` is NOT NULL DEFAULT 'platform' because the backfill is correct by --- construction, not merely convenient: BYOK does not exist yet, so every row --- already in this table was spend on the operator's own provider account. The --- default backfills them in place with the truth. Precedent for both shapes of --- ALTER: a CHECK on an added column in 037, NOT NULL DEFAULT in 032/034/041. --- No table rebuild — unlike `kind`, whose closed CHECK would need one to widen. --- --- The new index sorts by `created_at`, which is safe HERE specifically. --- `cost_records.created_at` is written exclusively by `recordCosts` --- (src/storage/costs.ts) as `new Date().toISOString()` — that function holds the --- only INSERT into this table anywhere in the tree — so every value is ISO 8601 --- with a 'T' separator. The hazard 042/044/047 warn about is a column mixing --- that format with `datetime('now')`'s space separator, which sorts differently --- as TEXT (' ' 0x20 < 'T' 0x54) and silently breaks range and ORDER BY --- comparisons. A single-format column has no such problem, so an owner's usage --- history can be paged newest-first in SQL rather than in JavaScript. --- --- ROLLBACK: drop `idx_costs_owner`; leave the three columns in place. SQLite's --- DROP COLUMN would rewrite the table for no gain, and the columns are inert --- when unwritten — reads name them explicitly, an absent owner is already the --- normal historical case, and `source` defaults itself. Code must therefore --- tolerate the columns being present but never populated. No migration in this --- repo has been reverted; this records the intent rather than building --- machinery for it. + +-- NULLABLE, and permanently so. Every historical row predates attribution, and +-- the mapping needed to backfill it lives in KV (the ProjectEntry's owner), not +-- in D1 — the same reason 025 backfilled no `project_id`. Live code writes NULL +-- too: `resolveBillingSubject` (src/storage/costs.ts) returns null rather than +-- throwing when an owner cannot be named, because `recordCosts` is best-effort +-- inside change creation, merge and deploy. So a NULL owner is a row nobody can +-- be billed for, which is honest, where an owner guessed from a project name +-- would be a row billed to possibly the wrong account. +ALTER TABLE cost_records ADD COLUMN owner_id TEXT; -- CHECK'd rather than free text: an owner_type outside this set is unreachable -- from `resolveBillingSubject`, which maps ProjectEntry's third kind ('agent') -- onto the user that owns the agent instead of passing it through. A row -- claiming a type no billing subject can have would aggregate under an account -- that will never be invoiced. -ALTER TABLE cost_records ADD COLUMN owner_id TEXT; ALTER TABLE cost_records ADD COLUMN owner_type TEXT CHECK(owner_type IN ('user','org')); + +-- NOT NULL DEFAULT because the backfill is correct by construction, not merely +-- convenient: BYOK does not exist yet, so every row already in this table was +-- spend on the operator's own provider account, and the default states that +-- truth in place. Precedent for both ALTER shapes: a CHECK on an added column in +-- 037, NOT NULL DEFAULT in 032/034/041. No table rebuild — unlike `kind`, whose +-- closed CHECK would need one to widen. ALTER TABLE cost_records ADD COLUMN source TEXT NOT NULL DEFAULT 'platform' CHECK(source IN ('platform','byok')); --- Covers the owner-scoped read this exists for: `WHERE owner_id = ? ORDER BY --- created_at DESC`. Rows with a NULL owner are still indexed (SQLite indexes --- NULLs) but are never selected by that predicate, which is the intent. +-- Serves the owner-scoped usage read this work is being built toward +-- (`WHERE owner_id = ? ORDER BY created_at DESC`); no caller issues it yet. +-- +-- Sorting on `created_at` is safe HERE specifically. Application code writes it +-- only through `recordCosts`, as `new Date().toISOString()`, so every row this +-- table gets in production is ISO 8601 with a 'T' separator. The hazard +-- 042/044/047 warn about is a column mixing that with `datetime('now')`'s space +-- separator, which sorts differently as TEXT (' ' 0x20 < 'T' 0x54) and silently +-- breaks every range comparison. A single-format column has no such problem. +-- +-- Rows with a NULL owner are still indexed (SQLite indexes NULLs) but are never +-- selected by that predicate, which is the intent. CREATE INDEX IF NOT EXISTS idx_costs_owner ON cost_records(owner_id, created_at DESC); diff --git a/src/deploy/runner.ts b/src/deploy/runner.ts index a4d136e1..5f0b107b 100644 --- a/src/deploy/runner.ts +++ b/src/deploy/runner.ts @@ -852,9 +852,9 @@ async function readTree( projectDefaultBranch(project), ); + const subject = await resolveBillingSubject(env.DB, logger, project); // Recorded whether or not the read succeeded: a clone that failed still cost // the round trip. - const subject = await resolveBillingSubject(env.DB, logger, project); await recordCosts( env.DB, logger, diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index ad4c733b..e5e2ebca 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -13,8 +13,9 @@ export interface EvalResult { * Structurally a subset of `CostSample` (src/storage/costs.ts) rather than an * import of it, so the evaluation layer does not depend on storage. `source` * is what lets an evaluator running on the project's own provider credential - * say so: both recording sites flatten this array straight into `recordCosts`, - * so a field missing here is a distinction that cannot be made at all. + * say so: every recording site flattens this array straight into + * `recordCosts`, so a field missing here is a distinction that cannot be made + * at all. * Omitted means `"platform"` — the operator paid. */ costs?: Array<{ diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index 89d627ef..f1077021 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -675,8 +675,8 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { // Recorded here, BEFORE the verdict is acted on, because the spend is // already incurred: a resolution the suite rejects burned exactly the same // model tokens as one it accepts, and the early return below would drop - // them. Mirrors the recording POST /changes/:id/evaluate does, minus the - // git_ops pair — `buildManualResolutionDiff` cloned once, above. + // them. Mirrors the recording POST /changes/:id/evaluate does, with one + // git_op rather than its pair — `buildManualResolutionDiff` cloned once. const resolveSubject = await resolveBillingSubject(c.env.DB, logger, project); await recordCosts( c.env.DB, diff --git a/src/storage/costs.ts b/src/storage/costs.ts index 1d33c97d..5934e0fd 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -52,7 +52,7 @@ interface SummaryRow { * * The single home for the `ownerType -> billing subject` mapping: every * recording site resolves through here so the agent walk and the - * cannot-be-attributed rule exist once rather than five times. Accepts a whole + * cannot-be-attributed rule exist once rather than at each of them. Accepts a whole * `ProjectEntry` (structurally) or a bare owner pair. * * An **agent is not a payer** — it belongs to one — so `"agent"` walks diff --git a/tests/cost-attribution.test.ts b/tests/cost-attribution.test.ts index b1e3ecb8..7515e273 100644 --- a/tests/cost-attribution.test.ts +++ b/tests/cost-attribution.test.ts @@ -209,7 +209,7 @@ describe("recordCosts attribution columns", () => { it("carries a byok sample through unchanged", async () => { // An evaluator running on the project's own credential is the only thing - // that can know this; both recording sites just flatten what it reports. + // that can know this; every recording site just flattens what it reports. const { db, raw } = makeSqliteD1(); const samples: CostSample[] = [ { kind: "llm_tokens", quantity: 800, estimated: false, source: "byok" }, diff --git a/tests/deploy-runner.test.ts b/tests/deploy-runner.test.ts index f1d4c8dd..4eb72720 100644 --- a/tests/deploy-runner.test.ts +++ b/tests/deploy-runner.test.ts @@ -437,9 +437,9 @@ describe("cost attribution", () => { it("still records the read when the tree could not be read", async () => { await run(MERGE_MESSAGE, { readFiles: unreadableTree }); - const costs = raw - .prepare("SELECT owner_id FROM cost_records") - .all() as unknown as Array<{ owner_id: string | null }>; + const costs = raw.prepare("SELECT owner_id FROM cost_records").all() as unknown as Array<{ + owner_id: string | null; + }>; expect(costs).toHaveLength(1); expect(costs[0]?.owner_id).toBe("usr_alice"); }); diff --git a/tests/post-merge.test.ts b/tests/post-merge.test.ts index 93a54d9c..bb65ac7b 100644 --- a/tests/post-merge.test.ts +++ b/tests/post-merge.test.ts @@ -394,9 +394,9 @@ describe("runPostMergeCheck cost attribution", () => { ); expect(result.status).toBe("passed"); - const rows = raw - .prepare("SELECT owner_id FROM cost_records") - .all() as unknown as Array<{ owner_id: string | null }>; + const rows = raw.prepare("SELECT owner_id FROM cost_records").all() as unknown as Array<{ + owner_id: string | null; + }>; expect(rows).toHaveLength(2); for (const row of rows) expect(row.owner_id).toBeNull(); }); From 27a6fccb29c08e106e38514458cd987cc0877364 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 01:15:27 +0000 Subject: [PATCH 07/26] Erase cost owner ids, and bill the resolve that skips the suite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Edge-case review of Task 2. Two defects in the attribution work itself. cost_records.owner_id was missing from IDENTITY_COLUMNS, so migration 048 introduced a column holding an account id that account erasure could not reach. A cost row survives an erasure whenever its project is org-owned, or agent-owned — only ownerType "user" projects are cascaded at all. The agent case is the sharp one: resolveBillingSubject's walk stamps the row with the owning user's id, and step 3 of the cascade then deletes the agent, destroying the only link that could re-derive the attribution. The table has no foreign key to users, so nothing else caught it, and a bare unresolvable account id stayed behind. Only owner_id is rewritten; quantity and kind are what the spend was, not who incurred it, and the row must stay readable as spend that happened. The existing test pins the column count exactly, which is what surfaced this. The new conflict-resolution recording sat inside the manual branch, but resolveConflict clones the project and the workspace fork and pushes for every strategy. So accept-project and accept-workspace resolutions recorded nothing at all, and manual under-recorded by the clone and push it had just paid for. The git work is now billed where it happens, two git_ops — the shape the single-merge route bills for the same clone and push — leaving the manual branch to record only the evaluator suite. Typecheck clean; lint clean; 3766 tests pass, up 2. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- src/routes/sync-management.ts | 20 +++++++++++++++ src/storage/deletion.ts | 10 ++++++++ tests/cost-attribution-paths.test.ts | 37 ++++++++++++++++++++++++++++ tests/deletion-account.test.ts | 6 ++++- 4 files changed, 72 insertions(+), 1 deletion(-) diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index f1077021..3d7b1c48 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -790,6 +790,26 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { const { commitSha } = resolveResult.data; + // `resolveConflict` clones the project AND the workspace fork and pushes, for + // EVERY strategy — so this is recorded here rather than in the `manual` branch + // above, which only ever sees the evaluator suite's spend. Without it, + // accept-project and accept-workspace resolutions record nothing at all, and + // manual under-records by the clone+push it just paid for. Two git_ops, the + // same shape the single-merge route bills for the same clone+push. + const resolveGitSubject = await resolveBillingSubject(c.env.DB, logger, project); + await recordCosts( + c.env.DB, + logger, + { + project: projectDisplayName(project), + projectId: project.id, + ...(conflictCtx.changeId ? { changeId: conflictCtx.changeId } : {}), + workspace: conflictCtx.workspaceName, + ...(resolveGitSubject ?? {}), + }, + [{ kind: "git_ops", quantity: 2 }], + ); + // Record history (non-throwing); delete conflict key regardless of history outcome await recordSyncHistory( c.env.DB, diff --git a/src/storage/deletion.ts b/src/storage/deletion.ts index 21536e0a..ef974396 100644 --- a/src/storage/deletion.ts +++ b/src/storage/deletion.ts @@ -673,6 +673,16 @@ const IDENTITY_COLUMNS: readonly [table: string, column: string][] = [ // agent's id under requested_by_type = 'agent' and is left alone. ["deployments", "requested_by_id"], ["deployments", "approved_by"], + // Cost rows outlive an erasure by two routes: the project is org-owned, or it + // is agent-owned and so never cascaded at all (only `ownerType === "user"` + // projects are). In the agent case the row was stamped with the *user* id by + // `resolveBillingSubject`'s agent walk, and step 3 of the cascade then deletes + // the agent — destroying the only link that could re-derive the attribution + // and leaving a bare, unresolvable account id behind. `cost_records` has no + // foreign key to `users`, so nothing else catches it. Only `owner_id` is + // rewritten: `quantity` and `kind` are what the spend was, not who incurred + // it, and the row must stay readable as spend that happened. + ["cost_records", "owner_id"], ]; /** diff --git a/tests/cost-attribution-paths.test.ts b/tests/cost-attribution-paths.test.ts index 718b4e25..aa8f6de5 100644 --- a/tests/cost-attribution-paths.test.ts +++ b/tests/cost-attribution-paths.test.ts @@ -453,6 +453,43 @@ describe("POST /api/projects/conflicts/:id/resolve records attributed costs", () expect(byKind.llm_tokens).toMatchObject({ owner_id: "user_alice", quantity: 1500 }); }); + it("records the clone and push for a strategy that runs no evaluator", async () => { + // accept-project and accept-workspace skip the evaluator suite entirely, but + // resolveConflict still clones both repos and pushes for every strategy. The + // suite's recording lives in the manual branch, so billing the git work + // there would have made these two resolutions free. + const request = new Request("http://localhost/api/projects/conflicts/conflict-abc/resolve", { + method: "POST", + headers: { ...USER_AUTH, "Content-Type": "application/json" }, + body: JSON.stringify({ strategy: "accept-project" }), + }); + + const res = await makeApp(syncManagementRouter).fetch(request, envWithConflict()); + + expect(res.status).toBe(200); + const rows = costRows(); + expect(rows).toHaveLength(1); + expect(rows[0]).toMatchObject({ + kind: "git_ops", + quantity: 2, + owner_id: "user_alice", + owner_type: "user", + project: "@alice/my-repo", + }); + }); + + it("bills the manual path for the resolve clone on top of the suite", async () => { + const res = await makeApp(syncManagementRouter).fetch(resolveRequest(), envWithConflict()); + + expect(res.status).toBe(200); + // One git_op for buildManualResolutionDiff's clone, two for resolveConflict's + // clone-both-and-push. Under-counting here was the whole defect. + const gitOps = costRows() + .filter((r) => r.kind === "git_ops") + .reduce((sum, r) => sum + r.quantity, 0); + expect(gitOps).toBe(3); + }); + it("leaves change_id null when the conflict predates that field", async () => { const { changeId: _dropped, ...legacy } = CONFLICT_CTX; const res = await makeApp(syncManagementRouter).fetch( diff --git a/tests/deletion-account.test.ts b/tests/deletion-account.test.ts index 99a0892b..a0c6ce23 100644 --- a/tests/deletion-account.test.ts +++ b/tests/deletion-account.test.ts @@ -80,7 +80,7 @@ describe("anonymizeUserContributions", () => { expect(result.success).toBe(true); const updates = executed.filter((s) => s.sql.startsWith("UPDATE")); - expect(updates.length).toBe(11); + expect(updates.length).toBe(12); for (const stmt of updates) { expect(stmt.sql).toMatch(/^UPDATE \w+ SET \w+ = \? WHERE \w+ = \?$/); expect(stmt.bindings).toEqual([DELETED_USER_SENTINEL, "usr_1"]); @@ -93,6 +93,10 @@ describe("anonymizeUserContributions", () => { expect(cols).toContain("UPDATE webhooks SET created_by = ? WHERE created_by = ?"); expect(cols).toContain("UPDATE project_secrets SET updated_by = ? WHERE updated_by = ?"); expect(cols).toContain("UPDATE deployments SET approved_by = ? WHERE approved_by = ?"); + // A cost row survives an erasure whenever its project is org- or + // agent-owned, and the agent it named is deleted outright — so without this + // the row keeps a bare account id that nothing can resolve or erase. + expect(cols).toContain("UPDATE cost_records SET owner_id = ? WHERE owner_id = ?"); }); }); From 02faa1d5e1e07e9a98342f96b9d3ae98ad809515 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 01:58:47 +0000 Subject: [PATCH 08/26] Keep the usage aggregate out of reach of project deletion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 3, and the fixes from Task 2's bug review. cost_records is project-scoped and hard-deleted by the project cascade, so an allowance computed from it is refunded by deleting the project: burn the month, delete, start again. usage_periods is owner-scoped, joins the account cascades instead, and is deliberately absent from PROJECT_SCOPED_TABLES. The regression test is the point of the table — record spend, delete the project through the real cascade, assert the ledger rows are gone and the aggregate is byte-identical. It accumulates inside recordCosts, so all six recording sites get it without any of them being able to forget. Its Result is deliberately not folded into recordCosts's: the ledger rows were written, so returning an error would tell the caller the spend went unrecorded when it did not. The upsert is one atomic statement rather than a read-then-write, and the test discriminates — the same test against a read-then-write yields 1 where the upsert yields 50. git_ops maps to no meter. It counts clones and pushes, several per change, and no allowance is set on it; routing it to deploys_month would be wrong in both directions. getProjectCostSummary is deleted rather than generalized: it was dead, and it aggregated the very table this table exists to stop aggregating from. The Task 2 review also found the conflict-resolution git recording I added sat after the failure branch, so a resolution that cloned both repos and then failed to push recorded nothing — the same loss the sibling recording twelve lines above has a comment arguing against. Now before it. And the merge attribution test asserted a branch production never takes. POST /changes/:id/merge routes to the merge queue whenever MERGE_QUEUE is bound — it is, in all three configs — and that branch returns before the recording. The test bound no queue, so a passing suite implied merges were metered when they are not. It now says which path it covers, and a sibling pins the gap so it cannot be mistaken for coverage again. Three assertions that passed for the wrong reason: an ordering test whose rows shared a millisecond and whose index check only asked whether the index existed, not whether the planner used it; a backfill test ending on an assertion that held against any schema; and a module-level logger mock that was never reset between tests. Typecheck clean; lint clean; 3788 tests pass, up 22. Coverage ratchet holds. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- migrations/049_usage_periods.sql | 65 +++++ src/routes/sync-management.ts | 19 +- src/storage/costs.ts | 83 +++--- src/storage/d1-backup.ts | 7 + src/storage/deletion.ts | 18 ++ src/storage/usage.ts | 215 ++++++++++++++ tests/cost-attribution-paths.test.ts | 50 +++- tests/cost-attribution.test.ts | 56 +++- tests/usage-periods.test.ts | 414 +++++++++++++++++++++++++++ 9 files changed, 864 insertions(+), 63 deletions(-) create mode 100644 migrations/049_usage_periods.sql create mode 100644 src/storage/usage.ts create mode 100644 tests/usage-periods.test.ts diff --git a/migrations/049_usage_periods.sql b/migrations/049_usage_periods.sql new file mode 100644 index 00000000..f4545476 --- /dev/null +++ b/migrations/049_usage_periods.sql @@ -0,0 +1,65 @@ +-- Migration 049: the owner-scoped usage aggregate a limit is compared against. +-- +-- `cost_records` cannot be that aggregate, and this is the whole reason for a +-- second table. It is in PROJECT_SCOPED_TABLES (src/storage/deletion.ts) and +-- hard-deleted by the project cascade, so any allowance summed from it is +-- refunded by deleting a project: burn the month's quota, delete the project, +-- get the quota back. This table keys on the billing subject instead and joins +-- only the ACCOUNT cascades, so the one thing that clears usage is erasing the +-- account that incurred it — at which point there is no subject left to bill. +-- +-- It is a derived aggregate, not a rival ledger. `cost_records` remains the +-- per-change evidence of what was spent; this is the running total. Both are +-- written by `recordCosts` (src/storage/costs.ts) from the same timestamp, so a +-- cost row and the period it counts toward can never disagree, and a lost +-- aggregate write is reconcilable from the ledger it derives from. +-- +-- Timestamps are ISO 8601 strings written by the application, NOT +-- `datetime('now')`, for the reason 042/044/047 give: the two formats sort +-- differently as TEXT (' ' 0x20 < 'T' 0x54), and a column mixing them silently +-- breaks every comparison over it. +-- +-- Rollback: drop the table. Writes are best-effort by contract and reads are +-- Result-typed, so an absent table degrades to logged failures rather than a +-- broken change flow — but every accumulated total is gone with it. + +CREATE TABLE IF NOT EXISTS usage_periods ( + -- The subject `resolveBillingSubject` named: a user or an org, never a + -- project and never an agent. NOT NULL, unlike `cost_records.owner_id`, + -- which 048 deliberately left nullable so an unattributable sample is still + -- recorded. There is no such thing here: a quantity nobody can be billed for + -- cannot be enforced against anyone, so `recordCosts` drops those samples + -- rather than accumulating a bucket no limit could ever apply to. + owner_id TEXT NOT NULL, + -- CHECK'd, and closed, for the same reason as `cost_records.owner_type`: an + -- owner_type outside this set is unreachable from `resolveBillingSubject`, + -- which maps ProjectEntry's third kind ('agent') onto the user that owns the + -- agent instead of passing it through. Unlike `meter` below, this set does + -- not grow — it is exactly what `BillingSubject` admits. + owner_type TEXT NOT NULL CHECK(owner_type IN ('user','org')), + -- 'YYYY-MM', UTC (PRD Open Question 1). Zero-padded, so lexical order is + -- chronological order and a period range is a plain TEXT comparison. + period TEXT NOT NULL, + -- Deliberately NOT CHECK'd — the one place this table departs from 048's + -- shape, and a decision rather than an omission. `cost_records.kind` carries + -- a closed CHECK from 021, and the bill for it is already recorded in this + -- feature's own scope: `storage_bytes` is a Non-Goal partly because adding a + -- kind would force a 043-style table rebuild. The meter set is expected to + -- grow (`deploys_month` has no recording site yet), so a CHECK here would buy + -- validation that the `MeterKey` union already gives at the only writer, at a + -- cost that has demonstrably bitten once. An unrecognised meter is an inert + -- row nothing reads; an unwritable one is a failed write on a hot path. + meter TEXT NOT NULL, + -- REAL, matching `cost_records.quantity`, which is what accumulates into it: + -- sandbox milliseconds and character-derived token estimates are not + -- guaranteed integral, and rounding at the aggregate would drift. + quantity REAL NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL, + -- `owner_type` is deliberately outside the key: it describes the owner, it is + -- not part of its identity. In the key it would let one account hold two rows + -- for the same meter and month, and a limit compared against either alone + -- would be wrong. The key doubles as the UPSERT's conflict target + -- (src/storage/usage.ts) and as the covering index for the only read + -- (`WHERE owner_id = ? AND period = ?`), so no separate index is added. + PRIMARY KEY (owner_id, period, meter) +); diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index 3d7b1c48..454a827e 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -783,19 +783,17 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { logger, ); - if (!resolveResult.success) { - const status = resolveResult.error.statusCode === 401 ? 401 : 422; - return c.json({ error: resolveResult.error.message, code: resolveResult.error.code }, status); - } - - const { commitSha } = resolveResult.data; - // `resolveConflict` clones the project AND the workspace fork and pushes, for // EVERY strategy — so this is recorded here rather than in the `manual` branch // above, which only ever sees the evaluator suite's spend. Without it, // accept-project and accept-workspace resolutions record nothing at all, and // manual under-records by the clone+push it just paid for. Two git_ops, the // same shape the single-merge route bills for the same clone+push. + // + // Before the failure branch below, for the same reason the manual branch + // records before its 422: the clones are already paid for when the push + // fails, and a failure that costs the operator the same as a success must not + // be the cheap way to use the platform. const resolveGitSubject = await resolveBillingSubject(c.env.DB, logger, project); await recordCosts( c.env.DB, @@ -810,6 +808,13 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { [{ kind: "git_ops", quantity: 2 }], ); + if (!resolveResult.success) { + const status = resolveResult.error.statusCode === 401 ? 401 : 422; + return c.json({ error: resolveResult.error.message, code: resolveResult.error.code }, status); + } + + const { commitSha } = resolveResult.data; + // Record history (non-throwing); delete conflict key regardless of history outcome await recordSyncHistory( c.env.DB, diff --git a/src/storage/costs.ts b/src/storage/costs.ts index 5934e0fd..9fc2663f 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -3,6 +3,7 @@ import { newId } from "../utils/ids"; import type { Logger } from "../utils/logger"; import { type Result, err, ok } from "../utils/result"; import { getAgent } from "./agents"; +import { type UsageDelta, meterForCostKind, upsertUsage, usagePeriod } from "./usage"; export type CostKind = "llm_tokens" | "sandbox_ms" | "git_ops"; @@ -102,6 +103,51 @@ export async function resolveBillingSubject( return { ownerId: agentResult.data.ownerId, ownerType: "user" }; } +/** + * Accumulate the metered half of a batch into `usage_periods`. + * + * Split out of `recordCosts` so all six recording sites get it without any of + * them knowing it exists — the aggregate must never be something a call site + * can forget, because a site that forgets is an owner whose spend does not + * count against their allowance. + * + * Samples with no billing subject are skipped. `cost_records` keeps them with a + * NULL owner because the spend still happened, but a quantity nobody can be + * billed for cannot be enforced against anyone, so there is no row here to + * write it into — `usage_periods.owner_id` is NOT NULL for exactly that reason. + * + * The period comes from the same `createdAt` the ledger rows carry, so a cost + * row and the month it counts toward cannot land on opposite sides of a + * boundary. + */ +async function accumulateUsage( + db: D1Database, + logger: Logger, + opts: { ownerId?: string; ownerType?: BillingSubject["ownerType"] }, + samples: CostSample[], + createdAt: string, +): Promise { + if (!opts.ownerId || !opts.ownerType) return; + const deltas: UsageDelta[] = []; + for (const sample of samples) { + // `git_ops` maps to no meter; see `meterForCostKind`. + const meter = meterForCostKind(sample.kind); + if (meter) deltas.push({ meter, quantity: sample.quantity }); + } + if (deltas.length === 0) return; + // The Result is deliberately not folded into `recordCosts`'s own: the ledger + // rows above ARE written, and reporting an error would tell the caller the + // spend went unrecorded when it did not. `upsertUsage` logs its own failure, + // and the aggregate is reconcilable from the ledger it derives from. + await upsertUsage( + db, + logger, + { ownerId: opts.ownerId, ownerType: opts.ownerType }, + usagePeriod(new Date(createdAt)), + deltas, + ); +} + /** * Record cost samples for a change. Best-effort: failures are logged and * reported, but callers treat cost recording as non-blocking. @@ -155,6 +201,9 @@ export async function recordCosts( changeId: opts.changeId, count: samples.length, }); + // Only after the ledger write succeeded: the aggregate derives from those + // rows, so it must never count spend the ledger does not record. + await accumulateUsage(db, logger, opts, samples, createdAt); return ok(undefined); } catch (error) { const appError = @@ -204,37 +253,3 @@ export async function getChangeCostSummary( return err(appError); } } - -export async function getProjectCostSummary( - db: D1Database, - logger: Logger, - project: string, -): Promise> { - try { - const result = await db - .prepare( - "SELECT kind, SUM(quantity) AS total, MAX(estimated) AS any_estimated FROM cost_records WHERE project = ? GROUP BY kind", - ) - .bind(project) - .all(); - return ok( - result.results.map((row) => ({ - kind: row.kind as CostKind, - total: row.total, - estimated: row.any_estimated === 1, - })), - ); - } catch (error) { - const appError = - error instanceof AppError - ? error - : new AppError( - error instanceof Error ? error.message : "Failed to summarize costs", - "DATABASE_ERROR", - 500, - { operation: "getProjectCostSummary", project }, - ); - logger.error("Failed to get project cost summary", appError, { project }); - return err(appError); - } -} diff --git a/src/storage/d1-backup.ts b/src/storage/d1-backup.ts index 5b594151..71fadf6a 100644 --- a/src/storage/d1-backup.ts +++ b/src/storage/d1-backup.ts @@ -43,6 +43,13 @@ export const BACKUP_TABLES: readonly string[] = [ "import_metrics", "failed_imports", "cost_records", + // The owner-scoped usage aggregate (migration 049). Backed up rather than + // excluded, and NOT recoverable from `cost_records` above: the project + // cascade hard-deletes cost rows while this table survives, so replaying the + // ledger after a restore would rebuild only the usage of projects that still + // exist and hand everybody else their month back. Declares no foreign keys, + // so its position here is only to sit beside the ledger it derives from. + "usage_periods", "events", "sync_history", "issues", diff --git a/src/storage/deletion.ts b/src/storage/deletion.ts index ef974396..f8d18c7c 100644 --- a/src/storage/deletion.ts +++ b/src/storage/deletion.ts @@ -683,6 +683,14 @@ const IDENTITY_COLUMNS: readonly [table: string, column: string][] = [ // rewritten: `quantity` and `kind` are what the spend was, not who incurred // it, and the row must stay readable as spend that happened. ["cost_records", "owner_id"], + // `usage_periods.owner_id` is deliberately NOT here, and its absence is a + // decision rather than an oversight. Rewriting it to the shared sentinel + // would collide two erased users' rows on PRIMARY KEY (owner_id, period, + // meter) and throw mid-erasure; and there is nothing worth preserving. A cost + // row is evidence that spend happened and reads perfectly well with the payer + // anonymized, whereas an aggregate exists only to be compared against a limit + // for a subject that no longer exists. `deleteAccountCascade` below deletes + // those rows outright instead. ]; /** @@ -785,6 +793,10 @@ async function resolveOrgOwnership( .bind(org.id) .run(); await db.prepare("DELETE FROM teams WHERE org_id = ?").bind(org.id).run(); + // The org's usage aggregate goes with the org, and ONLY on this branch: where + // a successor was promoted above, the org survives and so must its + // consumption — a change of owner is not a fresh monthly allowance. + await db.prepare("DELETE FROM usage_periods WHERE owner_id = ?").bind(org.id).run(); await db.prepare("DELETE FROM orgs WHERE id = ?").bind(org.id).run(); } @@ -859,6 +871,12 @@ export async function deleteAccountCascade( // erasure request is itself a retention bug. await db.prepare("DELETE FROM agents WHERE owner_id = ?").bind(userId).run(); await db.prepare("DELETE FROM api_tokens WHERE user_id = ?").bind(userId).run(); + // The owner-scoped usage aggregate (migration 049). Erasing the account is + // the ONLY event that may clear it: `usage_periods` is deliberately outside + // `PROJECT_SCOPED_TABLES` because an allowance a user can reset by deleting + // the project that burned it is not an allowance. Here the subject itself is + // going, so there is nothing left to bill and nothing left to enforce. + await db.prepare("DELETE FROM usage_periods WHERE owner_id = ?").bind(userId).run(); // Same reasoning, same trap: `oauth_tokens.user_id` and // `oauth_auth_codes.user_id` both REFERENCE users(id) (#349). An OAuth // grant is a live credential handed to an editor, so an erasure that left diff --git a/src/storage/usage.ts b/src/storage/usage.ts new file mode 100644 index 00000000..0e0b22f5 --- /dev/null +++ b/src/storage/usage.ts @@ -0,0 +1,215 @@ +import { AppError } from "../utils/errors"; +import type { Logger } from "../utils/logger"; +import { type Result, err, ok } from "../utils/result"; +import type { BillingSubject, CostKind } from "./costs"; + +/** + * A metered quantity an allowance can be set against, over a calendar month. + * + * The `_month` suffix is redundant with the row's `period` column and is kept + * anyway: these are the exact strings entitlements keys its limits on, so a + * limit lookup is `limits[row.meter]` rather than a second mapping table that + * could disagree with this one. + * + * `deploys_month` has no writer yet, deliberately — see `meterForCostKind`. + */ +export type MeterKey = "llm_tokens_month" | "sandbox_ms_month" | "deploys_month"; + +/** One meter's increment. Quantities are added, never replaced. */ +export interface UsageDelta { + meter: MeterKey; + quantity: number; +} + +export interface UsageSummaryEntry { + meter: MeterKey; + quantity: number; + /** ISO 8601, app-written. When this owner last consumed this meter. */ + updatedAt: string; +} + +interface UsageRow { + meter: string; + quantity: number; + updated_at: string; +} + +/** + * The meter a cost sample accumulates into, or `null` when it accumulates into + * none. + * + * `CostKind` and `MeterKey` are not the same vocabulary and are not meant to + * be. A cost kind names a *resource consumed by one change*; a meter names a + * *thing an allowance is set on*. Two align exactly: + * + * - `llm_tokens` -> `llm_tokens_month` + * - `sandbox_ms` -> `sandbox_ms_month` + * + * `git_ops` maps to nothing. It is the count of clones and pushes an operation + * performed — several per change — and there is no `git_ops_month` allowance + * in the PRD's meter set. The tempting misreading is to route it to + * `deploys_month`, which would be wrong in both directions: `post-merge.ts` + * records `git_ops` for a merge that deploys nothing, and `deploy/runner.ts` + * records one for a config read that may precede a deploy that never happens. + * `deploys_month` counts `deployments` rows, so its writer is the deploy + * consumer and not this function; the key exists in `MeterKey` now only so + * that adding that writer needs no migration. + */ +export function meterForCostKind(kind: CostKind): MeterKey | null { + switch (kind) { + case "llm_tokens": + return "llm_tokens_month"; + case "sandbox_ms": + return "sandbox_ms_month"; + case "git_ops": + return null; + } +} + +/** + * The 'YYYY-MM' UTC period a moment falls in (PRD Open Question 1). + * + * Derived from the ISO string rather than `getMonth()` so it is UTC by + * construction: a local-time month would put a spend just after midnight UTC + * into the previous period for half the world, and two counters keyed on + * different months are two allowances. + * + * @param at - The moment to classify; defaults to now + * @returns The calendar month in 'YYYY-MM' form, UTC + */ +export function usagePeriod(at: Date = new Date()): string { + return at.toISOString().slice(0, 7); +} + +/** + * Add to an owner's meters for one period, creating the rows if they do not + * exist yet. + * + * Each meter is a single SQLite UPSERT, which is where the correctness lives: + * `ON CONFLICT ... DO UPDATE SET quantity = quantity + excluded.quantity` reads + * and writes inside one statement, so concurrent recorders sum. A + * read-then-write would let N of them all read the same total and each write + * back "+1", losing N-1 increments — precisely the race + * `MagicLinkRateLimiter`'s doc comment describes for KV counters. + * + * Deliberately not `db.batch`: the meters are independent, so cross-meter + * atomicity buys nothing, and one statement per meter keeps the atomicity claim + * at the statement where it is actually true. Callers send at most two. + * + * Best-effort by contract, like `recordCosts` which drives it: this returns a + * Result and logs, and never throws into change creation, merge or deploy. + * + * @param subject - The account the usage is billed to + * @param period - 'YYYY-MM' UTC, from `usagePeriod` + * @param deltas - Increments to apply; same-meter entries are summed first + * @returns Success, or the database error that stopped the accumulation + */ +export async function upsertUsage( + db: D1Database, + logger: Logger, + subject: BillingSubject, + period: string, + deltas: readonly UsageDelta[], +): Promise> { + const merged = new Map(); + for (const delta of deltas) { + // A NaN or Infinity would poison the running total permanently: SQLite has + // no way back from it and every later comparison against a limit would be + // false. Drop the sample and keep the meter usable. + if (!Number.isFinite(delta.quantity)) { + logger.warn("Usage delta skipped: quantity is not finite", { + ownerId: subject.ownerId, + meter: delta.meter, + quantity: delta.quantity, + }); + continue; + } + merged.set(delta.meter, (merged.get(delta.meter) ?? 0) + delta.quantity); + } + if (merged.size === 0) return ok(undefined); + + const updatedAt = new Date().toISOString(); + try { + for (const [meter, quantity] of merged) { + await db + .prepare( + "INSERT INTO usage_periods (owner_id, owner_type, period, meter, quantity, updated_at) " + + "VALUES (?, ?, ?, ?, ?, ?) " + + "ON CONFLICT(owner_id, period, meter) DO UPDATE SET " + + "quantity = quantity + excluded.quantity, updated_at = excluded.updated_at", + ) + .bind(subject.ownerId, subject.ownerType, period, meter, quantity, updatedAt) + .run(); + } + logger.debug("Usage aggregate updated", { + ownerId: subject.ownerId, + period, + meters: merged.size, + }); + return ok(undefined); + } catch (error) { + const appError = + error instanceof AppError + ? error + : new AppError( + error instanceof Error ? error.message : "Failed to accumulate usage", + "DATABASE_ERROR", + 500, + { operation: "upsertUsage", ownerId: subject.ownerId, period }, + ); + logger.error("Failed to accumulate usage", appError, { ownerId: subject.ownerId, period }); + return err(appError); + } +} + +/** + * What one account consumed in one period, one entry per meter it touched. + * + * The period is a required argument rather than defaulted to the current month + * for the reason the PRD gives for `UsageMeter`'s `nowMs`: the caller decides + * which window it is asking about, so one request cannot straddle a boundary + * and read half of each. Meters with no rows are absent rather than zero — + * "never used" and "used nothing" are the same fact, and the caller knows the + * meter set it wants to display. + * + * @param ownerId - The billing subject's id + * @param period - 'YYYY-MM' UTC, from `usagePeriod` + * @returns The per-meter totals, or the database error + */ +export async function getOwnerUsageSummary( + db: D1Database, + logger: Logger, + ownerId: string, + period: string, +): Promise> { + try { + const result = await db + .prepare( + "SELECT meter, quantity, updated_at FROM usage_periods WHERE owner_id = ? AND period = ? ORDER BY meter", + ) + .bind(ownerId, period) + .all(); + return ok( + result.results.map((row) => ({ + // Cast, not validated: the column carries no CHECK on purpose (see + // migration 049), so a row written by a future meter this build does + // not know about reads through rather than failing the whole summary. + meter: row.meter as MeterKey, + quantity: row.quantity, + updatedAt: row.updated_at, + })), + ); + } catch (error) { + const appError = + error instanceof AppError + ? error + : new AppError( + error instanceof Error ? error.message : "Failed to summarize usage", + "DATABASE_ERROR", + 500, + { operation: "getOwnerUsageSummary", ownerId, period }, + ); + logger.error("Failed to get owner usage summary", appError, { ownerId, period }); + return err(appError); + } +} diff --git a/tests/cost-attribution-paths.test.ts b/tests/cost-attribution-paths.test.ts index aa8f6de5..95729e58 100644 --- a/tests/cost-attribution-paths.test.ts +++ b/tests/cost-attribution-paths.test.ts @@ -361,14 +361,24 @@ describe("POST /api/changes/:id/evaluate records attributed costs", () => { }); describe("POST /api/changes/:id/merge records attributed costs", () => { - it("attributes the merge's git operations to the project's owner", async () => { - const res = await makeApp(changesRouter).fetch( - new Request("http://localhost/api/changes/chg_abc123/merge", { - method: "POST", - headers: USER_AUTH, - }), - makeEnv(), - ); + function mergeRequest(query = ""): Request { + return new Request(`http://localhost/api/changes/chg_abc123/merge${query}`, { + method: "POST", + headers: USER_AUTH, + }); + } + + // NOTE ON WHAT THIS COVERS. `makeEnv()` binds no MERGE_QUEUE, which is the + // ONLY reason control reaches the recordCosts this asserts on. With the + // binding present — as it is in the top-level, production and staging configs + // — a `strategy=merge` request takes the MergeQueue branch at + // `changes.ts:589`, which returns before that recording. So this covers the + // cold path a self-hoster without the binding gets, and `?strategy=squash`. + // It does NOT cover a default merge on the hosted instance, which records + // nothing at all; the sibling test below pins that gap so it cannot be + // mistaken for coverage. + it("attributes the cold path's git operations to the project's owner", async () => { + const res = await makeApp(changesRouter).fetch(mergeRequest(), makeEnv()); expect(res.status).toBe(200); expect(costRows()).toHaveLength(1); @@ -380,6 +390,30 @@ describe("POST /api/changes/:id/merge records attributed costs", () => { change_id: "chg_abc123", }); }); + + it("records NOTHING when the merge queue is bound, which is every deployed config", async () => { + // Asserting the gap rather than the fix, deliberately. Closing it means + // recording inside the queue branch, which is its own change; until then a + // green suite must not imply merges are metered in production. Flip this to + // expect a row when that lands. + // `transitioned: false` is the queue branch's earliest clean exit — it + // returns 200 right after logging, skipping the event emit and post-merge + // work this test has no business exercising. Any successful outcome reaches + // the same conclusion: no recordCosts on the way out. + const queueEnv = makeEnv({ + MERGE_QUEUE: { + idFromName: (name: string) => name, + get: () => ({ + merge: async () => ({ success: true, transitioned: false, commit: "merged_sha" }), + }), + } as unknown as Env["MERGE_QUEUE"], + }); + + const res = await makeApp(changesRouter).fetch(mergeRequest(), queueEnv); + + expect(res.status).toBe(200); + expect(costRows()).toHaveLength(0); + }); }); // --------------------------------------------------------------------------- diff --git a/tests/cost-attribution.test.ts b/tests/cost-attribution.test.ts index 7515e273..5d9a8997 100644 --- a/tests/cost-attribution.test.ts +++ b/tests/cost-attribution.test.ts @@ -7,7 +7,7 @@ * that backfills `source`, and the owner index — rather than against a stub's * idea of it. */ -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { type CostSample, getChangeCostSummary, @@ -28,6 +28,10 @@ const logger: Logger = { child: vi.fn(() => logger), }; +// The mock is module-level, so without this a `toHaveBeenCalledWith` assertion +// can be satisfied by an earlier test's call and the suite passes on run order. +beforeEach(() => vi.clearAllMocks()); + interface CostRow { project: string; project_id: string | null; @@ -238,32 +242,56 @@ describe("recordCosts attribution columns", () => { const legacy = rows(raw)[0]; expect(legacy?.source).toBe("platform"); expect(legacy?.owner_id).toBeNull(); - // And the historical row still aggregates alongside attributed ones. + + // And it still aggregates alongside an attributed row. The legacy INSERT + // above carries no change_id, so it must be given one to be summarized at + // all — asserting only that the summary succeeds would pass against an + // empty result, and therefore against any schema. + raw.prepare("UPDATE cost_records SET change_id = ? WHERE id = ?").run("chg_1", "cost_legacy"); + await recordCosts( + db, + logger, + { project: "p", changeId: "chg_1", ownerId: "user_alice", ownerType: "user" }, + [{ kind: "llm_tokens", quantity: 250 }], + ); + const summary = await getChangeCostSummary(db, logger, "chg_1"); expect(summary.success).toBe(true); + if (!summary.success) return; + const llm = summary.data.find((entry) => entry.kind === "llm_tokens"); + expect(llm?.total).toBe(750); }); it("serves an owner-scoped, newest-first history from the new index", async () => { // What idx_costs_owner exists for. `created_at` is written only by // recordCosts as an ISO 8601 string, so ordering it in SQL is safe here — // the mixed-format hazard 042/044/047 warn about needs two writers. - const { db, raw } = makeSqliteD1(); - await recordCosts(db, logger, { project: "p", ownerId: "user_alice", ownerType: "user" }, [ - { kind: "git_ops", quantity: 1 }, - ]); - await recordCosts(db, logger, { project: "p", ownerId: "org_acme", ownerType: "org" }, [ - { kind: "git_ops", quantity: 5 }, - ]); + const { raw } = makeSqliteD1(); + // Written directly rather than through recordCosts: that function stamps + // `new Date().toISOString()`, and successive calls land inside the same + // millisecond, so an ordering assertion over them would hold whatever the + // ORDER BY said. Distinct timestamps are the whole point of the test. + const insert = raw.prepare( + "INSERT INTO cost_records (id, project, kind, quantity, estimated, created_at, owner_id, owner_type, source) VALUES (?, ?, ?, ?, 0, ?, ?, ?, 'platform')", + ); + insert.run("c_old", "p", "git_ops", 1, "2026-01-01T00:00:00.000Z", "user_alice", "user"); + insert.run("c_new", "p", "git_ops", 2, "2026-03-01T00:00:00.000Z", "user_alice", "user"); + insert.run("c_other", "p", "git_ops", 5, "2026-02-01T00:00:00.000Z", "org_acme", "org"); const mine = raw .prepare("SELECT quantity FROM cost_records WHERE owner_id = ? ORDER BY created_at DESC") .all("user_alice") as unknown as Array<{ quantity: number }>; - expect(mine.map((r) => r.quantity)).toEqual([1]); + // Newest first, and another subject's row is not in it. + expect(mine.map((r) => r.quantity)).toEqual([2, 1]); - const indexes = raw - .prepare("SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'cost_records'") - .all() as unknown as Array<{ name: string }>; - expect(indexes.map((i) => i.name)).toContain("idx_costs_owner"); + // The planner must actually reach for the index, not merely have it: an + // index nothing uses is a write cost with no read benefit. + const plan = raw + .prepare( + "EXPLAIN QUERY PLAN SELECT * FROM cost_records WHERE owner_id = ? ORDER BY created_at DESC", + ) + .all("user_alice") as unknown as Array<{ detail: string }>; + expect(plan.map((step) => step.detail).join(" ")).toContain("idx_costs_owner"); }); it("reports, and does not throw, when the owner_type CHECK rejects a row", async () => { diff --git a/tests/usage-periods.test.ts b/tests/usage-periods.test.ts new file mode 100644 index 00000000..41d2db60 --- /dev/null +++ b/tests/usage-periods.test.ts @@ -0,0 +1,414 @@ +/** + * The owner-scoped usage aggregate (migration 049) and what may erase it. + * + * The first describe block is the reason the table exists: `cost_records` is + * project-scoped and hard-deleted by the project cascade, so an allowance summed + * from it is refunded by deleting a project. Every other test here defends a + * property that regression depends on — that the aggregate is written at all, + * that concurrent writers sum, that a month boundary starts a new counter. + * + * D1 is real SQLite with the production migrations applied, so the PRIMARY KEY, + * the owner_type CHECK and the UPSERT's conflict target are the shipped ones + * rather than a stub's idea of them. + */ +import { describe, expect, it, vi } from "vitest"; +import { recordCosts } from "../src/storage/costs"; +import { + captureDeletionTarget, + deleteAccountCascade, + deleteProjectCascade, +} from "../src/storage/deletion"; +import { + getOwnerUsageSummary, + meterForCostKind, + upsertUsage, + usagePeriod, +} from "../src/storage/usage"; +import type { Env, ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { makeFakeKV } from "./helpers/fake-kv"; +import { makeSqliteD1, makeThrowingD1 } from "./helpers/sqlite-d1"; + +const logger: Logger = { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + child: vi.fn(() => logger), +}; + +const ALICE = { ownerId: "user_alice", ownerType: "user" } as const; +const PERIOD = "2026-09"; + +interface UsageRow { + owner_id: string; + owner_type: string; + period: string; + meter: string; + quantity: number; + updated_at: string; +} + +type Raw = ReturnType["raw"]; + +function usageRows(raw: Raw): UsageRow[] { + return raw + .prepare("SELECT * FROM usage_periods ORDER BY owner_id, period, meter") + .all() as unknown as UsageRow[]; +} + +function countCostRows(raw: Raw): number { + const row = raw.prepare("SELECT COUNT(*) AS n FROM cost_records").get() as { n: number }; + return row.n; +} + +function project(overrides: Partial = {}): ProjectEntry { + return { + id: "proj_abc", + name: "my-repo", + slug: "my-repo", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://acct.artifacts.cloudflare.net/git/@alice/my-repo.git", + createdAt: "2026-01-01T00:00:00.000Z", + ...overrides, + }; +} + +function makeEnv(db: D1Database, kv: KVNamespace): Env { + // No REPO_DO / MERGE_QUEUE: `purgeDurableObject` treats an absent binding as + // nothing to purge, which is what a minimal deployment does too. + return { + DB: db, + STATE: kv, + ARTIFACTS: { delete: async () => true } as unknown as Env["ARTIFACTS"], + } as Env; +} + +// --------------------------------------------------------------------------- +// 3.4 / 3.5 — the regression this table exists to prevent +// --------------------------------------------------------------------------- + +describe("deleting a project does not refund the allowance", () => { + it("keeps the owner's aggregate while destroying the project's cost rows", async () => { + // The defect in one test: burn the month's quota against a project, delete + // the project, and check the quota is still burnt. Before `usage_periods` + // the only aggregate lived in `cost_records`, which the cascade below + // hard-deletes, so this sequence handed the allowance straight back. + const { db, raw } = makeSqliteD1(); + const kv = makeFakeKV(); + const entry = project(); + await kv.put("project:@alice:my-repo", JSON.stringify(entry)); + const env = makeEnv(db, kv); + + await recordCosts( + db, + logger, + { project: "@alice/my-repo", projectId: entry.id, changeId: "chg_1", ...ALICE }, + [ + { kind: "llm_tokens", quantity: 12_000, estimated: true }, + { kind: "sandbox_ms", quantity: 4_500 }, + { kind: "git_ops", quantity: 2 }, + ], + ); + + const period = usagePeriod(); + const before = await getOwnerUsageSummary(db, logger, "user_alice", period); + expect(before.success && before.data).toEqual([ + { meter: "llm_tokens_month", quantity: 12_000, updatedAt: expect.any(String) }, + { meter: "sandbox_ms_month", quantity: 4_500, updatedAt: expect.any(String) }, + ]); + expect(countCostRows(raw)).toBe(3); + + const captured = await captureDeletionTarget(env, entry, logger); + expect(captured.success).toBe(true); + if (!captured.success) return; + const cascade = await deleteProjectCascade(env, captured.data, logger); + expect(cascade.success && cascade.data.residuals).toEqual([]); + + // The ledger is gone — that part is by design, the rows are project data. + expect(countCostRows(raw)).toBe(0); + // The allowance is not. + const after = await getOwnerUsageSummary(db, logger, "user_alice", period); + expect(after.success && after.data).toEqual(before.success && before.data); + }); + + it("leaves another owner's aggregate untouched when an account is erased", async () => { + // The mirror of the above: the cascade that IS allowed to clear usage must + // still be scoped to one subject. + const { db, raw } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, PERIOD, [{ meter: "llm_tokens_month", quantity: 10 }]); + await upsertUsage(db, logger, { ownerId: "org_acme", ownerType: "org" }, PERIOD, [ + { meter: "llm_tokens_month", quantity: 99 }, + ]); + + const env = makeEnv(db, makeFakeKV()); + const erased = await deleteAccountCascade(env, "user_alice", logger); + expect(erased.success && erased.data.residuals).toEqual([]); + + expect(usageRows(raw).map((r) => [r.owner_id, r.quantity])).toEqual([["org_acme", 99]]); + }); + + it("erases an org's aggregate only when the empty org itself is deleted", async () => { + // `resolveOrgOwnership` promotes a successor where one exists and deletes + // the org where none does. A change of owner is not a fresh allowance, so + // only the second branch may clear the counter. + async function eraseSoleOwner(withSurvivingMember: boolean): Promise { + const { db, raw } = makeSqliteD1(); + raw + .prepare("INSERT INTO users (id, email, token_hash, created_at) VALUES (?, ?, ?, ?)") + .run("user_alice", "alice@example.com", "h1", "2026-01-01T00:00:00.000Z"); + raw + .prepare("INSERT INTO orgs (id, name, slug, owner_id, created_at) VALUES (?, ?, ?, ?, ?)") + .run("org_acme", "Acme", "acme", "user_alice", "2026-01-01T00:00:00.000Z"); + if (withSurvivingMember) { + raw + .prepare("INSERT INTO users (id, email, token_hash, created_at) VALUES (?, ?, ?, ?)") + .run("user_bob", "bob@example.com", "h2", "2026-01-01T00:00:00.000Z"); + raw + .prepare("INSERT INTO org_members (org_id, user_id, role, joined_at) VALUES (?, ?, ?, ?)") + .run("org_acme", "user_bob", "admin", "2026-01-01T00:00:00.000Z"); + } + await upsertUsage(db, logger, { ownerId: "org_acme", ownerType: "org" }, PERIOD, [ + { meter: "sandbox_ms_month", quantity: 250 }, + ]); + await deleteAccountCascade(makeEnv(db, makeFakeKV()), "user_alice", logger); + return usageRows(raw); + } + + expect(await eraseSoleOwner(true)).toHaveLength(1); + expect(await eraseSoleOwner(false)).toEqual([]); + }); +}); + +// --------------------------------------------------------------------------- +// 3.2 — upsertUsage / getOwnerUsageSummary +// --------------------------------------------------------------------------- + +describe("upsertUsage", () => { + it("sums concurrent increments rather than losing all but one", async () => { + // The claim the single UPSERT statement buys, and the only reason it is + // written as one. A read-then-write would let all fifty callers observe the + // same total across the await between the SELECT and the UPDATE and each + // write back "+1", leaving 1 where 50 is correct. + const { db } = makeSqliteD1(); + await Promise.all( + Array.from({ length: 50 }, () => + upsertUsage(db, logger, ALICE, PERIOD, [{ meter: "llm_tokens_month", quantity: 1 }]), + ), + ); + const summary = await getOwnerUsageSummary(db, logger, "user_alice", PERIOD); + expect(summary.success && summary.data[0]?.quantity).toBe(50); + }); + + it("counts each calendar month separately", async () => { + const { db } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, "2026-09", [{ meter: "llm_tokens_month", quantity: 700 }]); + await upsertUsage(db, logger, ALICE, "2026-10", [{ meter: "llm_tokens_month", quantity: 5 }]); + + const september = await getOwnerUsageSummary(db, logger, "user_alice", "2026-09"); + const october = await getOwnerUsageSummary(db, logger, "user_alice", "2026-10"); + expect(september.success && september.data[0]?.quantity).toBe(700); + expect(october.success && october.data[0]?.quantity).toBe(5); + }); + + it("merges repeated meters in one call into a single row", async () => { + const { db, raw } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, PERIOD, [ + { meter: "llm_tokens_month", quantity: 10 }, + { meter: "llm_tokens_month", quantity: 5 }, + { meter: "sandbox_ms_month", quantity: 1 }, + ]); + expect(usageRows(raw).map((r) => [r.meter, r.quantity])).toEqual([ + ["llm_tokens_month", 15], + ["sandbox_ms_month", 1], + ]); + }); + + it("drops a non-finite quantity instead of poisoning the running total", async () => { + // SQLite has no way back from a NaN total, and every later comparison + // against a limit would silently be false. + const { db, raw } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, PERIOD, [{ meter: "llm_tokens_month", quantity: 100 }]); + const result = await upsertUsage(db, logger, ALICE, PERIOD, [ + { meter: "llm_tokens_month", quantity: Number.NaN }, + { meter: "sandbox_ms_month", quantity: Number.POSITIVE_INFINITY }, + ]); + expect(result.success).toBe(true); + expect(usageRows(raw).map((r) => [r.meter, r.quantity])).toEqual([["llm_tokens_month", 100]]); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining("not finite"), + expect.objectContaining({ meter: "llm_tokens_month" }), + ); + }); + + it("writes nothing, and succeeds, for an empty delta list", async () => { + const { db, raw } = makeSqliteD1(); + expect((await upsertUsage(db, logger, ALICE, PERIOD, [])).success).toBe(true); + expect(usageRows(raw)).toEqual([]); + }); + + it("stamps the owner type without letting it split the counter", async () => { + // owner_type is out of the PRIMARY KEY on purpose: one account must never + // hold two rows for the same meter and month, or a limit compared against + // either alone is wrong. + const { db, raw } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, PERIOD, [{ meter: "llm_tokens_month", quantity: 1 }]); + await upsertUsage(db, logger, { ownerId: "user_alice", ownerType: "org" }, PERIOD, [ + { meter: "llm_tokens_month", quantity: 1 }, + ]); + expect(usageRows(raw)).toHaveLength(1); + expect(usageRows(raw)[0]?.quantity).toBe(2); + }); + + it("reports a database failure as a Result rather than throwing", async () => { + // It runs inside change creation, merge and deploy; a dead D1 must cost the + // aggregate, never the change. + const result = await upsertUsage(makeThrowingD1("D1 down"), logger, ALICE, PERIOD, [ + { meter: "llm_tokens_month", quantity: 1 }, + ]); + expect(result.success).toBe(false); + expect(!result.success && result.error.code).toBe("DATABASE_ERROR"); + }); + + it("reports the owner_type CHECK as a Result rather than throwing", async () => { + const { db } = makeSqliteD1(); + const result = await upsertUsage( + db, + logger, + // Only reachable by lying to the compiler, which is the point: the + // database refuses an owner_type no billing subject can have. + { ownerId: "agt_bot", ownerType: "agent" as "user" }, + PERIOD, + [{ meter: "llm_tokens_month", quantity: 1 }], + ); + expect(result.success).toBe(false); + }); +}); + +describe("getOwnerUsageSummary", () => { + it("returns only the asked-for owner and period, ordered by meter", async () => { + const { db } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, PERIOD, [ + { meter: "sandbox_ms_month", quantity: 2 }, + { meter: "llm_tokens_month", quantity: 1 }, + ]); + await upsertUsage(db, logger, ALICE, "2026-08", [{ meter: "llm_tokens_month", quantity: 999 }]); + await upsertUsage(db, logger, { ownerId: "org_acme", ownerType: "org" }, PERIOD, [ + { meter: "llm_tokens_month", quantity: 999 }, + ]); + + const summary = await getOwnerUsageSummary(db, logger, "user_alice", PERIOD); + expect(summary.success && summary.data.map((e) => [e.meter, e.quantity])).toEqual([ + ["llm_tokens_month", 1], + ["sandbox_ms_month", 2], + ]); + }); + + it("returns an empty summary for an owner that has consumed nothing", async () => { + const { db } = makeSqliteD1(); + const summary = await getOwnerUsageSummary(db, logger, "user_nobody", PERIOD); + expect(summary.success && summary.data).toEqual([]); + }); + + it("reports a database failure as a Result rather than throwing", async () => { + const summary = await getOwnerUsageSummary(makeThrowingD1(), logger, "user_alice", PERIOD); + expect(summary.success).toBe(false); + }); +}); + +describe("usagePeriod", () => { + it("classifies a moment by its UTC month, not the local one", () => { + // A local-time month would put a spend just after midnight UTC into the + // previous period for half the world, and two counters keyed on different + // months are two allowances. + expect(usagePeriod(new Date("2026-09-30T23:59:59.999Z"))).toBe("2026-09"); + expect(usagePeriod(new Date("2026-10-01T00:00:00.000Z"))).toBe("2026-10"); + expect(usagePeriod(new Date("2026-01-01T00:00:00.000Z"))).toBe("2026-01"); + }); +}); + +// --------------------------------------------------------------------------- +// 3.3 — what recordCosts accumulates, and what it declines to +// --------------------------------------------------------------------------- + +describe("meterForCostKind", () => { + it("meters the two kinds an allowance is set on and no others", () => { + expect(meterForCostKind("llm_tokens")).toBe("llm_tokens_month"); + expect(meterForCostKind("sandbox_ms")).toBe("sandbox_ms_month"); + // Not `deploys_month`: git operations are recorded by paths that deploy + // nothing, so routing them there would count deploys that never happened. + expect(meterForCostKind("git_ops")).toBeNull(); + }); +}); + +describe("recordCosts accumulates usage alongside the ledger", () => { + it("writes both, from one period, for an attributed batch", async () => { + const { db, raw } = makeSqliteD1(); + const result = await recordCosts( + db, + logger, + { project: "@alice/my-repo", changeId: "chg_1", ...ALICE }, + [ + { kind: "llm_tokens", quantity: 1_500, estimated: true }, + { kind: "git_ops", quantity: 2 }, + ], + ); + expect(result.success).toBe(true); + expect(countCostRows(raw)).toBe(2); + + const rows = usageRows(raw); + expect(rows.map((r) => [r.meter, r.quantity])).toEqual([["llm_tokens_month", 1_500]]); + expect(rows[0]?.owner_type).toBe("user"); + // Ledger row and aggregate must agree about which month they belong to. + const created = raw.prepare("SELECT created_at FROM cost_records LIMIT 1").get() as { + created_at: string; + }; + expect(rows[0]?.period).toBe(usagePeriod(new Date(created.created_at))); + }); + + it("skips a sample no account can be billed for", async () => { + // The ledger keeps it with a NULL owner because the spend happened; the + // aggregate cannot, because a quantity nobody owes cannot be enforced + // against anyone. `usage_periods.owner_id` is NOT NULL for this reason. + const { db, raw } = makeSqliteD1(); + const result = await recordCosts(db, logger, { project: "@alice/my-repo" }, [ + { kind: "llm_tokens", quantity: 900, estimated: true }, + ]); + expect(result.success).toBe(true); + expect(countCostRows(raw)).toBe(1); + expect(usageRows(raw)).toEqual([]); + }); + + it("writes no aggregate row for a batch of unmetered kinds alone", async () => { + const { db, raw } = makeSqliteD1(); + await recordCosts(db, logger, { project: "@alice/my-repo", ...ALICE }, [ + { kind: "git_ops", quantity: 3 }, + ]); + expect(countCostRows(raw)).toBe(1); + expect(usageRows(raw)).toEqual([]); + }); + + it("accumulates across separate recordings, as a monthly counter must", async () => { + const { db, raw } = makeSqliteD1(); + for (const quantity of [100, 250, 40]) { + await recordCosts(db, logger, { project: "@alice/my-repo", ...ALICE }, [ + { kind: "llm_tokens", quantity }, + ]); + } + expect(usageRows(raw).map((r) => r.quantity)).toEqual([390]); + }); + + it("bills a BYOK sample to the same subject as a platform one", async () => { + // `source` records whose provider account paid; it does not change whose + // usage it is. An org running its own key still consumes org allowance. + const { db, raw } = makeSqliteD1(); + await recordCosts(db, logger, { project: "@acme/api", ownerId: "org_acme", ownerType: "org" }, [ + { kind: "llm_tokens", quantity: 800, source: "byok" }, + ]); + expect(usageRows(raw)[0]).toMatchObject({ owner_id: "org_acme", quantity: 800 }); + }); +}); From ada38d7224295521592c54123ef9aba79b27b91a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 02:49:17 +0000 Subject: [PATCH 09/26] Stop the usage aggregate being refundable, and add the entitlements seam MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 4, plus the review findings on Task 3. Two reviews of the aggregate found three ways it could still lose or give back what it counted. A negative delta subtracted from the month — a refund through the single statement the table exists to make non-refundable. Reachable without malice: sandbox_ms is wall-clock arithmetic, so a backwards clock step produces one. Now rejected in the guard and CHECK'd in the schema. The aggregate dropped `source`, collapsing BYOK and platform spend into one meter. An allowance limits what the OPERATOR spends, so a project paying its own provider bill would have been charged against a hosted allowance — and because this table is not rebuildable from cost_records, no month accumulated that way could ever have been separated again. The dimension is now part of the primary key, which was free to do today and impossible later. db.batch is one transaction and quantity is NOT NULL, while SQLite stores NaN as NULL. So one unusable sample did not cost itself, it aborted the write and took every other sample in the same evaluation with it. Samples are filtered before the batch now, which also explains why the aggregate's own guard is a backstop rather than the first line. Both the migration and the code claimed a lost aggregate write was reconcilable from the ledger. The backup module added in the same commit says the opposite, and it is right — the cascade deletes cost rows and leaves these standing. A lost upsert is a permanent under-count, and saying so is the point: it is accepted because failing a merge over an accounting write is worse. The conflict route billed a flat two git ops for every strategy on a comment claiming all three clone both repos. Only accept-workspace does; the other two clone once. The count is now strategy-derived, and the comment says plainly that it is the planned cost rather than a measurement, and what would have to change if git_ops ever gains a meter. Entitlements land as the seam the beta gate already models: inert unless both service vars are set, so a self-hoster resolves to unlimited and nothing is fetched. It fails OPEN where the invite gate fails closed, because a billing outage blocking a paying customer's merge is worse than minutes of unmetered usage, and the module says so where someone would otherwise "correct" it. Reads are cache-only; the network call is a separate refresh the auth path warms through waitUntil, so enforcement can never put a five-second dependency in front of a merge. Nothing consults it yet. Also: a squash merge test, because the note above it claimed that coverage and no test delivered it; and the UTC period test no longer mutates process.env.TZ, which is process-global while vitest runs several files per worker. Typecheck clean; lint clean; 3818 tests pass, up 30. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- migrations/049_usage_periods.sql | 49 ++- src/billing/entitlements.ts | 453 +++++++++++++++++++++++++++ src/middleware/auth.ts | 25 ++ src/middleware/config-guard.ts | 50 ++- src/routes/sync-management.ts | 28 +- src/storage/costs.ts | 36 ++- src/storage/usage.ts | 53 +++- src/types.ts | 11 + tests/config-guard.test.ts | 35 ++- tests/cost-attribution-paths.test.ts | 21 ++ tests/cost-attribution.test.ts | 21 ++ tests/entitlements.test.ts | 309 ++++++++++++++++++ tests/usage-periods.test.ts | 91 +++++- 13 files changed, 1122 insertions(+), 60 deletions(-) create mode 100644 src/billing/entitlements.ts create mode 100644 tests/entitlements.test.ts diff --git a/migrations/049_usage_periods.sql b/migrations/049_usage_periods.sql index f4545476..e35babd5 100644 --- a/migrations/049_usage_periods.sql +++ b/migrations/049_usage_periods.sql @@ -8,11 +8,19 @@ -- only the ACCOUNT cascades, so the one thing that clears usage is erasing the -- account that incurred it — at which point there is no subject left to bill. -- --- It is a derived aggregate, not a rival ledger. `cost_records` remains the --- per-change evidence of what was spent; this is the running total. Both are --- written by `recordCosts` (src/storage/costs.ts) from the same timestamp, so a --- cost row and the period it counts toward can never disagree, and a lost --- aggregate write is reconcilable from the ledger it derives from. +-- `cost_records` remains the per-change evidence of what was spent; this is the +-- running total. Both are written by `recordCosts` (src/storage/costs.ts) from +-- the same timestamp, so a cost row and the period it counts toward can never +-- disagree. +-- +-- This total is NOT recoverable from that ledger, which is the uncomfortable +-- half and the reason it is backed up in its own right (see BACKUP_TABLES in +-- src/storage/d1-backup.ts). The cascade that hard-deletes cost rows leaves +-- these standing, so replaying the ledger would rebuild usage only for projects +-- that still exist and hand every other owner their month back — the same +-- refund this table exists to prevent, arriving by a different route. A lost +-- upsert is therefore a permanent under-count, accepted because the alternative +-- is failing a change over an accounting write. -- -- Timestamps are ISO 8601 strings written by the application, NOT -- `datetime('now')`, for the reason 042/044/047 give: the two formats sort @@ -50,16 +58,37 @@ CREATE TABLE IF NOT EXISTS usage_periods ( -- cost that has demonstrably bitten once. An unrecognised meter is an inert -- row nothing reads; an unwritable one is a failed write on a hot path. meter TEXT NOT NULL, + -- Which provider account paid, carried into the aggregate rather than summed + -- away. An allowance is a limit on what the OPERATOR spends, so enforcement + -- reads the 'platform' rows alone: a project running the gate on its own key + -- costs the operator nothing and must not be charged against a hosted + -- allowance. Collapsing the two here would be a one-way door — this table is + -- not rebuildable from `cost_records` (see above), so any month accumulated + -- without the dimension could never have it recovered, and the usage page + -- could never show a project what its own key paid for. + source TEXT NOT NULL DEFAULT 'platform' CHECK(source IN ('platform','byok')), -- REAL, matching `cost_records.quantity`, which is what accumulates into it: -- sandbox milliseconds and character-derived token estimates are not -- guaranteed integral, and rounding at the aggregate would drift. - quantity REAL NOT NULL DEFAULT 0, + -- CHECK'd non-negative, because the UPSERT below adds rather than replaces + -- and nothing else stops a negative delta from REFUNDING the month — through + -- the one statement this whole table exists to make non-refundable. It is + -- reachable without malice: sandbox_ms is wall-clock arithmetic + -- (`Date.now() - startedAt` in post-merge.ts and sandbox-evaluator.ts), so a + -- backwards clock step yields one. Failing the write is correct here; the + -- caller treats accumulation as best-effort and an under-count is the safe + -- direction, where a refund is not. + quantity REAL NOT NULL DEFAULT 0 CHECK(quantity >= 0), updated_at TEXT NOT NULL, -- `owner_type` is deliberately outside the key: it describes the owner, it is -- not part of its identity. In the key it would let one account hold two rows -- for the same meter and month, and a limit compared against either alone - -- would be wrong. The key doubles as the UPSERT's conflict target - -- (src/storage/usage.ts) and as the covering index for the only read - -- (`WHERE owner_id = ? AND period = ?`), so no separate index is added. - PRIMARY KEY (owner_id, period, meter) + -- would be wrong. `source` IS in the key for the opposite reason — platform + -- and BYOK spend are two separately-answerable questions, and one limit + -- applies to only one of them. The key doubles as the UPSERT's conflict + -- target (src/storage/usage.ts), and SQLite's implicit index on it already + -- serves the owner+period lookup, so no separate index is added. (It is not a + -- *covering* index — `quantity` and `updated_at` sit outside the key — but the + -- search is on the key, which is what matters here.) + PRIMARY KEY (owner_id, period, meter, source) ); diff --git a/src/billing/entitlements.ts b/src/billing/entitlements.ts new file mode 100644 index 00000000..851ae300 --- /dev/null +++ b/src/billing/entitlements.ts @@ -0,0 +1,453 @@ +/** + * Entitlements — what one billing subject's plan allows (Stratum Cloud only). + * + * This is a cloud-only hook living in the AGPL tree, in the shape `src/beta/gate.ts` + * established: the plan service itself lives in the cloud layer, core just asks it. + * When `BILLING_SERVICE_URL` / `BILLING_SERVICE_SECRET` are unset — the default for + * every OSS self-hoster — every function here is inert and every owner resolves to + * `UnlimitedEntitlements`, which is exactly today's behavior. + * + * Three key sets, deliberately not one union, because they are not one kind of + * thing and do not live in one place: a `MeterKey` is a monthly *flow* accumulated + * in `usage_periods`, a `CountKey` is a *gauge* counted on demand, and a `RateKey` + * is a per-minute *rate* enforced out of KV. Collapsing them invites a limit to be + * looked up from the wrong store. + * + * ## Fail OPEN — deliberately the opposite of the beta gate + * + * `validateInviteCode` (`src/beta/gate.ts`) fails CLOSED on timeout, non-OK or a + * parse failure: a signup wall that opens under load is not a wall. Entitlements + * do the reverse. On any failure to reach or understand the billing service we + * serve the cached value, and with no cache `UnlimitedEntitlements`. A billing + * outage that blocks a paying customer's merge is strictly worse than minutes of + * unmetered usage; the hole is bounded by the cache TTL and is an accepted cost. + * + * Do not "correct" this to match the gate. The two hooks look alike and want + * opposite failure modes, which is why both say so in their own doc comments. + * + * ## Reading a limit + * + * `-1` means unlimited. **`0` means a hard block** — not "unset", not "unlimited"; + * an owner whose limit is `0` may not consume that resource at all. Both readings + * of `0` exist in the wild, so this file fixes one and callers can rely on it. + * Any other value that is not a non-negative integer — a string, `NaN`, `1.5`, + * `-7` — is DISCARDED in favor of the cached or default value for that field, and + * never coerced: a limit derived from garbage is worse than a limit we admit we + * do not have. + * + * NOTHING here enforces anything. Defining and fetching entitlements is separate + * from consulting them (PRD §8) precisely so the fetch can be observed for a month + * before it decides anything. + */ +import type { MeterKey } from "../storage/usage"; +import type { Env } from "../types"; +import { type AppError, ValidationError } from "../utils/errors"; +import type { Logger } from "../utils/logger"; +import { type Result, err, ok } from "../utils/result"; + +/** A gauge counted on demand, not a monthly flow — there is no `private_projects` meter. */ +export type CountKey = "private_projects"; + +/** A rate, enforced from KV by the request limiter rather than from a spend counter. */ +export type RateKey = "requests_per_minute"; + +/** Every meter a limit may be set on. Mirrors `MeterKey`; see `parseEntitlements`. */ +export const METER_KEYS: readonly MeterKey[] = [ + "llm_tokens_month", + "sandbox_ms_month", + "deploys_month", +]; + +export const COUNT_KEYS: readonly CountKey[] = ["private_projects"]; + +export const RATE_KEYS: readonly RateKey[] = ["requests_per_minute"]; + +/** + * One billing subject's allowances. + * + * Every limit is `-1` (unlimited), `0` (hard block) or a positive integer. The + * record is total over each key set on purpose: a missing key would make "no + * limit configured" and "limit of zero" indistinguishable at the call site. + */ +export interface Entitlements { + /** Opaque plan identifier as the billing service names it; for display and logs. */ + plan: string; + meters: Record; + counts: Record; + rates: Record; +} + +/** + * Where a resolved value came from. + * + * Carried because observe-only telemetry that cannot tell "unlimited because that + * is the plan" from "unlimited because the billing service timed out" cannot + * measure anything: a month of measurement would be indistinguishable from a + * month of outage. + */ +export type EntitlementsSource = "remote" | "cache" | "default"; + +export interface EntitlementsResolution { + entitlements: Entitlements; + source: EntitlementsSource; +} + +export interface EntitlementsProvider { + /** + * The allowances in force for one billing subject. + * + * A `Result` rather than a throw per AGENTS.md ("errors are values"). Because + * resolution fails open, the error arm is reserved for a caller mistake — an + * unidentified subject — which must not silently resolve to "unlimited". + */ + forOwner( + ownerId: string, + ownerType: "user" | "org", + ): Promise>; +} + +/** The value of a limit that is not a limit. */ +export const UNLIMITED = -1; + +/** The plan name reported when nobody has told us a real one. */ +const DEFAULT_PLAN = "unlimited"; + +/** + * The default in every sense: what a self-hoster always gets, and what a cloud + * request falls back to when the billing service cannot be reached and nothing + * is cached. Frozen because it is shared by every such resolution. + */ +export const UnlimitedEntitlements: Entitlements = Object.freeze({ + plan: DEFAULT_PLAN, + meters: Object.freeze({ + llm_tokens_month: UNLIMITED, + sandbox_ms_month: UNLIMITED, + deploys_month: UNLIMITED, + }), + counts: Object.freeze({ private_projects: UNLIMITED }), + rates: Object.freeze({ requests_per_minute: UNLIMITED }), +}); + +/** + * True only when the service is configured AND we hold a credential for it — + * mirroring `betaGateEnabled`. A URL without a secret is a misconfiguration that + * would produce a request the service rejects on every single hit, so it counts + * as off rather than as half-on. + */ +export function entitlementsEnabled(env: Env): boolean { + return !!env.BILLING_SERVICE_URL && !!env.BILLING_SERVICE_SECRET; +} + +/** + * Cap how long a warm can wait on the billing service. This runs off the response + * path via `waitUntil`, so the timeout is not protecting a user's latency — it is + * bounding how long a Worker invocation is kept alive by a hung dependency. + */ +const BILLING_TIMEOUT_MS = 5000; + +/** Positive cache lifetime (PRD §4). Consequence: an upgrade takes effect up to this late. */ +const CACHE_TTL_SECONDS = 300; + +/** + * Negative-cache and lock lifetime. Both are "as short as possible": KV's minimum + * `expirationTtl` is 60 seconds, so this is the floor, not a tuned value. The + * negative entry keeps a hard outage from being re-fetched once per request; the + * lock keeps a cold cache from stampeding the service. + */ +const FAILURE_TTL_SECONDS = 60; + +/** Versioned so a shape change cannot be read back as the old shape. */ +const KEY_PREFIX = "entitlements:v1:"; + +function cacheKey(ownerType: string, ownerId: string): string { + return `${KEY_PREFIX}${ownerType}:${ownerId}`; +} + +function negativeKey(ownerType: string, ownerId: string): string { + return `${KEY_PREFIX}neg:${ownerType}:${ownerId}`; +} + +function lockKey(ownerType: string, ownerId: string): string { + return `${KEY_PREFIX}lock:${ownerType}:${ownerId}`; +} + +function serviceUrl(env: Env, path: string): string { + return `${(env.BILLING_SERVICE_URL ?? "").replace(/\/$/, "")}${path}`; +} + +/** + * One limit, or `null` when the service sent something that is not one. + * + * `Number.isInteger` rejects strings, `NaN` and both infinities in one check, so + * a JSON `"100"` is discarded rather than becoming 100 — a limit we invented is + * indistinguishable, downstream, from one the customer is paying for. + */ +function parseLimit(value: unknown): number | null { + if (!Number.isInteger(value)) return null; + const limit = value as number; + if (limit < 0 && limit !== UNLIMITED) return null; + return limit; +} + +/** + * Build entitlements field by field, keeping `fallback`'s value for anything the + * payload does not supply validly. + * + * Field by field rather than all-or-nothing so one bad key does not throw away + * the rest of a customer's plan — and never partially: every key is present in + * the result because it was present in `fallback`. + */ +function parseEntitlements(payload: unknown, fallback: Entitlements): Entitlements { + const record = (typeof payload === "object" && payload !== null ? payload : {}) as Record< + string, + unknown + >; + const readGroup = ( + group: unknown, + keys: readonly K[], + defaults: Record, + ): Record => { + const source = (typeof group === "object" && group !== null ? group : {}) as Record< + string, + unknown + >; + const out = {} as Record; + for (const key of keys) { + out[key] = parseLimit(source[key]) ?? defaults[key]; + } + return out; + }; + + const plan = typeof record.plan === "string" && record.plan.trim() ? record.plan.trim() : null; + return { + plan: plan ?? fallback.plan, + meters: readGroup(record.meters, METER_KEYS, fallback.meters), + counts: readGroup(record.counts, COUNT_KEYS, fallback.counts), + rates: readGroup(record.rates, RATE_KEYS, fallback.rates), + }; +} + +/** + * Entitlements from the cloud billing service, cached in KV. + * + * `forOwner` is a CACHED READ and never touches the network: it is called on the + * request path, and a cold cache there would put a 5-second dependency in front + * of a merge. The network fetch happens in `refresh`, which `warmEntitlements` + * runs off the response path after auth — that is what makes the cached read hit. + */ +export class RemoteEntitlements implements EntitlementsProvider { + constructor( + private readonly env: Env, + private readonly logger: Logger, + ) {} + + async forOwner( + ownerId: string, + ownerType: "user" | "org", + ): Promise> { + if (!ownerId) { + // Not a fail-open case: an unidentified subject is a bug in the caller, and + // answering "unlimited" would hide it behind exactly the behavior it breaks. + return err(new ValidationError("Entitlements requested without an owner id", { ownerType })); + } + if (!entitlementsEnabled(this.env)) { + return ok({ entitlements: UnlimitedEntitlements, source: "default" }); + } + const cached = await this.readCache(ownerId, ownerType); + if (cached) return ok({ entitlements: cached, source: "cache" }); + return ok({ entitlements: UnlimitedEntitlements, source: "default" }); + } + + /** + * Ask the billing service and populate the cache. + * + * Fails open (see the module doc comment): a timeout, a non-OK status or an + * unparseable body all serve the cached value, or `UnlimitedEntitlements` when + * there is none, and record a negative-cache entry so an outage is not + * re-fetched once per request. + */ + async refresh( + ownerId: string, + ownerType: "user" | "org", + ): Promise> { + if (!ownerId) { + return err(new ValidationError("Entitlements refresh without an owner id", { ownerType })); + } + if (!entitlementsEnabled(this.env)) { + return ok({ entitlements: UnlimitedEntitlements, source: "default" }); + } + + const cached = await this.readCache(ownerId, ownerType); + const fallback: EntitlementsResolution = cached + ? { entitlements: cached, source: "cache" } + : { entitlements: UnlimitedEntitlements, source: "default" }; + + try { + const url = `${serviceUrl(this.env, "/api/billing/entitlements")}?ownerId=${encodeURIComponent( + ownerId, + )}&ownerType=${encodeURIComponent(ownerType)}`; + const res = await fetch(url, { + headers: { Authorization: `Bearer ${this.env.BILLING_SERVICE_SECRET ?? ""}` }, + signal: AbortSignal.timeout(BILLING_TIMEOUT_MS), + }); + if (!res.ok) { + this.logger.warn("Entitlements lookup returned non-OK; failing open", { + status: res.status, + ownerId, + source: fallback.source, + }); + await this.markFailure(ownerId, ownerType); + return ok(fallback); + } + const entitlements = parseEntitlements(await res.json(), fallback.entitlements); + await this.writeCache(ownerId, ownerType, entitlements); + return ok({ entitlements, source: "remote" }); + } catch (error) { + // Never swallowed: an outage that is invisible is one nobody fixes, and the + // free usage it hands out is only acceptable while it is being watched. + this.logger.error( + "Entitlements lookup failed; failing open", + error instanceof Error ? error : undefined, + { ownerId, ownerType, source: fallback.source }, + ); + await this.markFailure(ownerId, ownerType); + return ok(fallback); + } + } + + private async readCache(ownerId: string, ownerType: string): Promise { + if (!this.env.STATE) return null; + try { + const raw = await this.env.STATE.get(cacheKey(ownerType, ownerId)); + if (!raw) return null; + // Re-validated on the way out, not trusted because we wrote it: the cache + // outlives deploys, so an entry written by an older shape must degrade to + // the default rather than be cast into a limit nobody validated. + return parseEntitlements(JSON.parse(raw), UnlimitedEntitlements); + } catch (error) { + this.logger.warn("Entitlements cache read failed; treating as a miss", { + ownerId, + error: error instanceof Error ? error.message : String(error), + }); + return null; + } + } + + private async writeCache( + ownerId: string, + ownerType: string, + entitlements: Entitlements, + ): Promise { + if (!this.env.STATE) return; + try { + await this.env.STATE.put(cacheKey(ownerType, ownerId), JSON.stringify(entitlements), { + expirationTtl: CACHE_TTL_SECONDS, + }); + await this.env.STATE.delete(negativeKey(ownerType, ownerId)); + } catch (error) { + this.logger.warn("Entitlements cache write failed", { + ownerId, + error: error instanceof Error ? error.message : String(error), + }); + } + } + + /** + * Record that the service just failed, in its OWN key: a failure must never + * overwrite a good cached plan, or one blip would demote a paying customer to + * the default for a full TTL. + */ + private async markFailure(ownerId: string, ownerType: string): Promise { + if (!this.env.STATE) return; + try { + await this.env.STATE.put(negativeKey(ownerType, ownerId), String(Date.now()), { + expirationTtl: FAILURE_TTL_SECONDS, + }); + } catch { + // A failure to record a failure changes nothing about the answer we already + // have; the fetch failure itself was logged above. + } + } +} + +/** + * Refreshes in flight in THIS isolate, keyed by cache key. + * + * Two layers of single-flight because there are two stampedes. KV has no + * compare-and-set, so its lock is best-effort and cross-isolate only — two + * concurrent requests in one isolate can both read a null lock before either + * writes it. This map closes that window exactly where it is widest: one isolate + * handling a burst of requests for the same owner. + */ +const inFlightWarms = new Map>(); + +/** + * Warm one owner's entitlements cache off the response path. + * + * Called after auth resolves an identity, so `forOwner`'s cached read hits on + * subsequent requests — without this the cached read would be permanently cold + * and every enforcement site would see `source: "default"` forever. + * + * Fire-and-forget by construction: it schedules on `waitUntil` and returns + * nothing. With no `ExecutionContext` (unit tests, non-Workers runtimes) it does + * NOT fall back to awaiting inline — unlike the token `last_used` touches nearby, + * which are cheap local writes. This one is a cross-network call, and putting it + * on the request path is the exact latency this design avoids. Skipping it is + * safe: readers fail open. + */ +export function warmEntitlements( + env: Env, + waitUntil: ((promise: Promise) => void) | undefined, + subject: { ownerId: string; ownerType: "user" | "org" }, + logger: Logger, +): void { + if (!waitUntil) return; + if (!subject.ownerId) return; + if (!entitlementsEnabled(env)) return; + if (!env.STATE) return; + + const key = cacheKey(subject.ownerType, subject.ownerId); + const existing = inFlightWarms.get(key); + if (existing) { + waitUntil(existing); + return; + } + const run = warmOnce(env, subject, logger).finally(() => { + inFlightWarms.delete(key); + }); + inFlightWarms.set(key, run); + waitUntil(run); +} + +async function warmOnce( + env: Env, + subject: { ownerId: string; ownerType: "user" | "org" }, + logger: Logger, +): Promise { + const { ownerId, ownerType } = subject; + const kv = env.STATE; + try { + if (await kv.get(cacheKey(ownerType, ownerId))) return; + // A recent failure is respected rather than retried: re-fetching per request + // during an outage is how a fail-open dependency turns into a load generator. + if (await kv.get(negativeKey(ownerType, ownerId))) return; + if (await kv.get(lockKey(ownerType, ownerId))) return; + await kv.put(lockKey(ownerType, ownerId), "1", { expirationTtl: FAILURE_TTL_SECONDS }); + } catch (error) { + // The lock is an optimization. Losing it must not lose the warm, exactly as + // the backup runner reasons about its own KV lock. + logger.warn("Entitlements warm lock unavailable; proceeding without single-flight", { + ownerId, + error: error instanceof Error ? error.message : String(error), + }); + } + + await new RemoteEntitlements(env, logger).refresh(ownerId, ownerType); + + try { + // Released as soon as the answer is cached; the TTL is only the crash path. + await kv.delete(lockKey(ownerType, ownerId)); + } catch { + // Held locks expire on their own. A stuck one costs one TTL of staleness. + } +} diff --git a/src/middleware/auth.ts b/src/middleware/auth.ts index ace96837..6bcb0d1c 100644 --- a/src/middleware/auth.ts +++ b/src/middleware/auth.ts @@ -1,5 +1,6 @@ import type { MiddlewareHandler } from "hono"; import { getCookie } from "hono/cookie"; +import { warmEntitlements } from "../billing/entitlements"; import { isGitHttpPath } from "../routes/git-http"; import { isOAuthClientEndpoint } from "../routes/mcp-oauth"; import { getAgentByToken } from "../storage/agents"; @@ -95,6 +96,22 @@ export const authMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, ne c.set("logger", logger); + /** + * Pre-fetch the caller's plan limits for LATER requests (PRD §4). + * + * Called at every point below where an identity is established, because this + * is the only place that knows who is paying before the routes run. It is + * inert unless the cloud billing vars are set, never awaited (it schedules on + * `waitUntil` and is skipped outright when there is none), and nothing in this + * request reads the result — a warm that fails costs a later reader nothing + * but a fail-open default. + * + * An agent warms its OWNER, not itself: agents are not billing subjects. + */ + const warmOwner = (ownerId: string): void => { + warmEntitlements(c.env, getWaitUntil(c), { ownerId, ownerType: "user" }, logger); + }; + // The git smart-HTTP router authenticates over HTTP Basic itself; let it own // the challenge instead of rejecting the non-Bearer header here. if (isGitHttpPath(c.req.path)) { @@ -266,6 +283,8 @@ export const authMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, ne else await touch; } + warmOwner(user.id); + logger.debug("Auth success - user", { userId: user.id, username: user.username, scope }); await next(); return; @@ -307,6 +326,8 @@ export const authMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, ne // choice governs it — otherwise opting out could be defeated by routing // traffic through an agent. The owner row is already in hand. c.set("telemetryOptOut", ownerResult.data.telemetryOptOut === true); + warmOwner(agentResult.data.ownerId); + logger.debug("Auth success - agent", { agentId: agentResult.data.id, ownerId: agentResult.data.ownerId, @@ -395,6 +416,8 @@ export const authMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, ne if (waitUntil) waitUntil(touch); else await touch; + warmOwner(grant.data.user.id); + logger.debug("Auth success - OAuth grant", { userId: grant.data.user.id, clientId: grant.data.clientId, @@ -450,6 +473,8 @@ export const authMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, ne userResult.data.username || (userResult.data.email.split("@")[0] ?? "").toLowerCase().replace(/[^a-z0-9]/g, ""); c.set("username", username); + warmOwner(sessionResult.data.userId); + logger.debug("Auth success - session", { userId: sessionResult.data.userId, username }); } } else { diff --git a/src/middleware/config-guard.ts b/src/middleware/config-guard.ts index e3e6d8de..fae4ccb5 100644 --- a/src/middleware/config-guard.ts +++ b/src/middleware/config-guard.ts @@ -26,21 +26,55 @@ export function repoDoConfigError( return null; } -// Log the config problem at most once per isolate rather than on every request. +/** + * Returns a human-readable error when entitlement enforcement is switched on + * without the service it enforces against, else null. + * + * Structurally the same footgun as `REPO_DO_ENABLED` above, one step quieter: + * `ENTITLEMENTS_ENFORCE=1` with no `BILLING_SERVICE_URL` throws nothing and + * breaks nothing — every owner resolves to `UnlimitedEntitlements` and every + * decision admits. That is enforcement that looks on in the dashboard and does + * nothing in production, which is worse than being off, because nobody goes + * looking for a limit they believe is already applied. + * + * `BILLING_SERVICE_SECRET` is not checked here: it is a Wrangler secret rather + * than a var, so a deploy that has the URL and not the secret is an operational + * state to fix, not a config file to reread — and `entitlementsEnabled` already + * treats it as off. + */ +export function entitlementsConfigError( + env: Pick, +): string | null { + if (env.ENTITLEMENTS_ENFORCE === "1" && !env.BILLING_SERVICE_URL) { + return ( + "ENTITLEMENTS_ENFORCE is '1' but BILLING_SERVICE_URL is not set — no plan " + + "limits can be fetched, so every owner resolves to unlimited and every " + + "enforcement point admits. Set BILLING_SERVICE_URL (and the " + + "BILLING_SERVICE_SECRET secret) for the [env.] block, or unset " + + "ENTITLEMENTS_ENFORCE." + ); + } + return null; +} + +// Log the config problems at most once per isolate rather than on every request. let hasLoggedConfigError = false; /** - * Surfaces an incoherent RepoDO/R2 configuration loudly in Workers Logs on the - * first request after a bad deploy, instead of only when a later merge throws. - * Non-fatal: reads and the UI still work without the bucket, so we log rather - * than reject the request. + * Surfaces an incoherent RepoDO/R2 or entitlements configuration loudly in + * Workers Logs on the first request after a bad deploy, instead of only when a + * later merge throws — or, for entitlements, never. Non-fatal: reads and the UI + * still work in both cases, so we log rather than reject the request. */ export const configGuardMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, next) => { if (!hasLoggedConfigError) { - const problem = repoDoConfigError(c.env); - if (problem) { + const problems = [repoDoConfigError(c.env), entitlementsConfigError(c.env)].filter( + (problem): problem is string => problem !== null, + ); + if (problems.length > 0) { hasLoggedConfigError = true; - createLogger({ component: "config-guard" }).error(problem); + const logger = createLogger({ component: "config-guard" }); + for (const problem of problems) logger.error(problem); } } await next(); diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index 454a827e..925e6e22 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -783,17 +783,23 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { logger, ); - // `resolveConflict` clones the project AND the workspace fork and pushes, for - // EVERY strategy — so this is recorded here rather than in the `manual` branch - // above, which only ever sees the evaluator suite's spend. Without it, - // accept-project and accept-workspace resolutions record nothing at all, and - // manual under-records by the clone+push it just paid for. Two git_ops, the - // same shape the single-merge route bills for the same clone+push. + // `resolveConflict` runs for EVERY strategy, so this is recorded here rather + // than in the `manual` branch above, which only ever sees the evaluator + // suite's spend. Without it, accept-project and accept-workspace resolutions + // record nothing at all, and manual under-records the git work it paid for. // - // Before the failure branch below, for the same reason the manual branch - // records before its 422: the clones are already paid for when the push - // fails, and a failure that costs the operator the same as a success must not - // be the cheap way to use the platform. + // Recorded before the failure branch below, for the same reason the manual + // branch records before its 422: a resolution that clones and then fails to + // push cost the operator what a successful one did, and a failure must not be + // the cheap way to use the platform. + // + // The count is what the strategy WOULD do, not a measurement — `resolveConflict` + // reports only a commit sha, so there is nothing to measure. It therefore + // over-bills a resolution that dies before its first clone (a rejected token) + // and is exact otherwise. Tolerable while `git_ops` maps to no meter + // (`meterForCostKind`); if it ever gains one, `resolveConflict` has to report + // the operations it actually completed rather than being guessed at here. + const resolveGitOps = strategy === "accept-workspace" ? 3 : 2; const resolveGitSubject = await resolveBillingSubject(c.env.DB, logger, project); await recordCosts( c.env.DB, @@ -805,7 +811,7 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { workspace: conflictCtx.workspaceName, ...(resolveGitSubject ?? {}), }, - [{ kind: "git_ops", quantity: 2 }], + [{ kind: "git_ops", quantity: resolveGitOps }], ); if (!resolveResult.success) { diff --git a/src/storage/costs.ts b/src/storage/costs.ts index 9fc2663f..777474bd 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -106,7 +106,7 @@ export async function resolveBillingSubject( /** * Accumulate the metered half of a batch into `usage_periods`. * - * Split out of `recordCosts` so all six recording sites get it without any of + * Split out of `recordCosts` so every recording site gets it without any of * them knowing it exists — the aggregate must never be something a call site * can forget, because a site that forgets is an owner whose spend does not * count against their allowance. @@ -132,13 +132,21 @@ async function accumulateUsage( for (const sample of samples) { // `git_ops` maps to no meter; see `meterForCostKind`. const meter = meterForCostKind(sample.kind); - if (meter) deltas.push({ meter, quantity: sample.quantity }); + if (meter) { + deltas.push({ meter, quantity: sample.quantity, source: sample.source ?? "platform" }); + } } if (deltas.length === 0) return; // The Result is deliberately not folded into `recordCosts`'s own: the ledger // rows above ARE written, and reporting an error would tell the caller the - // spend went unrecorded when it did not. `upsertUsage` logs its own failure, - // and the aggregate is reconcilable from the ledger it derives from. + // spend went unrecorded when it did not. `upsertUsage` logs its own failure. + // + // Be clear about what that costs, because it is not free and there is no job + // that repairs it: this aggregate cannot be rebuilt from `cost_records` (the + // project cascade deletes those rows and leaves these standing — see + // migration 049), so a failed upsert is a permanent under-count of that + // owner's month. Accepted because the alternative is failing a merge over an + // accounting write, which is the worse trade in both directions. await upsertUsage( db, logger, @@ -170,6 +178,22 @@ export async function recordCosts( samples: CostSample[], ): Promise> { if (samples.length === 0) return ok(undefined); + + // Dropped BEFORE the batch, not inside the aggregate that also guards them. + // `db.batch` is one transaction and `cost_records.quantity` is NOT NULL, + // while SQLite stores NaN as NULL — so a single unusable sample does not cost + // itself, it aborts the write and loses every other sample in the same + // evaluation. Filtering here keeps one bad number from taking the ledger with + // it, and is why `upsertUsage`'s own guard is a backstop rather than the + // first line of defence. + const usable = samples.filter((sample) => Number.isFinite(sample.quantity)); + if (usable.length !== samples.length) { + logger.warn("Cost samples dropped: quantity is not finite", { + project: opts.project, + dropped: samples.length - usable.length, + }); + } + if (usable.length === 0) return ok(undefined); const createdAt = new Date().toISOString(); try { @@ -177,7 +201,7 @@ export async function recordCosts( "INSERT INTO cost_records (id, project, project_id, change_id, workspace, kind, quantity, estimated, created_at, owner_id, owner_type, source) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ); await db.batch( - samples.map((sample) => + usable.map((sample) => stmt.bind( newId("cost"), opts.project, @@ -203,7 +227,7 @@ export async function recordCosts( }); // Only after the ledger write succeeded: the aggregate derives from those // rows, so it must never count spend the ledger does not record. - await accumulateUsage(db, logger, opts, samples, createdAt); + await accumulateUsage(db, logger, opts, usable, createdAt); return ok(undefined); } catch (error) { const appError = diff --git a/src/storage/usage.ts b/src/storage/usage.ts index 0e0b22f5..e0dc1c59 100644 --- a/src/storage/usage.ts +++ b/src/storage/usage.ts @@ -1,13 +1,13 @@ import { AppError } from "../utils/errors"; import type { Logger } from "../utils/logger"; import { type Result, err, ok } from "../utils/result"; -import type { BillingSubject, CostKind } from "./costs"; +import type { BillingSubject, CostKind, CostSource } from "./costs"; /** * A metered quantity an allowance can be set against, over a calendar month. * * The `_month` suffix is redundant with the row's `period` column and is kept - * anyway: these are the exact strings entitlements keys its limits on, so a + * anyway: these are the strings an entitlement's limits are keyed on, so a * limit lookup is `limits[row.meter]` rather than a second mapping table that * could disagree with this one. * @@ -19,10 +19,18 @@ export type MeterKey = "llm_tokens_month" | "sandbox_ms_month" | "deploys_month" export interface UsageDelta { meter: MeterKey; quantity: number; + /** + * Whose provider account paid. Omitted means `"platform"` — the operator's, + * and so the only kind an allowance limits. Kept as a separate total rather + * than summed in, because a limit applies to one of the two and not the + * other; see migration 049. + */ + source?: CostSource; } export interface UsageSummaryEntry { meter: MeterKey; + source: CostSource; quantity: number; /** ISO 8601, app-written. When this owner last consumed this meter. */ updatedAt: string; @@ -30,6 +38,7 @@ export interface UsageSummaryEntry { interface UsageRow { meter: string; + source: string; quantity: number; updated_at: string; } @@ -94,7 +103,8 @@ export function usagePeriod(at: Date = new Date()): string { * * Deliberately not `db.batch`: the meters are independent, so cross-meter * atomicity buys nothing, and one statement per meter keeps the atomicity claim - * at the statement where it is actually true. Callers send at most two. + * at the statement where it is actually true. However many deltas arrive, only + * as many statements run as there are distinct (meter, source) pairs. * * Best-effort by contract, like `recordCosts` which drives it: this returns a * Result and logs, and never throws into change creation, merge or deploy. @@ -111,34 +121,48 @@ export async function upsertUsage( period: string, deltas: readonly UsageDelta[], ): Promise> { - const merged = new Map(); + // Keyed on both dimensions the row is keyed on, so platform and BYOK spend + // for one meter accumulate as two totals rather than one. + const merged = new Map(); for (const delta of deltas) { // A NaN or Infinity would poison the running total permanently: SQLite has // no way back from it and every later comparison against a limit would be - // false. Drop the sample and keep the meter usable. - if (!Number.isFinite(delta.quantity)) { - logger.warn("Usage delta skipped: quantity is not finite", { + // false. A NEGATIVE would do something worse than poison it — it would + // refund the month, through the single statement that exists to stop the + // month being refundable. Wall-clock arithmetic makes that reachable + // without malice (`Date.now() - startedAt` under a backwards clock step). + // Drop either and keep the meter usable; migration 049's CHECK is the + // backstop if one ever gets past here. + if (!Number.isFinite(delta.quantity) || delta.quantity < 0) { + logger.warn("Usage delta skipped: quantity is not a usable increment", { ownerId: subject.ownerId, meter: delta.meter, quantity: delta.quantity, }); continue; } - merged.set(delta.meter, (merged.get(delta.meter) ?? 0) + delta.quantity); + const source: CostSource = delta.source ?? "platform"; + const key = `${delta.meter}\u0000${source}`; + const existing = merged.get(key); + merged.set(key, { + meter: delta.meter, + source, + quantity: (existing?.quantity ?? 0) + delta.quantity, + }); } if (merged.size === 0) return ok(undefined); const updatedAt = new Date().toISOString(); try { - for (const [meter, quantity] of merged) { + for (const { meter, source, quantity } of merged.values()) { await db .prepare( - "INSERT INTO usage_periods (owner_id, owner_type, period, meter, quantity, updated_at) " + - "VALUES (?, ?, ?, ?, ?, ?) " + - "ON CONFLICT(owner_id, period, meter) DO UPDATE SET " + + "INSERT INTO usage_periods (owner_id, owner_type, period, meter, source, quantity, updated_at) " + + "VALUES (?, ?, ?, ?, ?, ?, ?) " + + "ON CONFLICT(owner_id, period, meter, source) DO UPDATE SET " + "quantity = quantity + excluded.quantity, updated_at = excluded.updated_at", ) - .bind(subject.ownerId, subject.ownerType, period, meter, quantity, updatedAt) + .bind(subject.ownerId, subject.ownerType, period, meter, source, quantity, updatedAt) .run(); } logger.debug("Usage aggregate updated", { @@ -185,7 +209,7 @@ export async function getOwnerUsageSummary( try { const result = await db .prepare( - "SELECT meter, quantity, updated_at FROM usage_periods WHERE owner_id = ? AND period = ? ORDER BY meter", + "SELECT meter, source, quantity, updated_at FROM usage_periods WHERE owner_id = ? AND period = ? ORDER BY meter, source", ) .bind(ownerId, period) .all(); @@ -195,6 +219,7 @@ export async function getOwnerUsageSummary( // migration 049), so a row written by a future meter this build does // not know about reads through rather than failing the whole summary. meter: row.meter as MeterKey, + source: row.source as CostSource, quantity: row.quantity, updatedAt: row.updated_at, })), diff --git a/src/types.ts b/src/types.ts index 789986fe..96a463b3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -154,6 +154,17 @@ export interface Env { BETA_GATE?: string; REFERRAL_SERVICE_URL?: string; REFERRAL_SERVICE_SECRET?: string; + /** Plan/entitlements service (Stratum Cloud only; OSS self-hosters leave these + * unset -> every owner resolves to `UnlimitedEntitlements` and nothing is + * fetched). BOTH are required for the hook to be live — see + * `entitlementsEnabled` in src/billing/entitlements.ts. */ + BILLING_SERVICE_URL?: string; + BILLING_SERVICE_SECRET?: string; + /** "1" makes entitlement decisions REFUSE rather than only record. Off by + * default: the limits are measured for a period before they gate anything. + * Setting it without BILLING_SERVICE_URL is enforcement that looks on and does + * nothing — `entitlementsConfigError` flags exactly that. */ + ENTITLEMENTS_ENFORCE?: string; ANALYTICS?: AnalyticsEngineDataset; SANDBOX?: SandboxBinding; AI?: AiBinding; diff --git a/tests/config-guard.test.ts b/tests/config-guard.test.ts index b6e5db23..c70a5b39 100644 --- a/tests/config-guard.test.ts +++ b/tests/config-guard.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { repoDoConfigError } from "../src/middleware/config-guard"; +import { entitlementsConfigError, repoDoConfigError } from "../src/middleware/config-guard"; import type { Env } from "../src/types"; const bucket = {} as Env["REPO_OBJECTS"]; @@ -20,3 +20,36 @@ describe("repoDoConfigError", () => { expect(repoDoConfigError({ REPO_DO_ENABLED: undefined, REPO_OBJECTS: undefined })).toBeNull(); }); }); + +describe("entitlementsConfigError", () => { + it("flags ENTITLEMENTS_ENFORCE='1' with no billing service URL", () => { + const problem = entitlementsConfigError({ + ENTITLEMENTS_ENFORCE: "1", + BILLING_SERVICE_URL: undefined, + }); + expect(problem).toContain("ENTITLEMENTS_ENFORCE"); + expect(problem).toContain("BILLING_SERVICE_URL"); + }); + + it("is silent when enforcement is on and the service is configured", () => { + expect( + entitlementsConfigError({ + ENTITLEMENTS_ENFORCE: "1", + BILLING_SERVICE_URL: "https://billing.test", + }), + ).toBeNull(); + }); + + it("is silent when enforcement is off, configured or not (the default)", () => { + expect( + entitlementsConfigError({ ENTITLEMENTS_ENFORCE: undefined, BILLING_SERVICE_URL: undefined }), + ).toBeNull(); + expect( + entitlementsConfigError({ ENTITLEMENTS_ENFORCE: "0", BILLING_SERVICE_URL: undefined }), + ).toBeNull(); + // Only the exact "1" arms enforcement, so "true" is off — and off is coherent. + expect( + entitlementsConfigError({ ENTITLEMENTS_ENFORCE: "true", BILLING_SERVICE_URL: undefined }), + ).toBeNull(); + }); +}); diff --git a/tests/cost-attribution-paths.test.ts b/tests/cost-attribution-paths.test.ts index 95729e58..1dd8cf0d 100644 --- a/tests/cost-attribution-paths.test.ts +++ b/tests/cost-attribution-paths.test.ts @@ -391,6 +391,27 @@ describe("POST /api/changes/:id/merge records attributed costs", () => { }); }); + it("still records for ?strategy=squash, which never takes the queue branch", async () => { + // The queue branch is gated on `strategy === "merge"`, so squash reaches the + // cold path even with the binding present. Written because the note above + // claims this coverage. + const queueEnv = makeEnv({ + MERGE_QUEUE: { + idFromName: (name: string) => name, + get: () => ({ + merge: async () => { + throw new Error("squash must not reach the merge queue"); + }, + }), + } as unknown as Env["MERGE_QUEUE"], + }); + + const res = await makeApp(changesRouter).fetch(mergeRequest("?strategy=squash"), queueEnv); + + expect(res.status).toBe(200); + expect(costRows()[0]).toMatchObject({ kind: "git_ops", owner_id: "user_alice" }); + }); + it("records NOTHING when the merge queue is bound, which is every deployed config", async () => { // Asserting the gap rather than the fix, deliberately. Closing it means // recording inside the queue branch, which is its own change; until then a diff --git a/tests/cost-attribution.test.ts b/tests/cost-attribution.test.ts index 5d9a8997..c63149a1 100644 --- a/tests/cost-attribution.test.ts +++ b/tests/cost-attribution.test.ts @@ -262,6 +262,27 @@ describe("recordCosts attribution columns", () => { expect(llm?.total).toBe(750); }); + it("drops one unusable sample instead of losing the whole batch to it", async () => { + // db.batch is a transaction and quantity is NOT NULL, while SQLite stores + // NaN as NULL — so an unfiltered NaN would not cost itself, it would abort + // the write and take every other sample in the same evaluation with it. + const { db, raw } = makeSqliteD1(); + const result = await recordCosts( + db, + logger, + { project: "p", changeId: "chg_1", ownerId: "user_alice", ownerType: "user" }, + [ + { kind: "llm_tokens", quantity: Number.NaN, estimated: true }, + { kind: "git_ops", quantity: 2 }, + ], + ); + + expect(result.success).toBe(true); + const kept = rows(raw); + expect(kept).toHaveLength(1); + expect(kept[0]?.kind).toBe("git_ops"); + }); + it("serves an owner-scoped, newest-first history from the new index", async () => { // What idx_costs_owner exists for. `created_at` is written only by // recordCosts as an ISO 8601 string, so ordering it in SQL is safe here — diff --git a/tests/entitlements.test.ts b/tests/entitlements.test.ts new file mode 100644 index 00000000..9dac847f --- /dev/null +++ b/tests/entitlements.test.ts @@ -0,0 +1,309 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { + type Entitlements, + RemoteEntitlements, + UnlimitedEntitlements, + entitlementsEnabled, + warmEntitlements, +} from "../src/billing/entitlements"; +import type { Env } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { makeFakeKV } from "./helpers/fake-kv"; + +const noopLogger: Logger = { + trace: () => {}, + debug: () => {}, + info: () => {}, + warn: () => {}, + error: () => {}, + fatal: () => {}, + child: () => noopLogger, +}; + +type FakeKV = ReturnType; + +function env(overrides: Partial = {}): { env: Env; kv: FakeKV } { + const kv = makeFakeKV(); + return { + env: { + STATE: kv, + BILLING_SERVICE_URL: "https://billing.test/", + BILLING_SERVICE_SECRET: "shh", + ...overrides, + } as Env, + kv, + }; +} + +/** Configured with a KV binding but no billing service — the self-hoster default. */ +function unconfiguredEnv(): Env { + return env({ BILLING_SERVICE_URL: undefined, BILLING_SERVICE_SECRET: undefined }).env; +} + +const CACHE_KEY = "entitlements:v1:user:usr_1"; +const NEGATIVE_KEY = "entitlements:v1:neg:user:usr_1"; +const LOCK_KEY = "entitlements:v1:lock:user:usr_1"; + +const paidPlan: Entitlements = { + plan: "pro", + meters: { llm_tokens_month: 1_000_000, sandbox_ms_month: 60_000, deploys_month: 100 }, + counts: { private_projects: 25 }, + rates: { requests_per_minute: 2_000 }, +}; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { status }); +} + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("entitlementsEnabled", () => { + it("is off by default (self-hoster, no env)", () => { + expect(entitlementsEnabled({} as Env)).toBe(false); + }); + + it("is off when only the URL is set", () => { + expect(entitlementsEnabled({ BILLING_SERVICE_URL: "https://billing.test" } as Env)).toBe(false); + }); + + it("is off when only the secret is set", () => { + expect(entitlementsEnabled({ BILLING_SERVICE_SECRET: "shh" } as Env)).toBe(false); + }); + + it("is on only when both are configured", () => { + expect( + entitlementsEnabled({ + BILLING_SERVICE_URL: "https://billing.test", + BILLING_SERVICE_SECRET: "shh", + } as Env), + ).toBe(true); + }); +}); + +describe("RemoteEntitlements.forOwner", () => { + it("returns the unlimited default, tagged 'default', when the hook is unconfigured", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const provider = new RemoteEntitlements(unconfiguredEnv(), noopLogger); + const result = await provider.forOwner("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + expect(result.success && result.data.entitlements).toEqual(UnlimitedEntitlements); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("errors rather than silently answering 'unlimited' for an unidentified owner", async () => { + const { env: e } = env(); + const result = await new RemoteEntitlements(e, noopLogger).forOwner("", "user"); + expect(result.success).toBe(false); + expect(!result.success && result.error.code).toBe("VALIDATION_ERROR"); + }); + + it("serves a warm cache and tags it 'cache'", async () => { + const { env: e, kv } = env(); + kv.store.set(CACHE_KEY, JSON.stringify(paidPlan)); + const result = await new RemoteEntitlements(e, noopLogger).forOwner("usr_1", "user"); + expect(result.success && result.data.source).toBe("cache"); + expect(result.success && result.data.entitlements).toEqual(paidPlan); + }); + + it("is a cached read only — a cold cache never reaches the network", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { env: e } = env(); + const result = await new RemoteEntitlements(e, noopLogger).forOwner("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + expect(result.success && result.data.entitlements).toEqual(UnlimitedEntitlements); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("treats an unparseable cache entry as a miss", async () => { + const { env: e, kv } = env(); + kv.store.set(CACHE_KEY, "{not json"); + const result = await new RemoteEntitlements(e, noopLogger).forOwner("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + }); +}); + +describe("RemoteEntitlements.refresh", () => { + it("returns the remote value tagged 'remote' and caches it", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(jsonResponse(paidPlan)); + const { env: e, kv } = env(); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.source).toBe("remote"); + expect(result.success && result.data.entitlements).toEqual(paidPlan); + expect(JSON.parse(String(kv.store.get(CACHE_KEY)))).toEqual(paidPlan); + + const url = String(fetchSpy.mock.calls[0]?.[0]); + expect(url).toBe("https://billing.test/api/billing/entitlements?ownerId=usr_1&ownerType=user"); + const headers = fetchSpy.mock.calls[0]?.[1]?.headers as Record; + expect(headers.Authorization).toBe("Bearer shh"); + }); + + it("keeps 0 as a hard block and discards string / NaN / negative limits", async () => { + vi.spyOn(globalThis, "fetch").mockResolvedValue( + jsonResponse({ + plan: "starter", + meters: { + llm_tokens_month: 0, + sandbox_ms_month: "50000", + deploys_month: Number.NaN, + }, + counts: { private_projects: -7 }, + rates: { requests_per_minute: 1.5 }, + }), + ); + const { env: e } = env(); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.source).toBe("remote"); + expect(result.data.entitlements).toEqual({ + plan: "starter", + // 0 survives: it is a hard block, not "unset". + meters: { llm_tokens_month: 0, sandbox_ms_month: -1, deploys_month: -1 }, + counts: { private_projects: -1 }, + rates: { requests_per_minute: -1 }, + }); + }); + + it("discards bad fields in favour of the CACHED value, not the default", async () => { + const { env: e, kv } = env(); + kv.store.set(CACHE_KEY, JSON.stringify(paidPlan)); + vi.spyOn(globalThis, "fetch").mockResolvedValue( + jsonResponse({ meters: { llm_tokens_month: "lots", sandbox_ms_month: 10 } }), + ); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.entitlements).toEqual({ + ...paidPlan, + meters: { ...paidPlan.meters, sandbox_ms_month: 10 }, + }); + }); + + it("fails OPEN to the warm cache on a non-OK response, without clobbering it", async () => { + const { env: e, kv } = env(); + kv.store.set(CACHE_KEY, JSON.stringify(paidPlan)); + vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response("nope", { status: 503 })); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.source).toBe("cache"); + expect(result.success && result.data.entitlements).toEqual(paidPlan); + expect(JSON.parse(String(kv.store.get(CACHE_KEY)))).toEqual(paidPlan); + expect(kv.store.get(NEGATIVE_KEY)).toBeDefined(); + }); + + it("fails OPEN to unlimited when the request throws and nothing is cached", async () => { + const { env: e, kv } = env(); + vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("The operation timed out")); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + expect(result.success && result.data.entitlements).toEqual(UnlimitedEntitlements); + expect(kv.store.get(NEGATIVE_KEY)).toBeDefined(); + }); + + it("fails OPEN on an unparseable body", async () => { + const { env: e } = env(); + vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response("oops")); + const result = await new RemoteEntitlements(e, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + }); + + it("does nothing and reports 'default' when the hook is unconfigured", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const result = await new RemoteEntitlements({} as Env, noopLogger).refresh("usr_1", "user"); + expect(result.success && result.data.source).toBe("default"); + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); + +describe("warmEntitlements", () => { + function collector(): { waitUntil: (p: Promise) => void; settle: () => Promise } { + const pending: Promise[] = []; + return { + waitUntil: (p) => { + pending.push(p); + }, + settle: async () => { + await Promise.all(pending); + }, + }; + } + + it("warms the cache so a later forOwner reads it", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(jsonResponse(paidPlan)); + const { env: e } = env(); + const ctx = collector(); + warmEntitlements(e, ctx.waitUntil, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + await ctx.settle(); + + expect(fetchSpy).toHaveBeenCalledTimes(1); + const read = await new RemoteEntitlements(e, noopLogger).forOwner("usr_1", "user"); + expect(read.success && read.data.source).toBe("cache"); + expect(read.success && read.data.entitlements).toEqual(paidPlan); + }); + + it("single-flights concurrent cold warms into one request", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue(jsonResponse(paidPlan)); + const { env: e } = env(); + const ctx = collector(); + const subject = { ownerId: "usr_1", ownerType: "user" as const }; + warmEntitlements(e, ctx.waitUntil, subject, noopLogger); + warmEntitlements(e, ctx.waitUntil, subject, noopLogger); + warmEntitlements(e, ctx.waitUntil, subject, noopLogger); + await ctx.settle(); + expect(fetchSpy).toHaveBeenCalledTimes(1); + }); + + it("respects a lock another isolate holds", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { env: e, kv } = env(); + kv.store.set(LOCK_KEY, "1"); + const ctx = collector(); + warmEntitlements(e, ctx.waitUntil, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + await ctx.settle(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("releases its own lock once the answer is cached", async () => { + vi.spyOn(globalThis, "fetch").mockResolvedValue(jsonResponse(paidPlan)); + const { env: e, kv } = env(); + const ctx = collector(); + warmEntitlements(e, ctx.waitUntil, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + await ctx.settle(); + expect(kv.store.has(LOCK_KEY)).toBe(false); + }); + + it("does not re-fetch while a negative-cache entry stands", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { env: e, kv } = env(); + kv.store.set(NEGATIVE_KEY, String(Date.now())); + const ctx = collector(); + warmEntitlements(e, ctx.waitUntil, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + await ctx.settle(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("skips a warm cache", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { env: e, kv } = env(); + kv.store.set(CACHE_KEY, JSON.stringify(paidPlan)); + const ctx = collector(); + warmEntitlements(e, ctx.waitUntil, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + await ctx.settle(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("is inert with no execution context, unconfigured env, or no owner", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { env: e } = env(); + const ctx = collector(); + warmEntitlements(e, undefined, { ownerId: "usr_1", ownerType: "user" }, noopLogger); + warmEntitlements( + unconfiguredEnv(), + ctx.waitUntil, + { ownerId: "usr_1", ownerType: "user" }, + noopLogger, + ); + warmEntitlements(e, ctx.waitUntil, { ownerId: "", ownerType: "user" }, noopLogger); + await ctx.settle(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/usage-periods.test.ts b/tests/usage-periods.test.ts index 41d2db60..5825ac92 100644 --- a/tests/usage-periods.test.ts +++ b/tests/usage-periods.test.ts @@ -11,7 +11,7 @@ * the owner_type CHECK and the UPSERT's conflict target are the shipped ones * rather than a stub's idea of them. */ -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { recordCosts } from "../src/storage/costs"; import { captureDeletionTarget, @@ -39,6 +39,10 @@ const logger: Logger = { child: vi.fn(() => logger), }; +// The mock is module-level, so without this a `toHaveBeenCalledWith` assertion +// can be satisfied by an earlier test's call and the suite passes on run order. +beforeEach(() => vi.clearAllMocks()); + const ALICE = { ownerId: "user_alice", ownerType: "user" } as const; const PERIOD = "2026-09"; @@ -47,6 +51,7 @@ interface UsageRow { owner_type: string; period: string; meter: string; + source: string; quantity: number; updated_at: string; } @@ -55,7 +60,7 @@ type Raw = ReturnType["raw"]; function usageRows(raw: Raw): UsageRow[] { return raw - .prepare("SELECT * FROM usage_periods ORDER BY owner_id, period, meter") + .prepare("SELECT * FROM usage_periods ORDER BY owner_id, period, meter, source") .all() as unknown as UsageRow[]; } @@ -118,8 +123,18 @@ describe("deleting a project does not refund the allowance", () => { const period = usagePeriod(); const before = await getOwnerUsageSummary(db, logger, "user_alice", period); expect(before.success && before.data).toEqual([ - { meter: "llm_tokens_month", quantity: 12_000, updatedAt: expect.any(String) }, - { meter: "sandbox_ms_month", quantity: 4_500, updatedAt: expect.any(String) }, + { + meter: "llm_tokens_month", + source: "platform", + quantity: 12_000, + updatedAt: expect.any(String), + }, + { + meter: "sandbox_ms_month", + source: "platform", + quantity: 4_500, + updatedAt: expect.any(String), + }, ]); expect(countCostRows(raw)).toBe(3); @@ -240,7 +255,7 @@ describe("upsertUsage", () => { expect(result.success).toBe(true); expect(usageRows(raw).map((r) => [r.meter, r.quantity])).toEqual([["llm_tokens_month", 100]]); expect(logger.warn).toHaveBeenCalledWith( - expect.stringContaining("not finite"), + expect.stringContaining("not a usable increment"), expect.objectContaining({ meter: "llm_tokens_month" }), ); }); @@ -321,12 +336,22 @@ describe("getOwnerUsageSummary", () => { }); describe("usagePeriod", () => { - it("classifies a moment by its UTC month, not the local one", () => { - // A local-time month would put a spend just after midnight UTC into the - // previous period for half the world, and two counters keyed on different - // months are two allowances. + it("puts either side of a UTC month boundary in a different period", () => { + // Honest about its own limits: on a UTC runner — which is what CI is — + // these assertions do NOT discriminate against a local-time implementation, + // because local and UTC agree. Forcing `process.env.TZ` here would make them + // discriminate, and was tried; it is process-global while vitest runs + // several files per worker, so it perturbs every other suite sharing the + // process. A flaky suite costs more than this test can buy. + // + // The guarantee therefore rests on `usagePeriod` deriving from + // `toISOString()`, which is UTC by construction. What these pin is the + // property that matters either way: one millisecond either side of a + // boundary must not land in the same period, or two counters would share an + // allowance across the rollover. expect(usagePeriod(new Date("2026-09-30T23:59:59.999Z"))).toBe("2026-09"); expect(usagePeriod(new Date("2026-10-01T00:00:00.000Z"))).toBe("2026-10"); + expect(usagePeriod(new Date("2025-12-31T23:59:59.999Z"))).toBe("2025-12"); expect(usagePeriod(new Date("2026-01-01T00:00:00.000Z"))).toBe("2026-01"); }); }); @@ -409,6 +434,52 @@ describe("recordCosts accumulates usage alongside the ledger", () => { await recordCosts(db, logger, { project: "@acme/api", ownerId: "org_acme", ownerType: "org" }, [ { kind: "llm_tokens", quantity: 800, source: "byok" }, ]); - expect(usageRows(raw)[0]).toMatchObject({ owner_id: "org_acme", quantity: 800 }); + expect(usageRows(raw)[0]).toMatchObject({ + owner_id: "org_acme", + quantity: 800, + source: "byok", + }); + }); + + it("refuses a negative delta rather than refunding the month", async () => { + // The refund vector the CHECK and the guard both exist for. sandbox_ms is + // wall-clock arithmetic, so a backwards clock step produces one without + // anybody trying — and it would subtract through the single statement this + // table exists to make non-refundable. + const { db, raw } = makeSqliteD1(); + await upsertUsage(db, logger, ALICE, "2026-09", [ + { meter: "sandbox_ms_month", quantity: 5_000 }, + ]); + const result = await upsertUsage(db, logger, ALICE, "2026-09", [ + { meter: "sandbox_ms_month", quantity: -4_000 }, + ]); + + expect(result.success).toBe(true); + expect(usageRows(raw).map((r) => r.quantity)).toEqual([5_000]); + expect(logger.warn).toHaveBeenCalledWith( + expect.stringContaining("not a usable increment"), + expect.objectContaining({ quantity: -4_000 }), + ); + }); + + it("keeps BYOK and platform spend as separate totals for the same meter", async () => { + // The reason `source` is in the PRIMARY KEY. An allowance limits what the + // OPERATOR spends, so enforcement reads the platform row alone — a project + // paying its own provider costs the operator nothing. Summed together these + // could never be separated again: this table is not rebuildable from + // cost_records, so the dimension has to be captured as it is written. + const { db, raw } = makeSqliteD1(); + await recordCosts(db, logger, { project: "@alice/my-repo", ...ALICE }, [ + { kind: "llm_tokens", quantity: 1000 }, + { kind: "llm_tokens", quantity: 250, source: "byok" }, + ]); + await recordCosts(db, logger, { project: "@alice/my-repo", ...ALICE }, [ + { kind: "llm_tokens", quantity: 500, source: "byok" }, + ]); + + const rows = usageRows(raw).filter((r) => r.meter === "llm_tokens_month"); + expect(rows).toHaveLength(2); + expect(rows.find((r) => r.source === "platform")?.quantity).toBe(1000); + expect(rows.find((r) => r.source === "byok")?.quantity).toBe(750); }); }); From 7f527a92f37f935e3b3ab868dc216e0f9054d622 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 10:06:28 +0000 Subject: [PATCH 10/26] Add the usage meter, keyed on who the limit is actually for MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 5. The counter that admits or refuses a spend, plus its bindings. Nothing calls it yet — enforcement is Task 8. Keyed on the enforcement subject rather than the billing subject, which is the correction revision 3 made: keying it on the recorded subject would have reopened the reset hole at this layer regardless of how correct the D1 sum above it is, because N free orgs would be N independent counters each holding a full limit. The DO never resolves its own subject — the caller passes one already resolved — since a DO that reached for the entitlements cache would make its serialization guarantee depend on a lookup that can be stale. Reserve-then-settle rather than a plain increment, because an LLM call's true cost is only known from its response: reserve an upper bound, settle the difference after. A negative delta is legitimate here and forbidden one layer down, where migration 049 CHECKs it — settle corrects a reservation, a negative aggregate delta refunds the month. Same sign, opposite meaning, and both are commented so the pair does not read as a contradiction. evaluations_per_hour lives here rather than in KV. KV's read staleness is about the length of the existing minute bucket, which is exactly why the current request limiter is sound; stretch the window to an hour and the counter degrades to per-colo-per-minute, admitting roughly sixty times the limit to a distributed caller. It is a sliding window of twelve five-minute buckets, so there is no fresh hour to burst into — a test spends the allowance at 11:59 and asserts refusal at 12:00, which a single hourly bucket would have admitted. Two storage keys, not the precedent's one: a month rolling must not reset an hourly window straddling the 1st, so alarm() sweeps what has expired and re-arms for what has not rather than wiping. Malformed input admits, opposite the magic-link limiter it is modelled on. Refusing is right for a send cap and wrong for a merge gate: a NaN limit from a bad billing payload must not block every merge. Every such admission is logged, because a limit that stopped limiting has to be visible rather than inferred from a bill. Deletion purges by a type-prefixed name, so "deleting one org must not purge a person's counter" is a property of the naming rather than of two id generators never colliding. The project cascade purges none, with a test asserting it. Typecheck clean; lint clean; 3877 tests pass, up 59. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- src/index.ts | 1 + src/queue/usage-meter.ts | 621 ++++++++++++++++++++++ src/storage/deletion.ts | 43 +- src/types.ts | 14 + tests/deletion-account.test.ts | 92 +++- tests/deletion.test.ts | 19 + tests/helpers/usage-meter.ts | 31 ++ tests/usage-meter.test.ts | 592 +++++++++++++++++++++ tests/wrangler-usage-meter-config.test.ts | 100 ++++ wrangler.toml | 76 +++ 10 files changed, 1584 insertions(+), 5 deletions(-) create mode 100644 src/queue/usage-meter.ts create mode 100644 tests/helpers/usage-meter.ts create mode 100644 tests/usage-meter.test.ts create mode 100644 tests/wrangler-usage-meter-config.test.ts diff --git a/src/index.ts b/src/index.ts index a660b1b9..ab06fc4a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -58,6 +58,7 @@ import { createLogger } from "./utils/logger"; export { MergeQueue } from "./queue/merge-queue"; export { MagicLinkRateLimiter } from "./queue/magic-link-limiter"; export { RepoDO } from "./queue/repo-do"; +export { UsageMeter } from "./queue/usage-meter"; const app = new Hono<{ Bindings: Env }>(); diff --git a/src/queue/usage-meter.ts b/src/queue/usage-meter.ts new file mode 100644 index 00000000..54a406e3 --- /dev/null +++ b/src/queue/usage-meter.ts @@ -0,0 +1,621 @@ +import { DurableObject } from "cloudflare:workers"; +import { UNLIMITED } from "../billing/entitlements"; +import type { MeterKey } from "../storage/usage"; +import type { Env } from "../types"; +import { type Logger, createLogger } from "../utils/logger"; + +/** + * The monthly counters, all of them, under one key: `{ period, counts }`. + * + * One key per meter would make a rollover something to sweep; one key holding + * the period makes it free, exactly as `MagicLinkRateLimiter`'s bucket does — + * the first write of a new month overwrites the record and the previous month's + * counts are gone without a delete. + */ +const PERIOD_KEY = "period"; + +/** + * The windowed rate counters, kept apart from {@link PERIOD_KEY} because their + * lifetimes are unrelated: an hourly window that straddles midnight on the 1st + * must not be reset by the month rolling, and a month's counter must not be + * dropped when the hour it happened to be written in expires. + */ +const RATE_KEY = "rate"; + +/** + * How long after a window closes the object waits before erasing what belongs + * to it. Longer than the precedent's minute because the window here is a + * calendar month: a late `settle` for a spend that started before midnight on + * the last day is ordinary, and erasing the record out from under it would + * discard a correction rather than a stale count. + */ +const CLEANUP_GRACE_MS = 60 * 60 * 1000; + +/** + * A rate meter: a bound on burst over a rolling window, not a bound on the + * month's spend. `reserve` with no `settle` is exactly that check. + * + * This lives here rather than in KV (PRD §4b) for one reason. KV reads come + * from a per-colo edge cache whose staleness is about a minute, which is + * tolerable for `rateLimitMiddleware` only because its bucket also rolls every + * minute — staleness and window are the same size. Stretch the window to an + * hour and staleness is 1/60th of it, so an N/hour limit admits roughly + * N x colos x 60 to a distributed caller: not a loose limit, no limit at all, + * against precisely the distributed burst this control exists to stop. + */ +export type RateMeterKey = "evaluations_per_hour"; + +/** + * Everything this object counts: the monthly flows of `MeterKey`, plus the + * rate meters above. + * + * The two are counted differently and the difference is deliberate. A monthly + * meter is keyed on the caller's `period` string and reconciled against + * `usage_periods`; a rate meter is keyed on wall-clock buckets and reconciled + * against nothing, because no D1 aggregate of it exists. + */ +export type UsageMeterKey = MeterKey | RateMeterKey; + +interface RateWindow { + /** How far back the limit looks. */ + windowSeconds: number; + /** + * The granularity the window is summed at. SHORT on purpose: one bucket per + * window is a fixed window, and a fixed window admits the whole allowance in + * the first second of every new one — twice the limit across the boundary, + * which is the burst the meter exists to bound. Twelve buckets make the + * window slide, so an hour's allowance is an hour's allowance wherever the + * caller starts. + */ + bucketSeconds: number; +} + +const RATE_WINDOWS: Record = { + evaluations_per_hour: { windowSeconds: 60 * 60, bucketSeconds: 5 * 60 }, +}; + +/** Bucket id (`floor(epochSeconds / bucketSeconds)`) -> quantity in that bucket. */ +type RateBuckets = Record; + +interface PeriodRecord { + /** 'YYYY-MM' UTC, from `usagePeriod`. The counts below belong to this month and no other. */ + period: string; + counts: Partial>; +} + +interface RateRecord { + meters: Partial>; +} + +export interface MeterReserveOutcome { + admitted: boolean; + /** The counter after this call. Reported for logging and telemetry; callers gate on `admitted`. */ + count: number; +} + +export interface MeterSnapshot { + /** Echoed back so a caller cannot mistake a rolled-over read for a live one. */ + period: string; + /** Empty when this object holds nothing for `period` — "never used" and "used nothing" are one fact. */ + counts: Partial>; +} + +/** + * Which instance an enforcement subject maps to. + * + * The type prefix is not decoration. This object is keyed on the ENFORCEMENT + * subject (see the class doc comment), so a user and an org can both name one, + * and the deletion cascade must be able to erase one without touching the + * other: deleting a single org may not purge a counter that belongs to a + * person. Distinct prefixes make that a property of the naming rather than of + * whichever id generator happens not to collide today. + * + * @param subjectType - Whether the enforcement subject is a user or an org + * @param subjectId - The subject's id + * @returns The Durable Object name for that subject's counters + */ +export function usageMeterName(subjectType: "user" | "org", subjectId: string): string { + return `${subjectType}:${subjectId}`; +} + +/** + * Serialized usage counters, one instance per ENFORCEMENT subject. + * + * ## What this is for + * + * A Workers KV read-modify-write lets N concurrent requests read the same total + * and each write back "+1", so the counter advances by one and the cap bounds + * sequential traffic only — the race `MagicLinkRateLimiter`'s doc comment + * describes, and the one `rateLimitMiddleware` still has. A Durable Object + * closes it: the runtime's input gate holds incoming events for an instance + * while one of its storage operations is in flight, so the read and the write + * below cannot interleave with another request against the same subject. + * Serialization is per instance, so different subjects still run in parallel. + * + * ## Reserve, then settle + * + * An LLM call's true cost is known only from the response, so a plain increment + * would leave enforcement doing check-then-act — the exact property the object + * is bought for. `reserve` takes an UPPER BOUND before the spend and admits or + * refuses it atomically; `settle` applies the signed difference once the real + * cost is known. `MagicLinkRateLimiter.refund` is the model for the second half. + * + * ## The subject is the ENFORCEMENT subject, and this object does not choose it + * + * PRD §4a: the subject a limit is CHECKED against is not the subject a spend is + * RECORDED against. Recording keeps naming the project's owner, because a + * ledger must say what happened; a limit is checked against the acting user (or + * an org positively known to hold a paid plan), because an allowance keyed on + * the recorded owner is resettable — create an org, create a project in it, get + * a fresh `usage_periods` subject and a fresh allowance. + * + * That resolution is a decision about plans and actors, and it is NOT made + * here. The caller passes an already-resolved subject id and this object counts + * against it. Keep it that way: a DO that resolved its own subject would need + * the entitlements cache and the request's actor, and the one property it + * exists to provide — that everything for one subject is serialized — would + * then depend on a lookup that can be stale. + * + * ## Limits are caller-supplied + * + * As in the precedent, so the policy numbers stay next to the enforcement point + * rather than being baked into storage. That is safe only because DO RPC is + * reachable from Worker code alone: no request can name its own limit. **Do not + * expose this class over `fetch`** — a `fetch` handler would make `limit` an + * attacker-controlled input. + * + * ## Malformed input ADMITS — a deliberate deviation from the precedent + * + * `MagicLinkRateLimiter.reserve` returns `{admitted:false}` on a malformed + * limit: fail closed, correct for a magic-link cap, where refusing a send is a + * minor inconvenience and an unlimited endpoint is an abuse vector. It is wrong + * here. The limits reaching this object come from a billing payload over the + * network, and a `NaN` limit that refuses would block EVERY merge on the + * instance for as long as the billing service kept sending it — an outage + * caused by the metering, in a system whose entitlements layer fails open on + * purpose for exactly this reason (`src/billing/entitlements.ts`). So a + * malformed limit, estimate, period or clock admits, and says so in the log. + * + * Well-formed limits are NOT this branch: `-1` is unlimited (admit and count), + * and `0` is a hard block (refuse). Only a value that is not a limit at all — + * `NaN`, `1.5`, `-7`, an infinity — takes the admitting path. + */ +export class UsageMeter extends DurableObject { + /** + * Atomically admits or refuses an upper bound on a spend, and counts it. + * + * The whole bound has to fit: a reservation is admitted only when + * `count + estimate` is within the limit, because the caller cannot know yet + * that the spend will come in cheaper. `settle` gives back the difference. + * + * A rate meter (see {@link RateMeterKey}) takes the windowed path instead, + * where `period` is not consulted at all — its window is wall-clock and does + * not reset when the month does. + * + * @param meter - Which counter to charge + * @param estimate - Upper bound on this spend; a non-usable value counts nothing and still admits + * @param limit - The subject's allowance: `-1` unlimited, `0` a hard block, otherwise a ceiling + * @param period - 'YYYY-MM' UTC, from `usagePeriod`; ignored by rate meters + * @param nowMs - Caller's clock, so one request's counters cannot straddle a window boundary + * @returns Whether the reservation was admitted, and the counter after this call + */ + async reserve( + meter: UsageMeterKey, + estimate: number, + limit: number, + period: string, + nowMs: number, + ): Promise { + const usable = usableQuantity(estimate); + if (usable === null) { + // Counting nothing is the safe half of the deviation: a poisoned counter + // (NaN never compares below a limit again) would block the subject for + // the rest of the month, which is the outcome the deviation exists to + // avoid. The limit still applies to what is already counted. + warn("Usage meter estimate is not a usable quantity; reserving zero", { meter, estimate }); + } + const quantity = usable ?? 0; + if (isRateMeter(meter)) return this.reserveRate(meter, quantity, limit, nowMs); + return this.reservePeriod(meter, quantity, limit, period); + } + + /** + * Applies the real cost of a spend that was already reserved. + * + * **A negative delta is legitimate here**, and that reads as a contradiction + * next to `upsertUsage`, which drops one, and migration 049, whose CHECK + * forbids one. Both are right, because they hold different things. A + * `usage_periods` row is the month's accumulated total, and a negative there + * REFUNDS THE MONTH through the one statement that exists to make the month + * non-refundable. This counter is the in-flight reservation ledger for a + * single period: the negative it takes is the unused remainder of a bound + * this same request reserved moments ago — 4,000 tokens reserved, 900 spent, + * 3,100 handed back — and refusing it would leave every over-estimate + * charged in full, which walls a subject off at a fraction of its allowance. + * The floor at zero is what keeps the difference safe: a settle can never + * hand back more than the counter holds. + * + * A settle for a period this object is no longer holding is dropped, as the + * precedent's `refund` drops one whose window has rolled. The reservation it + * corrects went with the rollover, so there is nothing to correct. + * + * @param meter - The counter the matching `reserve` charged + * @param delta - Signed difference between the true cost and the reserved bound + * @param period - The period the matching `reserve` used + * @param nowMs - Caller's clock; used by rate meters to find the live bucket + */ + async settle(meter: UsageMeterKey, delta: number, period: string, nowMs: number): Promise { + if (!Number.isFinite(delta)) { + warn("Usage meter settle ignored: delta is not a finite number", { meter, delta }); + return; + } + if (isRateMeter(meter)) { + await this.settleRate(meter, delta, nowMs); + return; + } + const record = await this.ctx.storage.get(PERIOD_KEY); + if (record?.period !== period) return; + const current = record.counts[meter] ?? 0; + await this.ctx.storage.put(PERIOD_KEY, { + period, + counts: { ...record.counts, [meter]: Math.max(0, current + delta) }, + }); + } + + /** + * Raises a counter to match the D1 aggregate. The reconcile write path. + * + * **Monotonic within a live period: it may raise, never lower** (PRD §3). The + * counter is ahead of `usage_periods` by design — a reservation is taken + * before the spend and the aggregate is written after it — so a reconcile + * that treated D1 as authoritative would hand back every in-flight + * reservation on the instance, once per reconcile. Lowering is a correction + * to a CLOSED period, and it belongs in D1, not in a live counter. + * + * A floor for a period older than the one held is dropped for the same + * reason: 'YYYY-MM' sorts chronologically, so a stale reconcile cannot + * clobber the live month with last month's total. + * + * Rate meters are not reconcilable and are ignored here: there is no D1 + * aggregate of `evaluations_per_hour` to reconcile against, by design — an + * hourly window is not something a monthly table can carry. + * + * WHICH rows the caller summed is a decision this object cannot make or + * check, and getting it wrong is silent. It must be the ENFORCEMENT + * subject's rows — the actor's, not the project owner's, since that is what + * this instance counts — and `source = 'platform'` alone: migration 049 put + * `source` in the primary key so that a project paying its own provider bill + * is not charged against the hosted allowance, and a sum that omits the + * filter hands that distinction straight back. + * + * @param meter - The monthly counter to raise + * @param quantity - The aggregate's total for this subject, meter and period + * @param period - 'YYYY-MM' UTC the quantity was summed over + */ + async setFloor(meter: UsageMeterKey, quantity: number, period: string): Promise { + if (isRateMeter(meter)) { + warn("Usage meter floor ignored: rate meters have no aggregate to reconcile", { meter }); + return; + } + const floor = usableQuantity(quantity); + if (floor === null) { + warn("Usage meter floor ignored: quantity is not a usable total", { meter, quantity }); + return; + } + const endsAt = periodEndMs(period); + if (endsAt === null) { + warn("Usage meter floor ignored: unparseable period", { meter, period }); + return; + } + const record = await this.ctx.storage.get(PERIOD_KEY); + if (record && record.period > period) return; + const counts = record?.period === period ? record.counts : {}; + const current = counts[meter] ?? 0; + if (record?.period === period && current >= floor) return; + await this.ctx.storage.put(PERIOD_KEY, { + period, + counts: { ...counts, [meter]: Math.max(current, floor) }, + }); + await this.armCleanup(endsAt + CLEANUP_GRACE_MS); + } + + /** + * This subject's monthly counters for one period. + * + * The period is an argument rather than "now" for the reason `nowMs` is one + * on `reserve`: the caller decides which window it is asking about, so a read + * cannot straddle a boundary and answer for half of each. + * + * Rate meters are deliberately absent. Which of their buckets are live + * depends on a clock this method is not given, and a rate total answered + * against the wrong instant is worse than no answer; `reserve` returns the + * live count to the one caller that needs it. + * + * @param period - 'YYYY-MM' UTC, from `usagePeriod` + * @returns The counters held for that period, empty when none are + */ + async read(period: string): Promise { + const record = await this.ctx.storage.get(PERIOD_KEY); + return { period, counts: record?.period === period ? record.counts : {} }; + } + + /** + * Deletion-cascade RPC: DO storage is only reachable from inside the class, + * so the ACCOUNT cascade calls this to drop the counters of a subject that no + * longer exists. Deliberately not called by the project cascade — an + * allowance a user can reset by deleting a project is not an allowance + * (`src/storage/deletion.ts`, and migration 049's header). + */ + async purge(): Promise { + await this.ctx.storage.deleteAll(); + } + + /** + * Erases what has expired, so a subject seen once does not leave durable + * state behind forever. Without it a per-owner monthly meter accumulates a + * record per subject for the life of the deployment; this is what + * `expirationTtl` did for the KV counters the precedent replaced. + * + * Selective rather than the precedent's flat `deleteAll` because two + * independent lifetimes share this object: a flat wipe on the month's alarm + * would erase a live hourly window, handing a caller a fresh burst allowance + * every time a month ended. Uses `Date.now()` rather than a caller's clock — + * nothing is being admitted or refused here, so there is no decision a + * straddled boundary could corrupt. + */ + override async alarm(): Promise { + const nowMs = Date.now(); + let nextAlarm: number | null = null; + + const record = await this.ctx.storage.get(PERIOD_KEY); + if (record) { + const endsAt = periodEndMs(record.period); + const expiresAt = endsAt === null ? null : endsAt + CLEANUP_GRACE_MS; + // An unparseable period cannot expire on its own, so it is swept now + // rather than pinning storage forever. + if (expiresAt === null || expiresAt <= nowMs) await this.ctx.storage.delete(PERIOD_KEY); + else nextAlarm = expiresAt; + } + + const rates = await this.ctx.storage.get(RATE_KEY); + if (rates) { + const pruned = pruneRates(rates, nowMs); + if (pruned === null) { + await this.ctx.storage.delete(RATE_KEY); + } else { + await this.ctx.storage.put(RATE_KEY, pruned.record); + nextAlarm = nextAlarm === null ? pruned.expiresAt : Math.max(nextAlarm, pruned.expiresAt); + } + } + + // Nothing left to expire: erase whatever remains (a key written by an + // older shape of this class, say) and leave no alarm armed, so the object + // costs nothing until it is next used. + if (nextAlarm === null) await this.ctx.storage.deleteAll(); + else await this.ctx.storage.setAlarm(nextAlarm); + } + + private async reservePeriod( + meter: UsageMeterKey, + quantity: number, + limit: number, + period: string, + ): Promise { + const endsAt = periodEndMs(period); + if (endsAt === null) { + warn("Usage meter reserve admitted: unparseable period", { meter, period }); + return { admitted: true, count: 0 }; + } + const record = await this.ctx.storage.get(PERIOD_KEY); + // A record from an earlier month reads as zero rather than being deleted: + // the `put` below overwrites it, so the stale record costs nothing. + const counts = record?.period === period ? record.counts : {}; + const current = counts[meter] ?? 0; + const decision = limitAdmits(current, quantity, limit); + if (decision.malformed) { + warn("Usage meter reserve admitted: limit is not a usable limit", { meter, limit }); + } + if (!decision.admitted) return { admitted: false, count: current }; + const next = current + quantity; + await this.ctx.storage.put(PERIOD_KEY, { + period, + counts: { ...counts, [meter]: next }, + }); + // Re-armed on every admission so the erase always trails the live period. + await this.armCleanup(endsAt + CLEANUP_GRACE_MS); + return { admitted: true, count: next }; + } + + private async reserveRate( + meter: RateMeterKey, + quantity: number, + limit: number, + nowMs: number, + ): Promise { + if (!Number.isFinite(nowMs)) { + warn("Usage meter reserve admitted: clock is not a usable timestamp", { meter, nowMs }); + return { admitted: true, count: 0 }; + } + const window = RATE_WINDOWS[meter]; + const bucket = bucketId(nowMs, window.bucketSeconds); + const record = await this.ctx.storage.get(RATE_KEY); + const live = liveBuckets(record?.meters[meter], bucket, window); + const total = sum(live); + const decision = limitAdmits(total, quantity, limit); + if (decision.malformed) { + warn("Usage meter reserve admitted: limit is not a usable limit", { meter, limit }); + } + if (!decision.admitted) return { admitted: false, count: total }; + const key = String(bucket); + live[key] = (live[key] ?? 0) + quantity; + await this.ctx.storage.put(RATE_KEY, { + meters: { ...record?.meters, [meter]: live }, + }); + await this.armCleanup(rateExpiresAt(bucket, window)); + return { admitted: true, count: total + quantity }; + } + + /** + * Rate meters are reserve-with-no-settle by convention, so this exists only + * for the caller that reserved an evaluation which then never ran. The delta + * lands on the live bucket rather than on whichever bucket the reservation + * was taken in: the limit reads the window's SUM, so the sum is corrected + * either way, and per-bucket floors of zero keep a refund from inventing + * headroom that was never reserved. + */ + private async settleRate(meter: RateMeterKey, delta: number, nowMs: number): Promise { + if (!Number.isFinite(nowMs)) { + warn("Usage meter settle ignored: clock is not a usable timestamp", { meter, nowMs }); + return; + } + const window = RATE_WINDOWS[meter]; + const bucket = bucketId(nowMs, window.bucketSeconds); + const record = await this.ctx.storage.get(RATE_KEY); + if (!record) return; + const live = liveBuckets(record.meters[meter], bucket, window); + const key = String(bucket); + live[key] = Math.max(0, (live[key] ?? 0) + delta); + await this.ctx.storage.put(RATE_KEY, { + meters: { ...record.meters, [meter]: live }, + }); + } + + /** + * Arms the cleanup alarm for `atMs` unless a later one is already set. A DO + * has one alarm and two lifetimes to sweep, so the later time wins and + * `alarm()` re-arms for whatever it did not erase. + */ + private async armCleanup(atMs: number): Promise { + const existing = await this.ctx.storage.getAlarm(); + if (existing === null || existing < atMs) await this.ctx.storage.setAlarm(atMs); + } +} + +function isRateMeter(meter: UsageMeterKey): meter is RateMeterKey { + return meter in RATE_WINDOWS; +} + +/** + * A quantity that can safely be added to a counter, or null. + * + * `NaN` or an infinity would poison the counter permanently — no later + * comparison against a limit could be true again — and a negative would refund + * the period, which is what `reserve` is here to stop. `settle` is the only + * signed path, and it clamps at zero. + */ +function usableQuantity(value: number): number | null { + return Number.isFinite(value) && value >= 0 ? value : null; +} + +/** + * Whether a reservation fits, and whether the limit was a limit at all. + * + * `malformed` is reported rather than folded into `admitted` so the caller's + * log can tell "admitted under the allowance" from "admitted because the + * billing payload was garbage" — the distinction `EntitlementsSource` exists + * for one layer up, and a month of observation that cannot make it is a month + * of unreadable data. + */ +function limitAdmits( + current: number, + quantity: number, + limit: number, +): { admitted: boolean; malformed: boolean } { + if (!Number.isInteger(limit) || limit < UNLIMITED) return { admitted: true, malformed: true }; + if (limit === UNLIMITED) return { admitted: true, malformed: false }; + // `0` is a hard block, not "unset" (see `src/billing/entitlements.ts`): a + // subject limited to zero may not consume the resource at all, so a + // zero-quantity reservation must not sneak past on the arithmetic below. + if (limit === 0) return { admitted: false, malformed: false }; + return { admitted: current + quantity <= limit, malformed: false }; +} + +/** The instant a 'YYYY-MM' UTC period ends, or null when it is not one. */ +function periodEndMs(period: string): number | null { + const match = /^(\d{4})-(\d{2})$/.exec(period); + if (!match) return null; + const year = Number(match[1]); + const month = Number(match[2]); + if (month < 1 || month > 12) return null; + // Date.UTC's month is 0-based, so passing the 1-based month lands on the + // first instant of the NEXT month — the end of this one. + return Date.UTC(year, month, 1); +} + +function bucketId(nowMs: number, bucketSeconds: number): number { + return Math.floor(nowMs / 1000 / bucketSeconds); +} + +function bucketsPerWindow(window: RateWindow): number { + return Math.max(1, Math.ceil(window.windowSeconds / window.bucketSeconds)); +} + +/** The buckets still inside the window ending at `bucket`, as a fresh object. */ +function liveBuckets( + buckets: RateBuckets | undefined, + bucket: number, + window: RateWindow, +): RateBuckets { + const oldest = bucket - bucketsPerWindow(window) + 1; + const live: RateBuckets = {}; + for (const [id, quantity] of Object.entries(buckets ?? {})) { + const parsed = Number(id); + // A bucket in the FUTURE is kept: a caller's clock that ran ahead once must + // not have its reservation dropped by the next caller's slower clock. + if (Number.isFinite(parsed) && parsed >= oldest) live[id] = quantity; + } + return live; +} + +function sum(buckets: RateBuckets): number { + let total = 0; + for (const quantity of Object.values(buckets)) total += quantity; + return total; +} + +/** When a bucket stops counting toward any window, plus the cleanup grace. */ +function rateExpiresAt(bucket: number, window: RateWindow): number { + const bucketMs = window.bucketSeconds * 1000; + return (bucket + bucketsPerWindow(window)) * bucketMs + CLEANUP_GRACE_MS; +} + +/** + * Drop every bucket that has fallen out of its window. Returns null when + * nothing is left to keep, so the caller can erase the key rather than store an + * empty shell. + */ +function pruneRates( + record: RateRecord, + nowMs: number, +): { record: RateRecord; expiresAt: number } | null { + const meters: Partial> = {}; + let expiresAt = 0; + for (const [name, buckets] of Object.entries(record.meters ?? {})) { + if (!isRateMeter(name as UsageMeterKey)) continue; + const meter = name as RateMeterKey; + const window = RATE_WINDOWS[meter]; + const live = liveBuckets(buckets, bucketId(nowMs, window.bucketSeconds), window); + const ids = Object.keys(live); + if (ids.length === 0) continue; + meters[meter] = live; + const newest = Math.max(...ids.map(Number)); + expiresAt = Math.max(expiresAt, rateExpiresAt(newest, window)); + } + return expiresAt === 0 ? null : { record: { meters }, expiresAt }; +} + +// One logger per isolate rather than one per call: these paths are rare (a +// malformed input) but sit on the merge path, and `createLogger` builds a pino +// instance every time it is called. +let meterLogger: Logger | undefined; + +/** + * Every deviation from "count it" is logged, because each one silently ADMITS + * (see the class doc comment) and AGENTS.md forbids swallowing a failure. A + * limit that stopped limiting must be visible in Workers Logs, not inferred + * from a bill. + */ +function warn(message: string, meta: Record): void { + meterLogger ??= createLogger({ component: "usage-meter" }); + meterLogger.warn(message, meta); +} diff --git a/src/storage/deletion.ts b/src/storage/deletion.ts index f8d18c7c..cac2384a 100644 --- a/src/storage/deletion.ts +++ b/src/storage/deletion.ts @@ -1,3 +1,4 @@ +import { usageMeterName } from "../queue/usage-meter"; import type { Env, ProjectEntry, WorkspaceEntry } from "../types"; import { type AppError, toAppError } from "../utils/errors"; import type { Logger } from "../utils/logger"; @@ -382,9 +383,17 @@ async function deleteArtifactsRepo(env: Env, name: string, logger: Logger): Prom return false; } -/** Purge a Durable Object's storage via its `purge()` RPC; false on failure. */ -async function purgeDurableObject( - namespace: DurableObjectNamespace | undefined, +/** + * Purge a Durable Object's storage via its `purge()` RPC; false on failure. + * + * Generic over the namespace's class because the bindings differ: MERGE_QUEUE + * and REPO_DO are untyped namespaces while USAGE_METER is parameterised by its + * class (`src/types.ts`), and the two are not assignable to one another. The + * `purge()` call is asserted below either way, since a typed stub is not enough + * to prove the method exists on an untyped one. + */ +async function purgeDurableObject( + namespace: DurableObjectNamespace | undefined, key: string, label: string, logger: Logger, @@ -497,6 +506,13 @@ export async function deleteProjectCascade( residuals.push(`do:MergeQueue:${key}`); } } + // USAGE_METER is deliberately absent from this cascade, and its absence is + // the point of the aggregate it counts against: the meter is keyed on an + // ENFORCEMENT subject (a user, or a paid org) and never on a project, so a + // project deletion must not touch it. Purging it here would refund the + // month by deleting a project — exactly what keeping `usage_periods` out of + // PROJECT_SCOPED_TABLES prevents (migration 049). The ACCOUNT cascade + // erases it, and nothing else does. // 5) KV last, project entry very last: as long as the entry exists, the // project is still resolvable and a re-drive can recapture everything. @@ -797,6 +813,16 @@ async function resolveOrgOwnership( // a successor was promoted above, the org survives and so must its // consumption — a change of owner is not a fresh monthly allowance. await db.prepare("DELETE FROM usage_periods WHERE owner_id = ?").bind(org.id).run(); + // ...and so does the org's enforcement counter, on this branch alone. The DO + // is named `org:`, never a bare id, so deleting one org can only reach + // that org's own counter: a person who spent inside it was charged against + // `user:` (PRD §4a — the enforcement subject is the actor unless the org + // is positively known to be paid), which is a different object this cannot + // name. Purging a shared counter here would erase a stranger's allowance. + const meterName = usageMeterName("org", org.id); + if (!(await purgeDurableObject(env.USAGE_METER, meterName, "UsageMeter", logger))) { + residuals.push(`do:UsageMeter:${meterName}`); + } await db.prepare("DELETE FROM orgs WHERE id = ?").bind(org.id).run(); } @@ -877,6 +903,17 @@ export async function deleteAccountCascade( // the project that burned it is not an allowance. Here the subject itself is // going, so there is nothing left to bill and nothing left to enforce. await db.prepare("DELETE FROM usage_periods WHERE owner_id = ?").bind(userId).run(); + // The same subject's live counters. The DO is keyed on the ENFORCEMENT + // subject (PRD §4a), which for a person is the person: everything this user + // ran — including inside orgs they are merely a member of — was checked + // against `user:`, so erasing the account erases all of it while those + // orgs keep their own rows and their own counters. A failure becomes a + // residual, which retains the users row for a re-drive rather than + // stranding a counter under a deleted account. + const userMeterName = usageMeterName("user", userId); + if (!(await purgeDurableObject(env.USAGE_METER, userMeterName, "UsageMeter", logger))) { + residuals.push(`do:UsageMeter:${userMeterName}`); + } // Same reasoning, same trap: `oauth_tokens.user_id` and // `oauth_auth_codes.user_id` both REFERENCE users(id) (#349). An OAuth // grant is a live credential handed to an editor, so an erasure that left diff --git a/src/types.ts b/src/types.ts index 96a463b3..81b9ee9b 100644 --- a/src/types.ts +++ b/src/types.ts @@ -2,6 +2,7 @@ // runner's message union cannot create a runtime import cycle. import type { DeployQueueMessage } from "./deploy/runner"; import type { MagicLinkRateLimiter } from "./queue/magic-link-limiter"; +import type { UsageMeter } from "./queue/usage-meter"; import type { LoggerContext } from "./utils/logger"; export type { LoggerContext }; @@ -179,6 +180,19 @@ export interface Env { * is type-only, so the cycle with magic-link-limiter.ts (which imports `Env` * from here) is erased. */ MAGIC_LINK_LIMITER?: DurableObjectNamespace; + /** Serialized usage counters, one instance per ENFORCEMENT subject (PRD §4a: + * the acting user, or an org positively known to be paid — NOT the recorded + * billing subject, which is resettable by creating an org). Optional only so + * tests and minimal deployments can omit it; every deploy from this repo's + * wrangler.toml binds it, and nothing here enforces anything until the + * enforcement call sites exist. + * + * Parameterised by the class for the same reason as the limiter above: the + * `.get()` stub is typed, so a change to `reserve`/`settle`/`setFloor` is a + * compile error at the call site rather than something a cast absorbs. The + * import is type-only, so the cycle with usage-meter.ts (which imports `Env` + * from here) is erased. */ + USAGE_METER?: DurableObjectNamespace; /** Content-addressed git object plane (ADR 004 Phase 2). */ REPO_OBJECTS?: R2Bucket; /** Durable backup store (D1 dumps, KV identity, repo packs). Optional: when diff --git a/tests/deletion-account.test.ts b/tests/deletion-account.test.ts index a0c6ce23..06dabea3 100644 --- a/tests/deletion-account.test.ts +++ b/tests/deletion-account.test.ts @@ -6,7 +6,7 @@ import { } from "../src/storage/deletion"; import type { Env, ProjectEntry } from "../src/types"; import type { Logger } from "../src/utils/logger"; -import { type ExecutedStatement, makeKvStub } from "./helpers/deletion-stubs"; +import { type ExecutedStatement, makeDoNamespaceStub, makeKvStub } from "./helpers/deletion-stubs"; const mockLogger: Logger = { trace: vi.fn(), @@ -65,11 +65,12 @@ function sqlIndex(executed: ExecutedStatement[], fragment: string): number { return executed.findIndex((s) => s.sql.includes(fragment)); } -function makeEnv(db: D1Database, kv: KVNamespace): Env { +function makeEnv(db: D1Database, kv: KVNamespace, usageMeter?: DurableObjectNamespace): Env { return { DB: db, STATE: kv, ARTIFACTS: { delete: async () => true } as unknown as Env["ARTIFACTS"], + USAGE_METER: usageMeter, } as Env; } @@ -247,3 +248,90 @@ describe("deleteAccountCascade", () => { expect(executed.some((s) => s.sql.includes("DELETE FROM users WHERE id = ?"))).toBe(true); }); }); + +describe("deleteAccountCascade and the UsageMeter Durable Object", () => { + // The counter is keyed on the ENFORCEMENT subject (PRD §4a): a person's own + // instance, or an org's when the org is positively known to be paid. Erasure + // therefore has to reach exactly one of those and never the other. + + it("purges the counter named for the erased user", async () => { + const kv = makeKvStub(50); + const { db } = makeAccountD1(() => []); + const meter = makeDoNamespaceStub(); + + const result = await deleteAccountCascade(makeEnv(db, kv.kv, meter.ns), "usr_1", mockLogger); + + expect(result.success && result.data.residuals).toEqual([]); + expect(meter.purged).toEqual(["user:usr_1"]); + }); + + it("purges an empty org's counter, and nothing belonging to a person", async () => { + const kv = makeKvStub(50); + const { db } = makeAccountD1((sql) => { + if (sql.includes("SELECT id, owner_id FROM orgs WHERE owner_id")) { + return [{ id: "org_1", owner_id: "usr_1" }]; + } + // No admin, no member: the org is empty and goes with the account. + return []; + }); + const meter = makeDoNamespaceStub(); + + await deleteAccountCascade(makeEnv(db, kv.kv, meter.ns), "usr_1", mockLogger); + + // Both, and only these two: the org's own counter and the erased user's. + // A bare id would have made these one instance, so deleting an org would + // have erased the allowance of everyone who ever spent inside it — people + // whose usage is charged to `user:` precisely so an org cannot be used + // to reset it. + expect(meter.purged.sort()).toEqual(["org:org_1", "user:usr_1"]); + }); + + it("leaves a surviving org's counter alone when ownership is only transferred", async () => { + const kv = makeKvStub(50); + const { db } = makeAccountD1((sql) => { + if (sql.includes("SELECT id, owner_id FROM orgs WHERE owner_id")) { + return [{ id: "org_1", owner_id: "usr_1" }]; + } + if (sql.includes("role = 'admin' AND user_id != ?")) return [{ user_id: "usr_2" }]; + return []; + }); + const meter = makeDoNamespaceStub(); + + await deleteAccountCascade(makeEnv(db, kv.kv, meter.ns), "usr_1", mockLogger); + + // The org survives with a new owner, so its consumption must survive too — + // a change of owner is not a fresh monthly allowance. Only the departing + // user's own counter goes. + expect(meter.purged).toEqual(["user:usr_1"]); + }); + + it("reports a failed purge as a residual and retains the user row", async () => { + const kv = makeKvStub(50); + const { db, executed } = makeAccountD1(() => []); + const failing = { + idFromName: (name: string) => ({ name }), + get: () => ({ + purge: async () => { + throw new Error("durable object unavailable"); + }, + }), + } as unknown as DurableObjectNamespace; + + const result = await deleteAccountCascade(makeEnv(db, kv.kv, failing), "usr_1", mockLogger); + + expect(result.success).toBe(true); + expect(result.success && result.data.residuals).toContain("do:UsageMeter:user:usr_1"); + // A counter stranded under a deleted account is not erasure; the row stays + // so a re-drive can finish the job. + expect(executed.some((s) => s.sql.includes("DELETE FROM users WHERE id = ?"))).toBe(false); + }); + + it("completes without the binding, as a minimal deployment has none", async () => { + const kv = makeKvStub(50); + const { db } = makeAccountD1(() => []); + + const result = await deleteAccountCascade(makeEnv(db, kv.kv), "usr_1", mockLogger); + + expect(result.success && result.data.residuals).toEqual([]); + }); +}); diff --git a/tests/deletion.test.ts b/tests/deletion.test.ts index e99d9b5d..d36b3265 100644 --- a/tests/deletion.test.ts +++ b/tests/deletion.test.ts @@ -64,6 +64,7 @@ function makeEnv(parts: { artifacts?: Env["ARTIFACTS"]; repoDo?: DurableObjectNamespace; mergeQueue?: DurableObjectNamespace; + usageMeter?: DurableObjectNamespace; }): Env { return { DB: parts.db, @@ -71,6 +72,7 @@ function makeEnv(parts: { ARTIFACTS: parts.artifacts ?? makeArtifactsStub().artifacts, REPO_DO: parts.repoDo, MERGE_QUEUE: parts.mergeQueue, + USAGE_METER: parts.usageMeter, } as Env; } @@ -327,6 +329,23 @@ describe("deleteProjectCascade", () => { expect(mergeQueue.purged.sort()).toEqual(["api", "api-display", "proj_1"]); }); + it("never purges the UsageMeter, whose allowance must survive the project", async () => { + const { db } = makeRecordingD1(); + const usageMeter = makeDoNamespaceStub(); + const result = await deleteProjectCascade( + makeEnv({ db, kv: kvStub.kv, usageMeter: usageMeter.ns }), + makeTarget(), + mockLogger, + ); + + expect(result.success).toBe(true); + // The meter is keyed on an enforcement subject, never on a project. Purging + // it here would be the burn-then-delete refund that `usage_periods` was + // kept out of PROJECT_SCOPED_TABLES to prevent — with the counter reset + // even faster than the aggregate. + expect(usageMeter.purged).toEqual([]); + }); + it("under a collision, purges only the project_id-keyed MergeQueue DO", async () => { const { db } = makeRecordingD1(); const repoDo = makeDoNamespaceStub(); diff --git a/tests/helpers/usage-meter.ts b/tests/helpers/usage-meter.ts new file mode 100644 index 00000000..ce64aa6f --- /dev/null +++ b/tests/helpers/usage-meter.ts @@ -0,0 +1,31 @@ +import { UsageMeter } from "../../src/queue/usage-meter"; +import type { Env } from "../../src/types"; +import { type FakeDurableObjectStorage, makeFakeDurableObjects } from "./fake-durable-object"; + +export interface FakeUsageMeters { + /** Typed exactly as `Env` declares the binding, so no call site casts. */ + namespace: NonNullable; + instances: Map; + storages: Map; + calls: string[]; +} + +/** + * A fake namespace backed by real {@link UsageMeter} instances. + * + * The one cast lives here rather than at each call site, for the reason + * `makeMagicLinkLimiters` gives: `makeFakeDurableObjects` cannot know the RPC + * brand of whichever class it constructs, and asserting it once means a + * signature change on the Durable Object is a compile error in the tests + * instead of being absorbed by hand-written stub types. + * + * @param opts - `gated` (default true) models workerd's input gate + * @returns The namespace plus the harness's instances, storages, and RPC log + */ +export function makeUsageMeters(opts: { gated?: boolean } = {}): FakeUsageMeters { + const fake = makeFakeDurableObjects((ctx) => new UsageMeter(ctx, {} as Env), opts); + return { + ...fake, + namespace: fake.namespace as unknown as NonNullable, + }; +} diff --git a/tests/usage-meter.test.ts b/tests/usage-meter.test.ts new file mode 100644 index 00000000..6f13c624 --- /dev/null +++ b/tests/usage-meter.test.ts @@ -0,0 +1,592 @@ +/** + * The UsageMeter Durable Object. + * + * The concurrency suite at the bottom is the object's entire justification: a + * KV counter admits N concurrent reservations against a cap of one, and every + * other property here could be had from a D1 UPSERT. Its ungated twin is what + * keeps that claim honest. + */ +import { describe, expect, it } from "vitest"; +import type { UsageMeter } from "../src/queue/usage-meter"; +import { usageMeterName } from "../src/queue/usage-meter"; +import type { Env } from "../src/types"; +import { makeUsageMeters } from "./helpers/usage-meter"; + +const PERIOD = "2026-09"; +const NEXT_PERIOD = "2026-10"; +/** Mid-month, so nothing here depends on a boundary except the tests that mean to. */ +const NOW = Date.UTC(2026, 8, 15, 12, 0, 0); +const HOUR_MS = 60 * 60 * 1000; +const MINUTE_MS = 60 * 1000; + +/** + * The object's RPC surface, derived from the class rather than restated, so a + * signature change there fails this suite instead of being absorbed by a + * hand-written interface that has quietly gone out of date. + */ +type MeterStub = Pick; + +function stubFor(namespace: NonNullable, name: string): MeterStub { + return namespace.get(namespace.idFromName(name)); +} + +function meter(opts: { gated?: boolean } = {}) { + const meters = makeUsageMeters(opts); + return { ...meters, stub: stubFor(meters.namespace, "user:usr_1") }; +} + +describe("UsageMeter.reserve", () => { + it("admits reservations while the whole bound fits, then refuses", async () => { + const { stub } = meter(); + + const first = await stub.reserve("llm_tokens_month", 400, 1000, PERIOD, NOW); + const second = await stub.reserve("llm_tokens_month", 400, 1000, PERIOD, NOW); + // 800 + 400 would exceed 1000. The bound is what is reserved, so it has to + // fit whole: the caller cannot yet know the spend will come in cheaper. + const third = await stub.reserve("llm_tokens_month", 400, 1000, PERIOD, NOW); + + expect([first, second, third].map((o) => o.admitted)).toEqual([true, true, false]); + expect([first, second, third].map((o) => o.count)).toEqual([400, 800, 800]); + }); + + it("refuses a single reservation larger than the whole allowance", async () => { + const { stub } = meter(); + + const outcome = await stub.reserve("llm_tokens_month", 5000, 1000, PERIOD, NOW); + + expect(outcome).toEqual({ admitted: false, count: 0 }); + // A refusal writes nothing, so the counter is untouched for the next caller. + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBeUndefined(); + }); + + it("counts meters independently", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 900, 1000, PERIOD, NOW); + const deploys = await stub.reserve("deploys_month", 1, 5, PERIOD, NOW); + + expect(deploys).toEqual({ admitted: true, count: 1 }); + expect((await stub.read(PERIOD)).counts).toEqual({ llm_tokens_month: 900, deploys_month: 1 }); + }); + + it("treats -1 as unlimited: admits and still counts", async () => { + const { stub } = meter(); + + const outcome = await stub.reserve("llm_tokens_month", 10_000_000, -1, PERIOD, NOW); + + // Counting under an unlimited plan is the point of the observe-only period: + // an unmetered "unlimited" would make a month of measurement worthless. + expect(outcome).toEqual({ admitted: true, count: 10_000_000 }); + }); + + it("treats 0 as a hard block, including for a zero-cost reservation", async () => { + const { stub } = meter(); + + expect((await stub.reserve("llm_tokens_month", 100, 0, PERIOD, NOW)).admitted).toBe(false); + // `0` means "may not consume this resource at all" (src/billing/entitlements.ts), + // so it must not be defeated by an estimate that happens to be zero. + expect((await stub.reserve("llm_tokens_month", 0, 0, PERIOD, NOW)).admitted).toBe(false); + }); +}); + +describe("UsageMeter malformed input", () => { + // The deliberate deviation from MagicLinkRateLimiter, which REFUSES on a bad + // limit. Here the limits arrive from a billing service over the network, so a + // refusing meter would turn one bad payload into a total merge outage. + + it.each([ + ["NaN", Number.NaN], + ["Infinity", Number.POSITIVE_INFINITY], + ["a fraction", 1.5], + ["a negative that is not -1", -7], + ])("admits when the limit is %s", async (_label, limit) => { + const { stub } = meter(); + + const outcome = await stub.reserve("llm_tokens_month", 500, limit as number, PERIOD, NOW); + + expect(outcome).toEqual({ admitted: true, count: 500 }); + }); + + it("keeps admitting under a malformed limit rather than degrading to a block", async () => { + const { stub } = meter(); + + for (let i = 0; i < 20; i++) { + expect((await stub.reserve("llm_tokens_month", 1000, Number.NaN, PERIOD, NOW)).admitted).toBe( + true, + ); + } + }); + + it("admits a malformed estimate and counts nothing, leaving the counter usable", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW); + const bad = await stub.reserve("llm_tokens_month", Number.NaN, 1000, PERIOD, NOW); + + // A NaN added to the counter would never compare below a limit again — the + // subject would be blocked for the rest of the month by a bad estimate. + expect(bad).toEqual({ admitted: true, count: 100 }); + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(100); + expect((await stub.reserve("llm_tokens_month", 950, 1000, PERIOD, NOW)).admitted).toBe(false); + }); + + it("admits, and writes nothing, when the period is not a period", async () => { + const { stub, storages } = meter(); + + expect((await stub.reserve("llm_tokens_month", 100, 10, "not-a-month", NOW)).admitted).toBe( + true, + ); + expect((await stub.reserve("llm_tokens_month", 100, 10, "2026-13", NOW)).admitted).toBe(true); + + const storage = storages.get("user:usr_1"); + expect(await storage?.get("period")).toBeUndefined(); + }); + + it("ignores a settle whose delta is not a finite number", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 500, 1000, PERIOD, NOW); + await stub.settle("llm_tokens_month", Number.NaN, PERIOD, NOW); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(500); + }); +}); + +describe("UsageMeter.settle", () => { + it("hands back the unused remainder of a reservation", async () => { + const { stub } = meter(); + + // The shape Task 8 will use: reserve maxDiffChars/4 + max output, then + // settle the difference once the provider reports real token counts. + const reserved = await stub.reserve("llm_tokens_month", 4000, 10_000, PERIOD, NOW); + expect(reserved.admitted).toBe(true); + await stub.settle("llm_tokens_month", -3100, PERIOD, NOW); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(900); + }); + + it("returns the whole reservation when the provider call fails", async () => { + const { stub } = meter(); + + // 10 reservations of 1000 exhaust the allowance... + for (let i = 0; i < 10; i++) { + expect((await stub.reserve("llm_tokens_month", 1000, 10_000, PERIOD, NOW)).admitted).toBe( + true, + ); + } + expect((await stub.reserve("llm_tokens_month", 1000, 10_000, PERIOD, NOW)).admitted).toBe( + false, + ); + + // ...and a provider that returned nothing spent nothing, so the caller + // settles the full negative delta and the allowance comes back. + await stub.settle("llm_tokens_month", -1000, PERIOD, NOW); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(9000); + expect((await stub.reserve("llm_tokens_month", 1000, 10_000, PERIOD, NOW)).admitted).toBe(true); + }); + + it("charges a spend that came in over the reserved bound", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 1000, 10_000, PERIOD, NOW); + await stub.settle("llm_tokens_month", 250, PERIOD, NOW); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(1250); + }); + + it("never drives a counter below zero", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 10_000, PERIOD, NOW); + await stub.settle("llm_tokens_month", -500, PERIOD, NOW); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(0); + }); + + it("drops a settle for a period this object is no longer holding", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 500, 10_000, PERIOD, NOW); + // The month rolled between the reserve and the settle; crediting the new + // month for last month's over-estimate would carry an allowance forward. + await stub.reserve("llm_tokens_month", 100, 10_000, NEXT_PERIOD, NOW + 31 * 24 * HOUR_MS); + await stub.settle("llm_tokens_month", -400, PERIOD, NOW); + + expect((await stub.read(NEXT_PERIOD)).counts.llm_tokens_month).toBe(100); + }); +}); + +describe("UsageMeter period rollover", () => { + it("starts the next period from zero and forgets the last one", async () => { + const { stub } = meter(); + + for (let i = 0; i < 10; i++) { + await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW); + } + expect((await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW)).admitted).toBe(false); + + const rolled = await stub.reserve( + "llm_tokens_month", + 100, + 1000, + NEXT_PERIOD, + NOW + 31 * 24 * HOUR_MS, + ); + + expect(rolled).toEqual({ admitted: true, count: 100 }); + // One key holds the period, so the rollover overwrote the old counts + // instead of leaving a record per month for something to sweep. + expect((await stub.read(PERIOD)).counts).toEqual({}); + }); + + it("reads a period it holds nothing for as empty, not as the live one", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW); + + expect(await stub.read(NEXT_PERIOD)).toEqual({ period: NEXT_PERIOD, counts: {} }); + }); +}); + +describe("UsageMeter.setFloor", () => { + it("raises the counter to the aggregate", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 10_000, PERIOD, NOW); + await stub.setFloor("llm_tokens_month", 4200, PERIOD); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(4200); + }); + + it("never lowers a live period, however far behind the aggregate is", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 9000, 10_000, PERIOD, NOW); + // D1 lags: the aggregate is written after the spend, the reservation before + // it. A reconcile that treated D1 as authoritative would hand back every + // in-flight reservation on the instance. + await stub.setFloor("llm_tokens_month", 10, PERIOD); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(9000); + expect((await stub.reserve("llm_tokens_month", 2000, 10_000, PERIOD, NOW)).admitted).toBe( + false, + ); + }); + + it("establishes the counter for a period it holds nothing for", async () => { + const { stub } = meter(); + + // The cold-start case: a fresh object reconciling against a month of spend + // that D1 already knows about. + await stub.setFloor("llm_tokens_month", 7500, PERIOD); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(7500); + expect((await stub.reserve("llm_tokens_month", 3000, 10_000, PERIOD, NOW)).admitted).toBe( + false, + ); + }); + + it("ignores a floor for a period older than the one it holds", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 10_000, NEXT_PERIOD, NOW + 31 * 24 * HOUR_MS); + await stub.setFloor("llm_tokens_month", 9999, PERIOD); + + // 'YYYY-MM' sorts chronologically, so a late reconcile for a closed month + // cannot clobber the live one with last month's total. + expect((await stub.read(NEXT_PERIOD)).counts.llm_tokens_month).toBe(100); + }); + + it("ignores a malformed quantity", async () => { + const { stub } = meter(); + + await stub.reserve("llm_tokens_month", 100, 10_000, PERIOD, NOW); + await stub.setFloor("llm_tokens_month", Number.NaN, PERIOD); + await stub.setFloor("llm_tokens_month", -5, PERIOD); + + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(100); + }); + + it("ignores a floor for a rate meter, which has no aggregate", async () => { + const { stub } = meter(); + + await stub.setFloor("evaluations_per_hour", 500, PERIOD); + + expect((await stub.read(PERIOD)).counts).toEqual({}); + expect((await stub.reserve("evaluations_per_hour", 1, 5, PERIOD, NOW)).count).toBe(1); + }); +}); + +describe("UsageMeter evaluations_per_hour", () => { + it("admits up to the hourly limit and then refuses", async () => { + const { stub } = meter(); + + const outcomes = []; + for (let i = 0; i < 5; i++) { + outcomes.push(await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW)); + } + + expect(outcomes.map((o) => o.admitted)).toEqual([true, true, true, false, false]); + expect(outcomes.map((o) => o.count)).toEqual([1, 2, 3, 3, 3]); + }); + + it("sums short buckets across the window instead of resetting on the hour", async () => { + const { stub } = meter(); + // One minute before the top of an hour, so a fixed hourly bucket would roll + // between these two calls. + const beforeTheHour = NOW - MINUTE_MS; + + for (let i = 0; i < 12; i++) { + expect( + (await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, beforeTheHour)).admitted, + ).toBe(true); + } + + // The window slides, so the allowance spent at 11:59 is still spent at + // 12:00. A single hourly bucket would have admitted 12 more here — 24 in + // one minute against a limit of 12 per hour, the burst this meter exists + // to bound. + expect((await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW)).admitted).toBe(false); + expect( + (await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW + 30 * MINUTE_MS)).admitted, + ).toBe(false); + }); + + it("gives the allowance back once the window has fully passed", async () => { + const { stub } = meter(); + + for (let i = 0; i < 12; i++) { + await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW); + } + expect((await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW)).admitted).toBe(false); + + const later = await stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW + HOUR_MS); + + expect(later).toEqual({ admitted: true, count: 1 }); + }); + + it("is not reset by the month rolling", async () => { + const { stub } = meter(); + + for (let i = 0; i < 3; i++) { + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW); + } + + // A rate window that straddles midnight on the 1st is still the same + // window; the period the caller passes has nothing to do with it. + expect((await stub.reserve("evaluations_per_hour", 1, 3, NEXT_PERIOD, NOW)).admitted).toBe( + false, + ); + }); + + it("keeps the month's counters out of the rate window and vice versa", async () => { + const { stub } = meter(); + + await stub.reserve("evaluations_per_hour", 1, 10, PERIOD, NOW); + await stub.reserve("llm_tokens_month", 500, 10_000, PERIOD, NOW); + + // The rate meter is windowed, so it is deliberately not in `read`. + expect((await stub.read(PERIOD)).counts).toEqual({ llm_tokens_month: 500 }); + }); + + it("returns a rate reservation on a negative settle", async () => { + const { stub } = meter(); + + for (let i = 0; i < 3; i++) { + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW); + } + expect((await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW)).admitted).toBe(false); + + await stub.settle("evaluations_per_hour", -1, PERIOD, NOW); + + expect((await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW)).admitted).toBe(true); + }); + + it("admits when the clock is not a usable timestamp", async () => { + const { stub } = meter(); + + const outcome = await stub.reserve("evaluations_per_hour", 1, 1, PERIOD, Number.NaN); + + expect(outcome).toEqual({ admitted: true, count: 0 }); + }); +}); + +describe("UsageMeter cleanup", () => { + it("arms the alarm past the end of the live period", async () => { + const { stub, storages } = meter(); + + await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW); + + const alarm = await storages.get("user:usr_1")?.getAlarm(); + expect(alarm).toBeGreaterThan(Date.UTC(2026, 9, 1)); + }); + + it("erases a closed period on the alarm, so a monthly meter is not forever", async () => { + const { stub, instances, storages } = meter(); + const closed = "2020-01"; + + await stub.reserve("llm_tokens_month", 100, 1000, closed, Date.UTC(2020, 0, 15)); + const instance = instances.get("user:usr_1"); + const storage = storages.get("user:usr_1"); + if (!instance || !storage) throw new Error("instance not created"); + + await instance.alarm(); + + expect(await storage.get("period")).toBeUndefined(); + expect(await storage.getAlarm()).toBeNull(); + }); + + it("keeps a live hourly window when a closed period is swept", async () => { + const { stub, instances, storages } = meter(); + const closed = "2020-01"; + + await stub.reserve("llm_tokens_month", 100, 1000, closed, Date.UTC(2020, 0, 15)); + // The rate reservation is live NOW; the precedent's flat deleteAll would + // erase it along with the closed month and hand back a fresh burst. + for (let i = 0; i < 3; i++) { + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, Date.now()); + } + const instance = instances.get("user:usr_1"); + const storage = storages.get("user:usr_1"); + if (!instance || !storage) throw new Error("instance not created"); + + await instance.alarm(); + + expect(await storage.get("period")).toBeUndefined(); + expect(await storage.get("rate")).toBeDefined(); + expect(await storage.getAlarm()).not.toBeNull(); + expect((await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, Date.now())).admitted).toBe( + false, + ); + }); + + it("erases an expired rate window", async () => { + const { stub, instances, storages } = meter(); + + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, Date.now() - 4 * HOUR_MS); + const instance = instances.get("user:usr_1"); + const storage = storages.get("user:usr_1"); + if (!instance || !storage) throw new Error("instance not created"); + + await instance.alarm(); + + expect(await storage.get("rate")).toBeUndefined(); + expect(await storage.getAlarm()).toBeNull(); + }); + + it("purges everything for the deletion cascade", async () => { + const { stub, storages } = meter(); + + await stub.reserve("llm_tokens_month", 100, 1000, PERIOD, NOW); + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW); + await stub.purge(); + + const storage = storages.get("user:usr_1"); + expect(await storage?.get("period")).toBeUndefined(); + expect(await storage?.get("rate")).toBeUndefined(); + expect((await stub.read(PERIOD)).counts).toEqual({}); + }); +}); + +describe("usageMeterName", () => { + it("puts users and orgs in separate namespaces", () => { + // The deletion cascade relies on this: purging one org must be unable to + // name a person's counter, and vice versa. Without the prefix, that would + // rest on two id generators never colliding. + expect(usageMeterName("user", "abc")).toBe("user:abc"); + expect(usageMeterName("org", "abc")).toBe("org:abc"); + expect(usageMeterName("user", "abc")).not.toBe(usageMeterName("org", "abc")); + }); + + it("keeps two subjects' counters in different instances", async () => { + const { namespace, instances } = makeUsageMeters(); + const user = stubFor(namespace, usageMeterName("user", "abc")); + const org = stubFor(namespace, usageMeterName("org", "abc")); + + await user.reserve("llm_tokens_month", 900, 1000, PERIOD, NOW); + const orgOutcome = await org.reserve("llm_tokens_month", 900, 1000, PERIOD, NOW); + + expect(orgOutcome).toEqual({ admitted: true, count: 900 }); + expect([...instances.keys()].sort()).toEqual(["org:abc", "user:abc"]); + }); +}); + +describe("UsageMeter concurrency", () => { + it("holds the allowance when 50 reservations run at once", async () => { + const { stub } = meter(); + + const outcomes = await Promise.all( + Array.from({ length: 50 }, () => stub.reserve("llm_tokens_month", 10, 100, PERIOD, NOW)), + ); + + expect(outcomes.filter((o) => o.admitted)).toHaveLength(10); + // Counts are 10..100 with no repeats — no two admissions read the same value. + expect( + outcomes + .filter((o) => o.admitted) + .map((o) => o.count) + .sort((a, b) => a - b), + ).toEqual(Array.from({ length: 10 }, (_, i) => (i + 1) * 10)); + }); + + it("over-admits without the input gate, proving the assertion above is not vacuous", async () => { + // The same 50 reservations against an UNGATED instance: a read-then-write, + // which is what a KV counter (or a Worker-side read followed by a DO write) + // gives you. If this ever stops over-admitting, the harness has stopped + // modelling concurrency and the test above proves nothing. + const { stub } = meter({ gated: false }); + + const outcomes = await Promise.all( + Array.from({ length: 50 }, () => stub.reserve("llm_tokens_month", 10, 100, PERIOD, NOW)), + ); + + expect(outcomes.filter((o) => o.admitted).length).toBeGreaterThan(10); + }); + + it("holds an hourly rate limit under a concurrent burst", async () => { + const { stub } = meter(); + + const outcomes = await Promise.all( + Array.from({ length: 40 }, () => stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW)), + ); + + expect(outcomes.filter((o) => o.admitted)).toHaveLength(12); + }); + + it("over-admits the rate limit without the input gate", async () => { + const { stub } = meter({ gated: false }); + + const outcomes = await Promise.all( + Array.from({ length: 40 }, () => stub.reserve("evaluations_per_hour", 1, 12, PERIOD, NOW)), + ); + + expect(outcomes.filter((o) => o.admitted).length).toBeGreaterThan(12); + }); + + it("keeps subjects independent, so one busy subject cannot refuse another", async () => { + const { namespace } = makeUsageMeters(); + const a = stubFor(namespace, usageMeterName("user", "a")); + const b = stubFor(namespace, usageMeterName("user", "b")); + + const outcomes = await Promise.all([ + ...Array.from({ length: 20 }, () => a.reserve("llm_tokens_month", 10, 100, PERIOD, NOW)), + ...Array.from({ length: 20 }, () => b.reserve("llm_tokens_month", 10, 100, PERIOD, NOW)), + ]); + + expect(outcomes.filter((o) => o.admitted)).toHaveLength(20); + }); +}); + +describe("UsageMeter surface", () => { + it("exposes no fetch handler", async () => { + const { instances, stub } = meter(); + await stub.read(PERIOD); + const instance = instances.get("user:usr_1"); + + // Limits are caller-supplied, which is safe only while the class is + // reachable from Worker code alone. A `fetch` handler would make `limit` an + // attacker-controlled input. + expect((instance as unknown as { fetch?: unknown })?.fetch).toBeUndefined(); + }); +}); diff --git a/tests/wrangler-usage-meter-config.test.ts b/tests/wrangler-usage-meter-config.test.ts new file mode 100644 index 00000000..e7c9ffa4 --- /dev/null +++ b/tests/wrangler-usage-meter-config.test.ts @@ -0,0 +1,100 @@ +/// +import { describe, expect, it } from "vitest"; +// Loaded via Vite's raw import (not node:fs) so this type-checks under the +// Workers tsconfig, matching tests/wrangler-telemetry-config.test.ts. +import wranglerToml from "../wrangler.toml?raw"; + +/** + * Guards the UsageMeter's configuration against the two wrangler footguns that + * make a binding look configured and behave as absent. + * + * 1. **Named environments inherit nothing.** A `[[durable_objects.bindings]]` + * or a `[vars]` entry declared only at the top level never reaches + * `wrangler deploy --env production` — the failure that made the telemetry + * kill switch inert (#257, and the sibling guard in + * `wrangler-telemetry-config.test.ts`). An unbound `USAGE_METER` is worse + * than a wrong value: `env.USAGE_METER` is optional, so every counter would + * silently no-op forever with nothing in the logs. + * 2. **`script_name` is per environment.** Staging's classes live on + * `stratum-staging`; pointing its binding at `stratum` aims the staging + * Worker at production's Durable Objects. + * + * The migration tag is checked here only for presence and uniqueness — the + * append-only rule that actually matters lives in + * `tests/wrangler-migration-chain.test.ts`, which compares against main. + */ + +interface TomlSection { + name: string; + body: string; +} + +/** Split into `[section]` / `[[array]]` blocks, keeping each raw body. */ +function splitSections(toml: string): TomlSection[] { + const sections: TomlSection[] = []; + let current: TomlSection = { name: "", body: "" }; + for (const line of toml.split("\n")) { + const header = /^\s*\[{1,2}([^\]]+)\]{1,2}\s*$/.exec(line); + if (header?.[1]) { + sections.push(current); + current = { name: header[1], body: "" }; + continue; + } + current.body += `${line}\n`; + } + sections.push(current); + return sections; +} + +const sections = splitSections(wranglerToml); + +function bodiesOf(name: string): string[] { + return sections.filter((section) => section.name === name).map((section) => section.body); +} + +/** The value of `key` where it is set for real (a commented line is not set). */ +function settingValue(body: string, key: string): string | null { + const match = new RegExp(`^\\s*${key}\\s*=\\s*"([^"]*)"`, "m").exec(body); + return match?.[1] ?? null; +} + +/** Every environment that declares Durable Object bindings, and its script name. */ +const DO_SCOPES: readonly [scope: string, scriptName: string][] = [ + ["durable_objects.bindings", "stratum"], + ["env.production.durable_objects.bindings", "stratum"], + ["env.staging.durable_objects.bindings", "stratum-staging"], +]; + +describe("wrangler.toml UsageMeter configuration", () => { + it("declares the class in exactly one migration tag", () => { + const declaring = bodiesOf("migrations").filter((body) => body.includes('"UsageMeter"')); + expect(declaring).toHaveLength(1); + // SQLite-backed like RepoDO and MagicLinkRateLimiter; `new_classes` would + // create a key-value object the storage code here does not target. + expect(declaring[0]).toMatch(/new_sqlite_classes\s*=\s*\[[^\]]*"UsageMeter"/); + }); + + it.each(DO_SCOPES)("binds USAGE_METER in [%s] with script_name %s", (scope, scriptName) => { + const binding = bodiesOf(scope).find((body) => settingValue(body, "name") === "USAGE_METER"); + expect(binding, `no USAGE_METER binding in [[${scope}]]`).toBeDefined(); + expect(settingValue(binding as string, "class_name")).toBe("UsageMeter"); + expect(settingValue(binding as string, "script_name")).toBe(scriptName); + }); + + it.each(["vars", "env.production.vars", "env.staging.vars"])( + "declares the entitlements switch and documents its companions in [%s]", + (scope) => { + const body = bodiesOf(scope)[0]; + expect(body, `no [${scope}] block`).toBeDefined(); + // Spelled out rather than left unset, so the enforcement switch is + // visible in every environment instead of inferred from an absence. + expect(settingValue(body as string, "ENTITLEMENTS_ENFORCE")).toBe("0"); + // The other three ship commented (a URL, a secret's name, an operator + // allowlist), so what is asserted is that each block documents them — + // the same shape-based assertion the telemetry guard makes. + for (const key of ["BILLING_SERVICE_URL", "BILLING_SERVICE_SECRET", "LLM_PROVIDERS"]) { + expect(body as string, `${key} undocumented in [${scope}]`).toContain(key); + } + }, + ); +}); diff --git a/wrangler.toml b/wrangler.toml index e457002a..9ba23209 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -36,6 +36,15 @@ new_sqlite_classes = ["RepoDO"] tag = "v3" new_sqlite_classes = ["MagicLinkRateLimiter"] +# UsageMeter: serialized usage counters, one instance per ENFORCEMENT subject +# (PRD §4a — the acting user, or an org positively known to be paid). SQLite-backed +# like the two above. APPEND ONLY: tags are recorded on the deployed script, so +# editing or reordering one makes wrangler replay the chain and every deploy of +# that environment fails with error 10074 — see tests/wrangler-migration-chain.test.ts. +[[migrations]] +tag = "v4" +new_sqlite_classes = ["UsageMeter"] + # Observability - Workers Logs and Traces [observability] enabled = true @@ -90,6 +99,11 @@ name = "MAGIC_LINK_LIMITER" class_name = "MagicLinkRateLimiter" script_name = "stratum" +[[durable_objects.bindings]] +name = "USAGE_METER" +class_name = "UsageMeter" +script_name = "stratum" + [[queues.producers]] binding = "EVENTS_QUEUE" queue = "stratum-events" @@ -196,6 +210,37 @@ OAUTH_REDIRECT_URI = "http://localhost:8787/auth/github/callback" # BETA_GATE = "1" # REFERRAL_SERVICE_URL = "https://your-referral-service.example.com" +# --- Usage metering, entitlements, and evaluator BYOK ------------------------- +# Every var in this group must be declared PER ENVIRONMENT: named envs do NOT +# inherit top-level [vars] (the same footgun as the telemetry kill switch just +# below), so setting one here only affects a bare `wrangler deploy`. +# +# Plan/entitlements service (hosted-only). The hook is live only when BOTH +# BILLING_SERVICE_URL and the BILLING_SERVICE_SECRET secret are set — see +# `entitlementsEnabled` in src/billing/entitlements.ts, which treats a URL +# without a credential as OFF rather than half-on. Unset, the default for every +# self-hoster: nothing is fetched, every owner resolves to unlimited, and the +# usage counters record without limiting anything. +# BILLING_SERVICE_URL = "https://your-billing-service.example.com" +# +# "1" makes entitlement decisions REFUSE rather than only record. Declared "0" +# explicitly rather than left unset so the enforcement switch is visible in each +# environment's config instead of inferred from an absence — the same reason +# DEV_LOGIN_ENABLED is spelled out per env. Setting it WITHOUT +# BILLING_SERVICE_URL is enforcement that looks on and does nothing; +# `entitlementsConfigError` (src/middleware/config-guard.ts) logs exactly that. +ENTITLEMENTS_ENFORCE = "0" +# +# The evaluator's provider allowlist, and the ONLY place a provider endpoint may +# be named: a project's .stratum/policy.yaml may SELECT one of these by name and +# can never supply a `baseUrl`, so a project owner cannot aim the operator's +# Worker at a host of their choosing. An unrecognized name fails the policy +# closed — a merge gate that cannot run must not pass. Unset (the default) means +# Workers AI (`env.AI`) only, which is today's behavior. Credentials are NOT +# here: they live per project in `project_secrets`. Shape illustrated, parsed by +# the provider seam: +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' + # Instance-wide telemetry kill switch. Not a secret — a plain var, listed here # next to POSTHOG_API_KEY because they are read together. Named envs do NOT # inherit top-level vars, so uncommenting this alone only affects a bare @@ -211,6 +256,7 @@ OAUTH_REDIRECT_URI = "http://localhost:8787/auth/github/callback" # GOOGLE_CLIENT_SECRET = # EMAIL_FROM_ADDRESS = "noreply@yourdomain.com" # for magic link emails # REFERRAL_SERVICE_SECRET = # closed-beta gate shared secret (must match the referral service) +# BILLING_SERVICE_SECRET = # entitlements service shared secret; without it BILLING_SERVICE_URL is inert # DEPLOY_SECRET_KEY = # master key for per-project deploy secrets; rotating it makes every stored secret undecryptable # ============================================================================= @@ -279,6 +325,11 @@ name = "MAGIC_LINK_LIMITER" class_name = "MagicLinkRateLimiter" script_name = "stratum" +[[env.production.durable_objects.bindings]] +name = "USAGE_METER" +class_name = "UsageMeter" +script_name = "stratum" + [[env.production.queues.producers]] binding = "EVENTS_QUEUE" queue = "stratum-events" @@ -362,6 +413,17 @@ GOOGLE_REDIRECT_URI = "https://app.usestratum.dev/auth/google/callback" # the gate ON so the invite path stays exercised end-to-end. BETA_GATE = "0" REFERRAL_SERVICE_URL = "https://referral.usestratum.dev" +# Usage metering / entitlements / BYOK. Declared here as well as at top level +# because named envs inherit no vars; see the top-level [vars] block for what +# each one does. BILLING_SERVICE_SECRET is a secret and is set out-of-band with +# `wrangler secret put BILLING_SERVICE_SECRET --env production` — secrets are +# per-environment too, so a top-level one never reaches this deploy. +# BILLING_SERVICE_URL = "https://" +# Observe-only: the limits are measured for a period before they gate anything +# (PRD Goal 3 — metering is always on, enforcement is a separate switch). +ENTITLEMENTS_ENFORCE = "0" +# Operator-configured LLM providers; unset = Workers AI only. +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' # ============================================================================= # [env.staging] @@ -426,6 +488,11 @@ name = "MAGIC_LINK_LIMITER" class_name = "MagicLinkRateLimiter" script_name = "stratum-staging" +[[env.staging.durable_objects.bindings]] +name = "USAGE_METER" +class_name = "UsageMeter" +script_name = "stratum-staging" + [[env.staging.r2_buckets]] binding = "REPO_OBJECTS" bucket_name = "stratum-repo-objects-staging" @@ -506,3 +573,12 @@ EMAIL_FROM_ADDRESS = "noreply@staging.app.usestratum.dev" # live stratum-landing Worker (see [env.production.vars] cutover note). BETA_GATE = "1" REFERRAL_SERVICE_URL = "https://referral.usestratum.dev" +# Usage metering / entitlements / BYOK; declared per env (see [env.production.vars] +# and the top-level [vars] block). BILLING_SERVICE_SECRET is set with +# `wrangler secret put BILLING_SERVICE_SECRET --env staging`. +# BILLING_SERVICE_URL = "https://" +# Staging is where enforcement gets validated first, as with every other flag +# here — but only once the billing service is live and a period of observation +# has produced real numbers. Off until then. +ENTITLEMENTS_ENFORCE = "0" +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' From 67e9acc156e27b4e5e35a04eee4d962aa042e8ff Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 10:38:17 +0000 Subject: [PATCH 11/26] Fix two meter bugs, and put the LLM evaluator behind a provider seam MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 6, plus the review findings on Task 5. Both bugs were reproduced by the reviewer against the real class, not inferred. A rate refund landed on the bucket the settle fell in, which is empty whenever the settle outlives the reservation's bucket — the normal case, since buckets are minutes and an evaluation is seconds to minutes. The clamp to zero then discarded it and the original charge sat in its own bucket for the rest of the window, so a provider call that failed three minutes in still burned its hourly slot. Refunds now draw down the live buckets oldest-first, which is where the charge actually sits, and a refund larger than the window holds is logged rather than silently clamped. The doc comment asserting the sum was corrected either way was simply wrong. A finite but absurd clock was worse. Microseconds mistaken for milliseconds passed the finiteness guard, landed a bucket a thousand times into the future where it counted against every later window, and armed the cleanup alarm past the year 58000 — refusing the subject permanently and stranding its storage, which is exactly the "metering causes the outage" failure the malformed-input deviation exists to prevent, arriving through a value that looked fine. The clock is now checked against an absolute plausibility band rather than finiteness, and the alarm has a horizon, since armCleanup only ever raises and could never be talked back down. Chasing that surfaced a fake that was kinder than production: deleteAll does not clear a pending alarm in workerd, but the test double did, so purge() looked correct while leaving a purged subject armed to wake. The double now matches, alarms are consumed by firing where a test fires one, and purge and the sweep drop theirs explicitly. Also corrected: wrangler.toml described LLM_PROVIDERS in the present tense as the allowlist a policy selects from, which nothing reads until Task 7 — an operator following it would have set a var that does nothing. Marked reserved, with the test that asserts its presence told not to pin the prose. The provider seam itself is small and deliberately inert: WorkersAi, Anthropic and OpenAI-compatible implementations behind one interface, with only WorkersAi wired in. Both HTTP providers take base URL and key from their caller and read no env, because where those come from is Task 7's problem and a policy-supplied URL is a request-forgery primitive. Provider error messages carry status and shape only, never a response body — a body can echo the failing credential and can quote the prompt, which contains the diff. A test drives a 401 whose body contains a key and asserts it reaches no result. Typecheck clean; lint clean; 3920 tests pass, up 43. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- src/evaluation/index.ts | 1 + src/evaluation/llm-evaluator.ts | 55 +- src/evaluation/llm-provider.ts | 305 +++++++++++ src/queue/usage-meter.ts | 181 ++++++- src/services/change-flow.ts | 4 +- tests/costs.test.ts | 13 +- tests/helpers/fake-durable-object.ts | 9 +- tests/llm-evaluator.test.ts | 51 +- tests/llm-provider.test.ts | 611 ++++++++++++++++++++++ tests/magic-link-rate-limit.test.ts | 4 + tests/usage-meter.test.ts | 53 ++ tests/wrangler-usage-meter-config.test.ts | 5 +- wrangler.toml | 19 +- 13 files changed, 1227 insertions(+), 84 deletions(-) create mode 100644 src/evaluation/llm-provider.ts create mode 100644 tests/llm-provider.test.ts diff --git a/src/evaluation/index.ts b/src/evaluation/index.ts index 5f61b7a5..c2e2b2f7 100644 --- a/src/evaluation/index.ts +++ b/src/evaluation/index.ts @@ -4,5 +4,6 @@ export * from "./webhook-evaluator"; export * from "./composite-evaluator"; export * from "./policy-loader"; export * from "./sandbox-evaluator"; +export * from "./llm-provider"; export * from "./llm-evaluator"; export * from "./secret-scanner"; diff --git a/src/evaluation/llm-evaluator.ts b/src/evaluation/llm-evaluator.ts index 71bc941d..42ef8068 100644 --- a/src/evaluation/llm-evaluator.ts +++ b/src/evaluation/llm-evaluator.ts @@ -1,9 +1,10 @@ -import type { AiBinding } from "../types"; import type { AppError } from "../utils/errors"; import { ExternalServiceError } from "../utils/errors"; import type { Logger } from "../utils/logger"; import type { Result } from "../utils/result"; import { err, ok } from "../utils/result"; +import type { LlmProvider } from "./llm-provider"; +import { LlmProviderResponseError } from "./llm-provider"; import { sanitizePolicy } from "./sanitize-policy"; import type { EvalPolicy, EvalResult, Evaluator } from "./types"; @@ -30,7 +31,7 @@ const SYSTEM_PROMPT = [ ].join(" "); export class LLMEvaluator implements Evaluator { - constructor(private ai: AiBinding) {} + constructor(private provider: LlmProvider) {} async evaluate( diff: string, @@ -81,23 +82,39 @@ export class LLMEvaluator implements Evaluator { ]; const promptChars = messages.reduce((sum, m) => sum + m.content.length, 0); - const raw = await this.ai.run(model, { messages }); - - if (raw instanceof ReadableStream) { - logger.error("LLM evaluation failed: unexpected stream response"); - return ok({ - score: 0, - passed: false, - reason: "LLM evaluator error: unexpected stream response", - }); + const runResult = await this.provider.run(model, messages); + + if (!runResult.success) { + // A response that arrived but cannot be used fails the gate closed with + // a readable reason; a transport failure stays an error the caller + // surfaces. Either way the provider's own message is metadata only — + // it is copied into a user-visible reason by `runEvaluation`. + if (runResult.error instanceof LlmProviderResponseError) { + logger.error(`LLM evaluation failed: ${runResult.error.message}`); + return ok({ + score: 0, + passed: false, + reason: `LLM evaluator error: ${runResult.error.message}`, + }); + } + logger.error("LLM evaluation failed", runResult.error); + return err(runResult.error); } - const responseText = raw.response; - // Workers AI does not report token usage; ~4 chars/token is the standard estimate. - const estimatedTokens = Math.ceil((promptChars + (responseText?.length ?? 0)) / 4); - const costs: EvalResult["costs"] = [ - { kind: "llm_tokens", quantity: estimatedTokens, estimated: true }, - ]; + const responseText = runResult.data.text; + const usage = runResult.data.usage; + // Real counts when the provider reports them, otherwise ~4 chars/token — + // the distinction is per provider, not per source. Workers AI reports no + // usage at all, so that path stays estimated. + const costs: EvalResult["costs"] = usage + ? [{ kind: "llm_tokens", quantity: usage.inputTokens + usage.outputTokens }] + : [ + { + kind: "llm_tokens", + quantity: Math.ceil((promptChars + responseText.length) / 4), + estimated: true, + }, + ]; // A gate whose verdict can't be read has not passed. Never infer a score // from prose ("LGTM") — that let unparseable output half-approve a merge. @@ -110,7 +127,7 @@ export class LLMEvaluator implements Evaluator { passed: false, reason: `LLM evaluator failed closed: ${why}`, issues: [ - `Model response (${responseText?.length ?? 0} chars) was not a valid verdict object`, + `Model response (${responseText.length} chars) was not a valid verdict object`, ...(truncationNote ? [truncationNote] : []), ], costs, @@ -119,7 +136,7 @@ export class LLMEvaluator implements Evaluator { let parsed: { score: unknown; passed: unknown; reason: unknown; issues?: unknown }; try { - parsed = JSON.parse(responseText ?? "") as { + parsed = JSON.parse(responseText) as { score: unknown; passed: unknown; reason: unknown; diff --git a/src/evaluation/llm-provider.ts b/src/evaluation/llm-provider.ts new file mode 100644 index 00000000..c26c049d --- /dev/null +++ b/src/evaluation/llm-provider.ts @@ -0,0 +1,305 @@ +import type { AiBinding } from "../types"; +import { AppError, ExternalServiceError } from "../utils/errors"; +import type { Result } from "../utils/result"; +import { err, ok } from "../utils/result"; + +/** + * One turn of the conversation handed to a provider. Structurally identical to + * the shape `AiBinding.run` already takes, so `WorkersAiProvider` forwards it + * unchanged. + */ +export interface Message { + role: string; + content: string; +} + +/** Token counts a provider actually reported. Never an estimate — see `LlmResponse.usage`. */ +export interface LlmUsage { + inputTokens: number; + outputTokens: number; +} + +export interface LlmResponse { + /** The assistant's text. Empty string when the provider returned a well-formed + * response carrying no text, which the evaluator then fails closed on. */ + text: string; + /** + * Present only when the provider reported real token counts. Its absence is + * load-bearing: the evaluator falls back to the `~4 chars/token` estimate and + * marks the cost sample `estimated`. That distinction is **per provider, not + * per source** — Workers AI reports nothing, and an OpenAI-compatible + * endpoint that omits `usage` is estimated too. + */ + usage?: LlmUsage; +} + +/** + * The seam between the LLM evaluator and whoever actually runs the model. + * + * Errors are values: a provider never throws across this boundary. Two error + * shapes are distinguished, because they mean different things to the gate: + * `LlmProviderResponseError` for a response that arrived but cannot be used + * (which fails the evaluation closed with a readable reason), and any other + * `AppError` for a transport failure. + */ +export interface LlmProvider { + run(model: string, messages: Message[]): Promise>; +} + +/** + * The provider answered, but the answer is unusable — a non-2xx status, a body + * that is not JSON, a JSON body of the wrong shape, or a stream where a + * completion was expected. + * + * **The message is metadata only.** It reaches the user: `runEvaluation` + * (`services/change-flow.ts`) copies `error.message` straight into the recorded + * `EvalResult.reason`. A provider's error body can quote the request — which + * contains the diff — and some providers echo the credential that failed, so + * the body is never interpolated into it. Status codes and byte counts only. + */ +export class LlmProviderResponseError extends AppError { + constructor(message: string) { + super(message, "LLM_PROVIDER_RESPONSE", 502); + this.name = "LlmProviderResponseError"; + } +} + +function transportError(error: unknown): AppError { + const message = error instanceof Error ? error.message : String(error); + return new ExternalServiceError( + "LLM", + message, + error instanceof Error ? error : undefined, + ) as AppError; +} + +/** + * The Workers AI binding, and the default. Behaviour here is deliberately + * identical to what `LLMEvaluator` did inline before the seam existed: + * a `ReadableStream` is rejected rather than consumed, and a response object + * with no `response` field yields `""` so the evaluator fails it closed as + * unparseable output rather than treating it as a distinct condition. + * + * It reports no token usage, which is why the Workers AI path keeps the + * `~4 chars/token` estimate and `estimated: true`. + */ +export class WorkersAiProvider implements LlmProvider { + constructor(private ai: AiBinding) {} + + async run(model: string, messages: Message[]): Promise> { + let raw: Awaited>; + try { + raw = await this.ai.run(model, { messages }); + } catch (error) { + return err(transportError(error)); + } + + if (raw instanceof ReadableStream) { + return err(new LlmProviderResponseError("unexpected stream response")); + } + + return ok({ text: raw.response ?? "" }); + } +} + +/** How long a hosted provider gets before the gate gives up on it. */ +const DEFAULT_TIMEOUT_MS = 30_000; + +export interface HttpProviderOptions { + /** Abort the request after this many milliseconds. */ + timeoutMs?: number; +} + +/** + * Credentials and endpoints for both HTTP providers are **constructor + * arguments supplied by the caller** — never read from `env` here, and never + * taken from `.stratum/policy.yaml`. + * + * The omission is deliberate, not a gap. A base URL a repository's own policy + * file can choose turns this Worker into a request-forgery primitive aimed at + * whatever host the policy names, and the prompt it would receive contains the + * diff (PRD §7). Resolving a provider name against an operator-configured + * allowlist is Task 7's job; these classes only take what they are handed. + */ +abstract class HttpLlmProvider implements LlmProvider { + constructor( + protected readonly baseUrl: string, + protected readonly apiKey: string, + protected readonly options: HttpProviderOptions = {}, + ) {} + + abstract run(model: string, messages: Message[]): Promise>; + + /** + * POST a JSON body and parse a JSON response, converting every failure into + * a `Result`. Returns the decoded body as `unknown`: each provider validates + * its own shape, because a body that parsed is not a body that can be read. + */ + protected async post(path: string, headers: Record, body: unknown) { + const controller = new AbortController(); + const timer = setTimeout( + () => controller.abort(), + this.options.timeoutMs ?? DEFAULT_TIMEOUT_MS, + ); + + try { + const response = await fetch(`${this.baseUrl.replace(/\/+$/, "")}${path}`, { + method: "POST", + headers: { "content-type": "application/json", ...headers }, + body: JSON.stringify(body), + signal: controller.signal, + }); + + if (!response.ok) { + // The status, never the body: an error body can echo the API key that + // failed, and this message becomes a user-visible evaluation reason. + return err( + new LlmProviderResponseError(`provider returned HTTP ${response.status}`) as AppError, + ); + } + + try { + return ok((await response.json()) as unknown); + } catch { + return err( + new LlmProviderResponseError("provider response was not valid JSON") as AppError, + ); + } + } catch (error) { + return err(transportError(error)); + } finally { + clearTimeout(timer); + } + } +} + +/** + * Anthropic's Messages API: `POST {baseUrl}/messages`, `x-api-key` plus the + * required `anthropic-version` header, assistant text in `content[]` blocks of + * type `text`, real token counts in `usage.input_tokens` / `usage.output_tokens`. + * + * Docs: https://platform.claude.com/docs/en/api/messages/create and + * https://platform.claude.com/docs/en/api/overview (headers table). + */ +export class AnthropicProvider extends HttpLlmProvider { + /** Pinned rather than configurable: the API is versioned by this header, and + * a value the caller could vary is a shape this code has not been written for. */ + static readonly API_VERSION = "2023-06-01"; + /** The evaluator asks for one small JSON verdict; `max_tokens` is required. */ + static readonly MAX_TOKENS = 1024; + + async run(model: string, messages: Message[]): Promise> { + // Anthropic takes the system prompt as a top-level parameter, not as a + // message with role "system" — a system-role message is rejected. + const system = messages + .filter((m) => m.role === "system") + .map((m) => m.content) + .join("\n\n"); + const turns = messages.filter((m) => m.role !== "system"); + + const posted = await this.post( + "/messages", + { + "x-api-key": this.apiKey, + "anthropic-version": AnthropicProvider.API_VERSION, + }, + { + model, + max_tokens: AnthropicProvider.MAX_TOKENS, + ...(system.length > 0 ? { system } : {}), + messages: turns, + }, + ); + if (!posted.success) return posted; + + const body = posted.data as { + content?: unknown; + usage?: { input_tokens?: unknown; output_tokens?: unknown }; + } | null; + + if (body === null || typeof body !== "object" || !Array.isArray(body.content)) { + return err( + new LlmProviderResponseError("provider response had no content array") as AppError, + ); + } + + const text = body.content + .filter( + (block): block is { type: "text"; text: string } => + typeof block === "object" && + block !== null && + (block as { type?: unknown }).type === "text" && + typeof (block as { text?: unknown }).text === "string", + ) + .map((block) => block.text) + .join(""); + + return ok({ text, ...usageOrNothing(body.usage?.input_tokens, body.usage?.output_tokens) }); + } +} + +/** + * The OpenAI Chat Completions shape, which OpenAI, OpenRouter, Groq, Together + * and a self-hosted vLLM server all speak: `POST {baseUrl}/chat/completions` + * with `Authorization: Bearer`, assistant text at `choices[0].message.content`, + * token counts at `usage.prompt_tokens` / `usage.completion_tokens`. + * + * Docs: https://developers.openai.com/api/reference/chat-completions/overview, + * https://openrouter.ai/docs/api-reference/chat-completion, + * https://console.groq.com/docs/api-reference, + * https://docs.vllm.ai/en/latest/serving/openai_compatible_server/. + */ +export class OpenAiCompatibleProvider extends HttpLlmProvider { + async run(model: string, messages: Message[]): Promise> { + // No `max_tokens`: the four vendors this one endpoint serves disagree about + // it (newer OpenAI models take `max_completion_tokens` and reject + // `max_tokens`), and an unrecognised field is a 400 from the strict ones. + const posted = await this.post( + "/chat/completions", + { authorization: `Bearer ${this.apiKey}` }, + { model, messages }, + ); + if (!posted.success) return posted; + + const body = posted.data as { + choices?: unknown; + usage?: { prompt_tokens?: unknown; completion_tokens?: unknown }; + } | null; + + if (body === null || typeof body !== "object" || !Array.isArray(body.choices)) { + return err( + new LlmProviderResponseError("provider response had no choices array") as AppError, + ); + } + + const content = (body.choices[0] as { message?: { content?: unknown } } | undefined)?.message + ?.content; + if (typeof content !== "string") { + return err( + new LlmProviderResponseError("provider response had no assistant message text") as AppError, + ); + } + + return ok({ + text: content, + ...usageOrNothing(body.usage?.prompt_tokens, body.usage?.completion_tokens), + }); + } +} + +/** + * Report usage only when both counts are real numbers. A partial or malformed + * `usage` object is treated as absent rather than half-trusted: a wrong token + * count recorded as exact is worse than an estimate labelled as one. + */ +function usageOrNothing(input: unknown, output: unknown): { usage?: LlmUsage } { + if ( + typeof input !== "number" || + typeof output !== "number" || + !Number.isFinite(input) || + !Number.isFinite(output) + ) { + return {}; + } + return { usage: { inputTokens: input, outputTokens: output } }; +} diff --git a/src/queue/usage-meter.ts b/src/queue/usage-meter.ts index 54a406e3..69c07bde 100644 --- a/src/queue/usage-meter.ts +++ b/src/queue/usage-meter.ts @@ -63,9 +63,15 @@ interface RateWindow { * The granularity the window is summed at. SHORT on purpose: one bucket per * window is a fixed window, and a fixed window admits the whole allowance in * the first second of every new one — twice the limit across the boundary, - * which is the burst the meter exists to bound. Twelve buckets make the - * window slide, so an hour's allowance is an hour's allowance wherever the - * caller starts. + * which is the burst the meter exists to bound. + * + * Twelve buckets do not eliminate that, they bound it: this is the ordinary + * bucketed approximation of a sliding window, so a caller spending the whole + * allowance at the very end of one bucket gets it back one bucket-width + * before the hour is out — up to 2x the limit inside 55 minutes rather than + * inside a second. Halving `bucketSeconds` halves the error and doubles the + * stored keys. Say the bound out loud rather than claiming an exactness the + * arithmetic does not have. */ bucketSeconds: number; } @@ -212,7 +218,10 @@ export class UsageMeter extends DurableObject { // (NaN never compares below a limit again) would block the subject for // the rest of the month, which is the outcome the deviation exists to // avoid. The limit still applies to what is already counted. - warn("Usage meter estimate is not a usable quantity; reserving zero", { meter, estimate }); + warn("Usage meter estimate is not a usable quantity; reserving zero", { + meter, + estimate: String(estimate), + }); } const quantity = usable ?? 0; if (isRateMeter(meter)) return this.reserveRate(meter, quantity, limit, nowMs); @@ -254,7 +263,18 @@ export class UsageMeter extends DurableObject { return; } const record = await this.ctx.storage.get(PERIOD_KEY); - if (record?.period !== period) return; + if (record?.period !== period) { + // The period rolled between reserve and settle, so the charge being + // refunded is not in the live record. Logged rather than dropped: a + // reservation that outlives its month leaves the old counter high, and + // that has to be diagnosable from Workers Logs. + warn("Usage meter settle ignored: period has rolled since the reservation", { + meter, + settlePeriod: period, + livePeriod: record?.period ?? null, + }); + return; + } const current = record.counts[meter] ?? 0; await this.ctx.storage.put(PERIOD_KEY, { period, @@ -348,6 +368,11 @@ export class UsageMeter extends DurableObject { */ async purge(): Promise { await this.ctx.storage.deleteAll(); + // `deleteAll` does NOT clear a pending alarm in workerd, so without this a + // purged subject still gets one wake-up on a DO with no state. The test + // fake clears it inside deleteAll, which is exactly why this cannot be + // left to the fake to prove. + await this.ctx.storage.deleteAlarm(); } /** @@ -390,9 +415,13 @@ export class UsageMeter extends DurableObject { // Nothing left to expire: erase whatever remains (a key written by an // older shape of this class, say) and leave no alarm armed, so the object - // costs nothing until it is next used. - if (nextAlarm === null) await this.ctx.storage.deleteAll(); - else await this.ctx.storage.setAlarm(nextAlarm); + // costs nothing until it is next used. `deleteAll` does NOT clear a pending + // alarm in workerd, so the alarm is dropped explicitly — relying on the + // sweep to do it left the object waking forever on empty storage. + if (nextAlarm === null) { + await this.ctx.storage.deleteAll(); + await this.ctx.storage.deleteAlarm(); + } else await this.ctx.storage.setAlarm(nextAlarm); } private async reservePeriod( @@ -413,7 +442,10 @@ export class UsageMeter extends DurableObject { const current = counts[meter] ?? 0; const decision = limitAdmits(current, quantity, limit); if (decision.malformed) { - warn("Usage meter reserve admitted: limit is not a usable limit", { meter, limit }); + warn("Usage meter reserve admitted: limit is not a usable limit", { + meter, + limit: String(limit), + }); } if (!decision.admitted) return { admitted: false, count: current }; const next = current + quantity; @@ -432,8 +464,15 @@ export class UsageMeter extends DurableObject { limit: number, nowMs: number, ): Promise { - if (!Number.isFinite(nowMs)) { - warn("Usage meter reserve admitted: clock is not a usable timestamp", { meter, nowMs }); + if (!isUsableClock(nowMs)) { + // Admit and count nothing, as with every other malformed input. Counting + // it is what the earlier version did, and a bucket ~1000x in the future + // (microseconds mistaken for milliseconds) then counted against every + // window from here on, refusing the subject permanently. + warn("Usage meter reserve admitted: clock is not a usable timestamp", { + meter, + nowMs: String(nowMs), + }); return { admitted: true, count: 0 }; } const window = RATE_WINDOWS[meter]; @@ -443,7 +482,10 @@ export class UsageMeter extends DurableObject { const total = sum(live); const decision = limitAdmits(total, quantity, limit); if (decision.malformed) { - warn("Usage meter reserve admitted: limit is not a usable limit", { meter, limit }); + warn("Usage meter reserve admitted: limit is not a usable limit", { + meter, + limit: String(limit), + }); } if (!decision.admitted) return { admitted: false, count: total }; const key = String(bucket); @@ -457,24 +499,63 @@ export class UsageMeter extends DurableObject { /** * Rate meters are reserve-with-no-settle by convention, so this exists only - * for the caller that reserved an evaluation which then never ran. The delta - * lands on the live bucket rather than on whichever bucket the reservation - * was taken in: the limit reads the window's SUM, so the sum is corrected - * either way, and per-bucket floors of zero keep a refund from inventing - * headroom that was never reserved. + * for the caller that reserved an evaluation which then never ran. + * + * A refund DRAWS DOWN the live buckets rather than landing on the current + * one. The earlier version added the delta to the bucket `nowMs` falls in, + * which is empty whenever the settle outlives the reservation's bucket — the + * normal case, since buckets are minutes and an LLM gate is seconds to + * minutes. `Math.max(0, …)` then clamped the refund away and the original + * charge sat in its own bucket for the rest of the window, so an evaluation + * that failed three minutes after it was reserved still burned its slot for + * the next hour. + * + * Oldest bucket first: the reservation is by definition older than its + * settle, so that is where the charge most likely sits, and it is the choice + * that cannot free capacity a later reservation is still holding. Per-bucket + * floors of zero keep a refund from inventing headroom that was never taken. */ private async settleRate(meter: RateMeterKey, delta: number, nowMs: number): Promise { - if (!Number.isFinite(nowMs)) { - warn("Usage meter settle ignored: clock is not a usable timestamp", { meter, nowMs }); + if (!isUsableClock(nowMs)) { + warn("Usage meter settle ignored: clock is not a usable timestamp", { + meter, + nowMs: String(nowMs), + }); return; } const window = RATE_WINDOWS[meter]; const bucket = bucketId(nowMs, window.bucketSeconds); const record = await this.ctx.storage.get(RATE_KEY); - if (!record) return; + if (!record) { + warn("Usage meter rate settle ignored: no record for this subject", { meter, delta }); + return; + } const live = liveBuckets(record.meters[meter], bucket, window); - const key = String(bucket); - live[key] = Math.max(0, (live[key] ?? 0) + delta); + + if (delta >= 0) { + const key = String(bucket); + live[key] = (live[key] ?? 0) + delta; + } else { + let owed = -delta; + for (const key of Object.keys(live).sort((a, b) => Number(a) - Number(b))) { + if (owed <= 0) break; + const taken = Math.min(owed, live[key] ?? 0); + live[key] = (live[key] ?? 0) - taken; + owed -= taken; + } + if (owed > 0) { + // More was given back than the window is holding — a settle without a + // matching reserve, or one whose reservation has already aged out. + // Refusing to go negative is right; saying so is what makes it + // diagnosable rather than a counter that quietly disagrees with D1. + warn("Usage meter rate refund exceeded the window's charge", { + meter, + delta, + unapplied: owed, + }); + } + } + await this.ctx.storage.put(RATE_KEY, { meters: { ...record.meters, [meter]: live }, }); @@ -486,11 +567,52 @@ export class UsageMeter extends DurableObject { * `alarm()` re-arms for whatever it did not erase. */ private async armCleanup(atMs: number): Promise { + // Capped, because `armCleanup` only ever RAISES the alarm: one bad input + // that armed it centuries out could never be lowered by a later correct + // call, and the subject's storage would never be swept. + const ceiling = Date.now() + MAX_ALARM_HORIZON_MS; + const at = Math.min(atMs, ceiling); + if (at !== atMs) { + warn("Usage meter cleanup alarm clamped to the horizon", { requestedAtMs: atMs, atMs: at }); + } const existing = await this.ctx.storage.getAlarm(); - if (existing === null || existing < atMs) await this.ctx.storage.setAlarm(atMs); + if (existing === null || existing < at) await this.ctx.storage.setAlarm(at); } } +/** + * The band a caller's `nowMs` must fall in to be a timestamp at all. + * + * Deliberately an ABSOLUTE range rather than a tolerance around this object's + * own clock. What this rejects is the wrong UNIT or the wrong epoch — seconds + * (which land in 1970), microseconds or nanoseconds (which land tens of + * thousands of years out) — not skew, which is real and must keep working. A + * tolerance around `Date.now()` would also make the check depend on wall time, + * so a test with a fixed synthetic clock would pass or fail by the calendar. + * + * The failure this closes: a microsecond value passes `Number.isFinite`, lands + * a bucket ~1000x into the future where it counts against every later window + * forever, and re-arms the cleanup alarm past the year 58000. The subject's + * evaluations are refused permanently and its storage is stranded — the same + * "metering causes the outage" failure the malformed-input deviation exists to + * prevent, arriving through a value that looked finite. + */ +const MIN_PLAUSIBLE_MS = Date.UTC(2020, 0, 1); +const MAX_PLAUSIBLE_MS = Date.UTC(2200, 0, 1); + +/** + * The furthest ahead a cleanup alarm may ever be set. Longest month plus slack: + * nothing this object sweeps has a lifetime beyond one period, so an alarm + * further out than this is arithmetic on a bad input, and `armCleanup` only + * ever RAISES — so one bad value could never be lowered by a later good one. + */ +const MAX_ALARM_HORIZON_MS = 40 * 24 * 60 * 60 * 1000; + +/** Is `nowMs` a timestamp this object can do window arithmetic with? */ +function isUsableClock(nowMs: number): boolean { + return Number.isFinite(nowMs) && nowMs >= MIN_PLAUSIBLE_MS && nowMs <= MAX_PLAUSIBLE_MS; +} + function isRateMeter(meter: UsageMeterKey): meter is RateMeterKey { return meter in RATE_WINDOWS; } @@ -610,10 +732,15 @@ function pruneRates( let meterLogger: Logger | undefined; /** - * Every deviation from "count it" is logged, because each one silently ADMITS - * (see the class doc comment) and AGENTS.md forbids swallowing a failure. A - * limit that stopped limiting must be visible in Workers Logs, not inferred - * from a bill. + * Every path that admits without counting, and every settle that cannot be + * applied, logs. Each one silently ADMITS or silently keeps a charge (see the + * class doc comment), and AGENTS.md forbids swallowing a failure. A limit that + * stopped limiting, or a refund that never landed, must be visible in Workers + * Logs rather than inferred from a bill. + * + * Values that may be `NaN` or an infinity are stringified by the caller: JSON + * collapses all three to `null`, which would make the log unable to report the + * exact value it exists to report. */ function warn(message: string, meta: Record): void { meterLogger ??= createLogger({ component: "usage-meter" }); diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index 89cd71d9..af4daec7 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -5,6 +5,7 @@ import { SandboxEvaluator, SecretScanEvaluator, WebhookEvaluator, + WorkersAiProvider, diffTouchesProtectedConfig, loadPolicy, } from "../evaluation"; @@ -174,7 +175,8 @@ export function buildEvaluators( case "webhook": return [{ type: "webhook", evaluator: new WebhookEvaluator() }]; case "llm": - if (env.AI) return [{ type: "llm", evaluator: new LLMEvaluator(env.AI) }]; + if (env.AI) + return [{ type: "llm", evaluator: new LLMEvaluator(new WorkersAiProvider(env.AI)) }]; return [ { type: "llm", diff --git a/tests/costs.test.ts b/tests/costs.test.ts index 5c7481ec..a35b7876 100644 --- a/tests/costs.test.ts +++ b/tests/costs.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import { LLMEvaluator } from "../src/evaluation/llm-evaluator"; +import { WorkersAiProvider } from "../src/evaluation/llm-provider"; import { SandboxEvaluator } from "../src/evaluation/sandbox-evaluator"; import type { EvalPolicy } from "../src/evaluation/types"; import { getChangeCostSummary, recordCosts } from "../src/storage/costs"; @@ -131,7 +132,11 @@ describe("evaluator cost reporting", () => { }; const policy: EvalPolicy = { evaluators: [{ type: "llm" }] }; - const result = await new LLMEvaluator(ai).evaluate("diff content", policy, mockLogger); + const result = await new LLMEvaluator(new WorkersAiProvider(ai)).evaluate( + "diff content", + policy, + mockLogger, + ); expect(result.success).toBe(true); if (!result.success) return; expect(result.data.costs).toHaveLength(1); @@ -147,7 +152,11 @@ describe("evaluator cost reporting", () => { }; const policy: EvalPolicy = { evaluators: [{ type: "llm" }] }; - const result = await new LLMEvaluator(ai).evaluate("diff", policy, mockLogger); + const result = await new LLMEvaluator(new WorkersAiProvider(ai)).evaluate( + "diff", + policy, + mockLogger, + ); expect(result.success).toBe(true); if (!result.success) return; expect(result.data.costs?.[0]?.kind).toBe("llm_tokens"); diff --git a/tests/helpers/fake-durable-object.ts b/tests/helpers/fake-durable-object.ts index 0b841d58..dbe0fe89 100644 --- a/tests/helpers/fake-durable-object.ts +++ b/tests/helpers/fake-durable-object.ts @@ -18,6 +18,7 @@ export interface FakeDurableObjectStorage { deleteAll(): Promise; setAlarm(scheduledTime: number): Promise; getAlarm(): Promise; + deleteAlarm(): Promise; } export interface FakeDurableObjects { @@ -54,14 +55,20 @@ function makeStorage(): FakeDurableObjectStorage { values.set(key, structuredClone(value)); }, delete: async (key: string) => values.delete(key), + // Does NOT clear the alarm, matching workerd: `deleteAll` erases stored + // values and leaves a pending alarm armed, so a purged object still wakes + // once. Clearing it here would hide that from every test — a fake that is + // kinder than production proves the wrong thing. deleteAll: async () => { values.clear(); - alarm = null; }, setAlarm: async (scheduledTime: number) => { alarm = scheduledTime; }, getAlarm: async () => alarm, + deleteAlarm: async () => { + alarm = null; + }, }; } diff --git a/tests/llm-evaluator.test.ts b/tests/llm-evaluator.test.ts index 317adc42..95752d2a 100644 --- a/tests/llm-evaluator.test.ts +++ b/tests/llm-evaluator.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import { LLMEvaluator } from "../src/evaluation/llm-evaluator"; +import { WorkersAiProvider } from "../src/evaluation/llm-provider"; import type { EvalPolicy } from "../src/evaluation/types"; import type { AiBinding } from "../src/types"; import type { Logger } from "../src/utils/logger"; @@ -32,7 +33,7 @@ describe("LLMEvaluator — valid JSON responses", () => { const ai = makeMockAi( JSON.stringify({ score: 0.9, passed: true, reason: "Looks good", issues: [] }), ); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -46,7 +47,7 @@ describe("LLMEvaluator — valid JSON responses", () => { const ai = makeMockAi( JSON.stringify({ score: 0.5, passed: true, reason: "Mediocre", issues: [] }), ); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -57,7 +58,7 @@ describe("LLMEvaluator — valid JSON responses", () => { it("score 0.5 with explicit threshold 0.4 → passed: true", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.5, passed: true, reason: "OK", issues: [] })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", threshold: 0.4 }] }); const result = await evaluator.evaluate("diff content", policy, mockLogger); expect(result.success).toBe(true); @@ -70,7 +71,7 @@ describe("LLMEvaluator — valid JSON responses", () => { describe("LLMEvaluator — unparseable responses fail closed", () => { it("AI returns non-JSON text → score 0, failed, no throw", async () => { const ai = makeMockAi("This diff looks fine overall."); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -85,7 +86,7 @@ describe("LLMEvaluator — unparseable responses fail closed", () => { it('"LGTM" prose is not treated as approval — still fails closed at 0', async () => { const ai = makeMockAi("LGTM, no issues found."); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -97,7 +98,7 @@ describe("LLMEvaluator — unparseable responses fail closed", () => { it("JSON with wrong field types → fails closed with field reason", async () => { const ai = makeMockAi(JSON.stringify({ score: "high", passed: "yes", reason: 42 })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -109,7 +110,7 @@ describe("LLMEvaluator — unparseable responses fail closed", () => { it("fail-closed result still records estimated token costs", async () => { const ai = makeMockAi("not json"); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -125,7 +126,7 @@ describe("LLMEvaluator — error handling", () => { const ai: AiBinding = { run: vi.fn().mockRejectedValue(new Error("network failure")), }; - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(false); if (!result.success) { @@ -137,7 +138,7 @@ describe("LLMEvaluator — error handling", () => { describe("LLMEvaluator — score clamping", () => { it("score above 1 is clamped to 1", async () => { const ai = makeMockAi(JSON.stringify({ score: 1.5, passed: true, reason: "Great" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -147,7 +148,7 @@ describe("LLMEvaluator — score clamping", () => { it("score below 0 is clamped to 0", async () => { const ai = makeMockAi(JSON.stringify({ score: -0.2, passed: false, reason: "Terrible" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -163,7 +164,7 @@ describe("LLMEvaluator — issues array", () => { const ai = makeMockAi( JSON.stringify({ score: 0.4, passed: false, reason: "Problems found", issues }), ); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -173,7 +174,7 @@ describe("LLMEvaluator — issues array", () => { it("issues omitted from result when not present in JSON", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "Clean" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -195,7 +196,7 @@ function sentDiffPortion(ai: AiBinding): string { describe("LLMEvaluator — diff truncation", () => { it("diff longer than the 24000-char default is truncated before being sent to AI", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const longDiff = "a".repeat(30000); await evaluator.evaluate(longDiff, makePolicy(), mockLogger); expect(sentDiffPortion(ai).length).toBe(24000); @@ -203,14 +204,14 @@ describe("LLMEvaluator — diff truncation", () => { it("diff shorter than the window is sent whole", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); await evaluator.evaluate("a".repeat(10000), makePolicy(), mockLogger); expect(sentDiffPortion(ai).length).toBe(10000); }); it("policy maxDiffChars overrides the default window", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", maxDiffChars: 2000 }] }); await evaluator.evaluate("a".repeat(5000), policy, mockLogger); expect(sentDiffPortion(ai).length).toBe(2000); @@ -218,7 +219,7 @@ describe("LLMEvaluator — diff truncation", () => { it("a tiny maxDiffChars is raised to the 1000-char floor, never an empty diff", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", maxDiffChars: 5 }] }); await evaluator.evaluate("a".repeat(5000), policy, mockLogger); expect(sentDiffPortion(ai).length).toBe(1000); @@ -226,7 +227,7 @@ describe("LLMEvaluator — diff truncation", () => { it("a fractional maxDiffChars is floored to an integer window", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", maxDiffChars: 2000.7 }] }); await evaluator.evaluate("a".repeat(5000), policy, mockLogger); expect(sentDiffPortion(ai).length).toBe(2000); @@ -234,7 +235,7 @@ describe("LLMEvaluator — diff truncation", () => { it("a fractional value below one still sends a non-empty diff", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", maxDiffChars: 0.5 }] }); await evaluator.evaluate("a".repeat(5000), policy, mockLogger); expect(sentDiffPortion(ai).length).toBe(1000); @@ -242,7 +243,7 @@ describe("LLMEvaluator — diff truncation", () => { it("policy maxDiffChars is capped at the 100k ceiling", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "llm", maxDiffChars: 5_000_000 }] }); await evaluator.evaluate("a".repeat(200_000), policy, mockLogger); expect(sentDiffPortion(ai).length).toBe(100_000); @@ -250,7 +251,7 @@ describe("LLMEvaluator — diff truncation", () => { it("a truncated evaluation says so in the result issues", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("a".repeat(30000), makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -260,7 +261,7 @@ describe("LLMEvaluator — diff truncation", () => { it("an untruncated evaluation carries no truncation note", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("small diff", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -272,7 +273,7 @@ describe("LLMEvaluator — diff truncation", () => { describe("LLMEvaluator — non-finite scores", () => { it("JSON 1e999 parses to Infinity and fails closed instead of clamping to a pass", async () => { const ai = makeMockAi('{"score": 1e999, "passed": true, "reason": "great"}'); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -288,7 +289,7 @@ describe("LLMEvaluator — model verdict is honored", () => { const ai = makeMockAi( JSON.stringify({ score: 0.9, passed: false, reason: "Looks risky despite score" }), ); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", makePolicy(), mockLogger); expect(result.success).toBe(true); if (result.success) { @@ -301,7 +302,7 @@ describe("LLMEvaluator — model verdict is honored", () => { describe("LLMEvaluator — policy context bound", () => { it("an oversize policy fails closed before any model call", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const policy = makePolicy({ evaluators: [{ type: "webhook", url: `https://ci.example.com/${"a".repeat(9000)}` }], }); @@ -319,7 +320,7 @@ describe("LLMEvaluator — policy context bound", () => { describe("LLMEvaluator — prompt", () => { it("sends a real reviewer system prompt demanding strict JSON", async () => { const ai = makeMockAi(JSON.stringify({ score: 0.9, passed: true, reason: "OK" })); - const evaluator = new LLMEvaluator(ai); + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); await evaluator.evaluate("diff content", makePolicy(), mockLogger); const runMock = ai.run as ReturnType; const calledMessages = ( diff --git a/tests/llm-provider.test.ts b/tests/llm-provider.test.ts new file mode 100644 index 00000000..61b8322f --- /dev/null +++ b/tests/llm-provider.test.ts @@ -0,0 +1,611 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { LLMEvaluator } from "../src/evaluation/llm-evaluator"; +import { + AnthropicProvider, + type LlmProvider, + LlmProviderResponseError, + OpenAiCompatibleProvider, + WorkersAiProvider, +} from "../src/evaluation/llm-provider"; +import type { EvalPolicy } from "../src/evaluation/types"; +import type { AiBinding } from "../src/types"; +import { ExternalServiceError } from "../src/utils/errors"; +import type { Logger } from "../src/utils/logger"; +import { ok } from "../src/utils/result"; + +const mockLogger: Logger = { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + child: vi.fn(() => mockLogger), +}; + +const POLICY: EvalPolicy = { evaluators: [{ type: "llm" }] }; + +/** A JSON body from a provider, as `fetch` would hand it back. */ +function jsonResponse(status: number, body: unknown) { + return { + ok: status >= 200 && status < 300, + status, + json: async () => body, + }; +} + +/** A response whose body is not JSON — `json()` rejects, as the real one does. */ +function nonJsonResponse(status = 200) { + return { + ok: status >= 200 && status < 300, + status, + json: async () => { + throw new SyntaxError("Unexpected token < in JSON at position 0"); + }, + }; +} + +function lastFetchCall() { + const fetchMock = globalThis.fetch as unknown as ReturnType; + const [url, init] = fetchMock.mock.calls[fetchMock.mock.calls.length - 1] as [ + string, + { method: string; headers: Record; body: string }, + ]; + return { url, init, body: JSON.parse(init.body) as Record }; +} + +beforeEach(() => { + vi.restoreAllMocks(); +}); + +describe("WorkersAiProvider", () => { + it("returns the binding's response text and reports no usage", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: "hello" }) }; + const result = await new WorkersAiProvider(ai).run("@cf/model", [ + { role: "user", content: "hi" }, + ]); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe("hello"); + // Absent usage is what keeps the Workers AI path on the char estimate. + expect(result.data.usage).toBeUndefined(); + }); + + it("forwards the model and messages to the binding unchanged", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: "{}" }) }; + const messages = [ + { role: "system", content: "sys" }, + { role: "user", content: "usr" }, + ]; + await new WorkersAiProvider(ai).run("@cf/meta/llama", messages); + expect(ai.run).toHaveBeenCalledWith("@cf/meta/llama", { messages }); + }); + + it("a response object with no `response` field yields empty text, not an error", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue({}) }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe(""); + }); + + it("rejects a ReadableStream response rather than consuming it", async () => { + const ai: AiBinding = { + run: vi.fn().mockResolvedValue(new ReadableStream()), + }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("unexpected stream response"); + }); + + it("a throwing binding becomes a transport error, not a response error", async () => { + const ai: AiBinding = { run: vi.fn().mockRejectedValue(new Error("network failure")) }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("network failure"); + }); + + it("a non-Error rejection still becomes a readable transport error", async () => { + const ai: AiBinding = { run: vi.fn().mockRejectedValue("boom") }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toContain("boom"); + }); +}); + +describe("AnthropicProvider — request shape", () => { + it("POSTs /messages with the api-key, version header and required body fields", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "verdict" }], + usage: { input_tokens: 10, output_tokens: 5 }, + }), + ), + ); + + const provider = new AnthropicProvider("https://api.anthropic.com/v1", "sk-ant-secret"); + await provider.run("claude-sonnet-4-5", [ + { role: "system", content: "be strict" }, + { role: "user", content: "review this" }, + ]); + + const { url, init, body } = lastFetchCall(); + expect(url).toBe("https://api.anthropic.com/v1/messages"); + expect(init.method).toBe("POST"); + expect(init.headers["x-api-key"]).toBe("sk-ant-secret"); + expect(init.headers["anthropic-version"]).toBe("2023-06-01"); + expect(init.headers["content-type"]).toBe("application/json"); + expect(body.model).toBe("claude-sonnet-4-5"); + expect(typeof body.max_tokens).toBe("number"); + // The system prompt is a top-level parameter; a system-role message is rejected. + expect(body.system).toBe("be strict"); + expect(body.messages).toEqual([{ role: "user", content: "review this" }]); + }); + + it("omits `system` entirely when no system message was supplied", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { content: [{ type: "text", text: "x" }] })), + ); + await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", [ + { role: "user", content: "u" }, + ]); + expect(lastFetchCall().body).not.toHaveProperty("system"); + }); + + it("strips a trailing slash from the base URL rather than doubling it", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { content: [{ type: "text", text: "x" }] })), + ); + await new AnthropicProvider("https://api.anthropic.com/v1/", "k").run("m", []); + expect(lastFetchCall().url).toBe("https://api.anthropic.com/v1/messages"); + }); +}); + +describe("AnthropicProvider — response handling", () => { + it("joins every text block and maps usage to real token counts", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [ + { type: "thinking", thinking: "hmm" }, + { type: "text", text: "part one " }, + { type: "text", text: "part two" }, + ], + usage: { input_tokens: 100, output_tokens: 20 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe("part one part two"); + expect(result.data.usage).toEqual({ inputTokens: 100, outputTokens: 20 }); + }); + + it("a content array with no text blocks yields empty text, not an error", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { content: [{ type: "tool_use", id: "t" }] })), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe(""); + }); + + it("non-200 fails with the status only — never the error body", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(401, { + error: { message: "invalid x-api-key: sk-ant-LEAKED" }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("provider returned HTTP 401"); + expect(result.error.message).not.toContain("sk-ant-LEAKED"); + }); + + it("a 200 whose body is not JSON fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(nonJsonResponse())); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("provider response was not valid JSON"); + }); + + it("a 200 of the wrong shape fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(jsonResponse(200, { completion: "text" }))); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response had no content array"); + }); + + it("a 200 whose body is JSON `null` fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(jsonResponse(200, null))); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response had no content array"); + }); + + it("a response with no usage object reports no usage", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { content: [{ type: "text", text: "x" }] })), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a half-filled usage object is treated as absent, never half-trusted", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "x" }], + usage: { input_tokens: 100 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a non-finite token count is treated as absent", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "x" }], + usage: { input_tokens: Number.NaN, output_tokens: 5 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a network failure becomes a transport error, not a response error", async () => { + vi.stubGlobal("fetch", vi.fn().mockRejectedValue(new TypeError("fetch failed"))); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("fetch failed"); + }); + + it("aborts a request that outlives the timeout", async () => { + vi.stubGlobal( + "fetch", + vi.fn( + (_url: string, init: { signal: AbortSignal }) => + new Promise((_resolve, reject) => { + init.signal.addEventListener("abort", () => + reject(new DOMException("The operation was aborted.", "AbortError")), + ); + }), + ), + ); + const provider = new AnthropicProvider("https://api.anthropic.com/v1", "k", { timeoutMs: 5 }); + const result = await provider.run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toContain("aborted"); + }); +}); + +describe("OpenAiCompatibleProvider — request shape", () => { + it("POSTs /chat/completions with a bearer token and the messages verbatim", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + choices: [{ message: { content: "verdict" } }], + usage: { prompt_tokens: 7, completion_tokens: 3 }, + }), + ), + ); + + const messages = [ + { role: "system", content: "be strict" }, + { role: "user", content: "review this" }, + ]; + await new OpenAiCompatibleProvider("https://api.openai.com/v1", "sk-secret").run( + "gpt-4o-mini", + messages, + ); + + const { url, init, body } = lastFetchCall(); + expect(url).toBe("https://api.openai.com/v1/chat/completions"); + expect(init.method).toBe("POST"); + expect(init.headers.authorization).toBe("Bearer sk-secret"); + expect(body.model).toBe("gpt-4o-mini"); + // Unlike Anthropic, the system turn stays in the messages array. + expect(body.messages).toEqual(messages); + }); + + it("works against a self-hosted base URL with a trailing slash", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { choices: [{ message: { content: "x" } }] })), + ); + await new OpenAiCompatibleProvider("http://vllm.internal:8000/v1/", "token-abc").run("m", []); + expect(lastFetchCall().url).toBe("http://vllm.internal:8000/v1/chat/completions"); + }); +}); + +describe("OpenAiCompatibleProvider — response handling", () => { + it("reads the first choice's message content and maps usage", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + choices: [{ message: { content: "the verdict" } }], + usage: { prompt_tokens: 40, completion_tokens: 12, total_tokens: 52 }, + }), + ), + ); + const result = await new OpenAiCompatibleProvider("https://api.groq.com/openai/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe("the verdict"); + expect(result.data.usage).toEqual({ inputTokens: 40, outputTokens: 12 }); + }); + + it("non-200 fails with the status only — never the error body", async () => { + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValue( + jsonResponse(429, { error: { message: "rate limited for key sk-live-LEAKED" } }), + ), + ); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider returned HTTP 429"); + expect(result.error.message).not.toContain("sk-live-LEAKED"); + }); + + it("a body that is not JSON fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(nonJsonResponse())); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response was not valid JSON"); + }); + + it("a 200 with no choices array fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(jsonResponse(200, { object: "error" }))); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response had no choices array"); + }); + + it("an empty choices array fails closed", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(jsonResponse(200, { choices: [] }))); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response had no assistant message text"); + }); + + it("a non-string message content fails closed", async () => { + vi.stubGlobal( + "fetch", + vi + .fn() + .mockResolvedValue( + jsonResponse(200, { choices: [{ message: { content: [{ type: "text" }] } }] }), + ), + ); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider response had no assistant message text"); + }); + + it("a response with no usage reports none, keeping the caller on an estimate", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(jsonResponse(200, { choices: [{ message: { content: "x" } }] })), + ); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a network failure becomes a transport error", async () => { + vi.stubGlobal("fetch", vi.fn().mockRejectedValue(new TypeError("fetch failed"))); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("fetch failed"); + }); + + it("aborts a request that outlives the timeout", async () => { + vi.stubGlobal( + "fetch", + vi.fn( + (_url: string, init: { signal: AbortSignal }) => + new Promise((_resolve, reject) => { + init.signal.addEventListener("abort", () => + reject(new DOMException("The operation was aborted.", "AbortError")), + ); + }), + ), + ); + const provider = new OpenAiCompatibleProvider("https://api.openai.com/v1", "k", { + timeoutMs: 5, + }); + const result = await provider.run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toContain("aborted"); + }); +}); + +const VERDICT = JSON.stringify({ score: 0.9, passed: true, reason: "Looks good" }); + +/** A provider that answers with a fixed verdict and whatever usage it is given. */ +function stubProvider(usage?: { inputTokens: number; outputTokens: number }): LlmProvider { + return { run: vi.fn().mockResolvedValue(ok({ text: VERDICT, usage })) }; +} + +describe("LLMEvaluator — cost accounting is per provider, not per source", () => { + it("a provider reporting usage yields exact counts with no `estimated` flag", async () => { + const evaluator = new LLMEvaluator(stubProvider({ inputTokens: 900, outputTokens: 100 })); + const result = await evaluator.evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs).toEqual([{ kind: "llm_tokens", quantity: 1000 }]); + expect(result.data.costs?.[0]?.estimated).toBeUndefined(); + }); + + it("a provider omitting usage stays on the ~4 chars/token estimate", async () => { + const evaluator = new LLMEvaluator(stubProvider()); + const result = await evaluator.evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs?.[0]?.estimated).toBe(true); + expect(result.data.costs?.[0]?.quantity).toBeGreaterThan(0); + }); + + it("Workers AI, which reports nothing, is on the estimate", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: VERDICT }) }; + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); + const result = await evaluator.evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs?.[0]?.estimated).toBe(true); + }); +}); + +describe("LLMEvaluator — provider failures fail the gate closed", () => { + it("an unusable provider response scores 0 and does not pass", async () => { + const provider: LlmProvider = { + run: vi.fn().mockResolvedValue({ + success: false, + error: new LlmProviderResponseError("provider returned HTTP 503"), + }), + }; + const result = await new LLMEvaluator(provider).evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.score).toBe(0); + expect(result.data.passed).toBe(false); + expect(result.data.reason).toBe("LLM evaluator error: provider returned HTTP 503"); + }); + + it("a transport failure surfaces as an error, still never a pass", async () => { + const provider: LlmProvider = { + run: vi.fn().mockResolvedValue({ + success: false, + error: new ExternalServiceError("LLM", "connection reset"), + }), + }; + const result = await new LLMEvaluator(provider).evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toContain("connection reset"); + }); + + it("a provider error body carrying an API key never reaches the evaluation result", async () => { + // The whole HTTP path, not a stub: a real 401 body with a credential in it. + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(401, { + error: { message: "Incorrect API key provided: sk-live-DEADBEEF12345" }, + }), + ), + ); + const provider = new OpenAiCompatibleProvider( + "https://api.openai.com/v1", + "sk-live-DEADBEEF12345", + ); + const result = await new LLMEvaluator(provider).evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.reason).toBe("LLM evaluator error: provider returned HTTP 401"); + expect(JSON.stringify(result.data)).not.toContain("sk-live-DEADBEEF12345"); + expect(JSON.stringify(result.data)).not.toContain("Incorrect API key"); + }); + + it("a stream response from Workers AI fails closed with the historical reason", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue(new ReadableStream()) }; + const result = await new LLMEvaluator(new WorkersAiProvider(ai)).evaluate( + "diff content", + POLICY, + mockLogger, + ); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.score).toBe(0); + expect(result.data.passed).toBe(false); + expect(result.data.reason).toBe("LLM evaluator error: unexpected stream response"); + // Byte-identical to the pre-seam behaviour: no costs on this path. + expect(result.data.costs).toBeUndefined(); + }); + + it("a provider that throws is caught rather than escaping the evaluator", async () => { + const provider: LlmProvider = { + run: vi.fn().mockRejectedValue(new Error("isolate exceeded CPU")), + }; + const result = await new LLMEvaluator(provider).evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toContain("isolate exceeded CPU"); + }); +}); diff --git a/tests/magic-link-rate-limit.test.ts b/tests/magic-link-rate-limit.test.ts index f6b71081..69fbd172 100644 --- a/tests/magic-link-rate-limit.test.ts +++ b/tests/magic-link-rate-limit.test.ts @@ -109,6 +109,10 @@ describe("MagicLinkRateLimiter", () => { const windowEndMs = (Math.floor(now / 1000 / WINDOW) + 1) * HOUR_MS; expect(alarm).toBeGreaterThan(windowEndMs); + // workerd consumes an alarm by firing it, so the fake clears it here rather + // than inside `deleteAll` — which does NOT clear a pending alarm in + // production, and modelling it as if it did hid that from every test. + await storage.deleteAlarm(); await instance.alarm(); expect(await storage.getAlarm()).toBeNull(); // Storage erased, so the next window starts from nothing. diff --git a/tests/usage-meter.test.ts b/tests/usage-meter.test.ts index 6f13c624..4cb25a51 100644 --- a/tests/usage-meter.test.ts +++ b/tests/usage-meter.test.ts @@ -403,6 +403,59 @@ describe("UsageMeter evaluations_per_hour", () => { expect((await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW)).admitted).toBe(true); }); + it("returns a rate reservation settled in a LATER bucket than it was taken in", async () => { + // The refund used to be added to the bucket `nowMs` falls in, which is + // empty once the settle outlives the reservation's bucket — the normal + // case, since buckets are minutes and an evaluation is seconds to minutes. + // Math.max(0, …) then clamped the refund away and the charge sat in its own + // bucket for the rest of the hour, so a provider call that failed three + // minutes in still burned its slot until the window rolled. + const { stub } = meter(); + const sixMinutes = 6 * 60 * 1000; + + for (let i = 0; i < 3; i++) { + await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW); + } + expect((await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW)).admitted).toBe(false); + + await stub.settle("evaluations_per_hour", -1, PERIOD, NOW + sixMinutes); + + const after = await stub.reserve("evaluations_per_hour", 1, 3, PERIOD, NOW + sixMinutes); + expect(after.admitted).toBe(true); + }); + + it("refuses to refund more than the window is holding", async () => { + const { stub } = meter(); + await stub.reserve("evaluations_per_hour", 1, 5, PERIOD, NOW); + + await stub.settle("evaluations_per_hour", -4, PERIOD, NOW); + + // Back to empty, not negative — a refund must not invent headroom. + const outcome = await stub.reserve("evaluations_per_hour", 5, 5, PERIOD, NOW); + expect(outcome.admitted).toBe(true); + expect(outcome.count).toBe(5); + }); + + it.each([ + ["microseconds mistaken for milliseconds", NOW * 1000], + ["seconds mistaken for milliseconds", Math.floor(NOW / 1000)], + ])("admits and counts nothing for a clock in %s", async (_label, badClock) => { + // Finite, so the old `Number.isFinite` guard let it through: the bucket + // landed ~1000x into the future, counted against every later window, and + // refused the subject permanently while arming a cleanup alarm past the + // year 58000. Admitting and counting nothing is the same treatment every + // other malformed input gets. + const { stub } = meter(); + + const bad = await stub.reserve("evaluations_per_hour", 1, 1, PERIOD, badClock); + expect(bad).toEqual({ admitted: true, count: 0 }); + + // And the good clock that follows is unaffected — the whole allowance is + // still there, which is what "counts nothing" has to mean. + const good = await stub.reserve("evaluations_per_hour", 1, 1, PERIOD, NOW); + expect(good.admitted).toBe(true); + }); + it("admits when the clock is not a usable timestamp", async () => { const { stub } = meter(); diff --git a/tests/wrangler-usage-meter-config.test.ts b/tests/wrangler-usage-meter-config.test.ts index e7c9ffa4..4a086a15 100644 --- a/tests/wrangler-usage-meter-config.test.ts +++ b/tests/wrangler-usage-meter-config.test.ts @@ -91,7 +91,10 @@ describe("wrangler.toml UsageMeter configuration", () => { expect(settingValue(body as string, "ENTITLEMENTS_ENFORCE")).toBe("0"); // The other three ship commented (a URL, a secret's name, an operator // allowlist), so what is asserted is that each block documents them — - // the same shape-based assertion the telemetry guard makes. + // the same shape-based assertion the telemetry guard makes. This checks + // presence only: it deliberately does NOT assert what the prose says, + // because LLM_PROVIDERS is reserved and unread until Task 7, and a test + // that pinned its description would lock a false claim into place. for (const key of ["BILLING_SERVICE_URL", "BILLING_SERVICE_SECRET", "LLM_PROVIDERS"]) { expect(body as string, `${key} undocumented in [${scope}]`).toContain(key); } diff --git a/wrangler.toml b/wrangler.toml index 9ba23209..21c2b663 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -231,14 +231,17 @@ OAUTH_REDIRECT_URI = "http://localhost:8787/auth/github/callback" # `entitlementsConfigError` (src/middleware/config-guard.ts) logs exactly that. ENTITLEMENTS_ENFORCE = "0" # -# The evaluator's provider allowlist, and the ONLY place a provider endpoint may -# be named: a project's .stratum/policy.yaml may SELECT one of these by name and -# can never supply a `baseUrl`, so a project owner cannot aim the operator's -# Worker at a host of their choosing. An unrecognized name fails the policy -# closed — a merge gate that cannot run must not pass. Unset (the default) means -# Workers AI (`env.AI`) only, which is today's behavior. Credentials are NOT -# here: they live per project in `project_secrets`. Shape illustrated, parsed by -# the provider seam: +# NOT YET IMPLEMENTED — reserved, and documented here so the shape is settled +# before anything reads it. Setting it today does nothing: no code parses it. +# +# When Task 7 lands it becomes the evaluator's provider allowlist, and the ONLY +# place a provider endpoint may be named: a project's .stratum/policy.yaml will +# SELECT one of these by name and will never be able to supply a `baseUrl`, so a +# project owner cannot aim the operator's Worker at a host of their choosing. An +# unrecognized name will fail the policy closed — a merge gate that cannot run +# must not pass. Unset (the default) means Workers AI (`env.AI`) only, which is +# both today's behavior and the behavior until that work ships. Credentials are +# NOT here: they live per project in `project_secrets`. Intended shape: # LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' # Instance-wide telemetry kill switch. Not a secret — a plain var, listed here From aa7b2ccf878237e3c942ebf53357f85b1bfdcef7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 11:16:38 +0000 Subject: [PATCH 12/26] Let a project bring its own key, without letting it choose the endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 7, plus the review findings on the provider seam. The threat this is built around: .stratum/policy.yaml is repository content. Policy loads from the project's default branch, so a proposed edit never governs its own evaluation — it lands at merge, behind approvals and touches_protected_config. The reachable attacker is therefore the project owner using the operator's Worker as an exfiltration proxy, not a stranger opening a change. So a policy may SELECT a provider by name from an operator-configured allowlist and can never supply a baseUrl, and an unrecognized name fails the policy closed — the evaluator-side convention, deliberately opposite to the deploy sanitizer two files away, because a gate that cannot run must not pass. The allowlist is validated where it is parsed, not per request: https only, and loopback, unspecified, link-local, unique-local, IPv4-mapped and every RFC1918 range refused. Link-local coverage is what keeps the cloud metadata endpoint out. Credential failures fail closed with three distinct reasons — key unset, secret missing, row undecryptable — and never fall back to env.AI. Falling back would put the spend back on the operator while leaving the gate looking like it ran, which is the failure this whole PRD calls its worst. The derived key is cached per isolate and the secret is not loaded at all unless a policy names a provider, because PBKDF2 at 100k iterations now sits on a path that already does two clones and a diff. Stripping `provider` from the prompt turned out to close the policy-size hazard structurally rather than by sizing: the serialized policy does not grow when a project opts in, so adding a provider block cannot push an existing project past the context limit. From the seam's review: usage counts were accepted without checking sign or integrality, so a provider returning a negative token count wrote a negative quantity to the ledger and marked it exact. Guarded at the seam and again in recordCosts. A dropped connection mid-body was reported as malformed JSON, blaming the provider for a network fault; the body is now read and parsed in separate steps so the two stay distinguishable, and the SyntaxError is kept as a cause rather than discarded. And the seam claimed in three places that Workers AI reports no token usage. It does — Cloudflare's generated types name prompt_tokens and completion_tokens — this repo's AiBinding just never declared them. So every production cost sample was estimated while exact counts sat one type declaration away, in a feature about metering accuracy. The binding now surfaces them. Typecheck clean; lint clean; 3994 tests pass, up 74. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- src/deploy/runner.ts | 18 +- src/evaluation/index.ts | 2 + src/evaluation/llm-byok.ts | 142 +++++ src/evaluation/llm-evaluator.ts | 40 +- src/evaluation/llm-provider.ts | 201 +++++-- src/evaluation/llm-providers.ts | 281 +++++++++ src/evaluation/policy-loader.ts | 134 ++++- src/evaluation/sanitize-policy.ts | 73 ++- src/evaluation/types.ts | 29 +- src/middleware/config-guard.ts | 39 +- src/routes/changes.ts | 12 +- src/routes/sync-management.ts | 11 +- src/services/change-flow.ts | 43 +- src/storage/costs.ts | 13 +- src/storage/project-secrets.ts | 19 +- src/types.ts | 27 +- tests/cost-attribution.test.ts | 21 + tests/evaluator-billing-context.test.ts | 31 +- tests/llm-byok.test.ts | 730 ++++++++++++++++++++++++ tests/llm-provider.test.ts | 231 +++++++- tests/sandbox-evaluator.test.ts | 10 +- wrangler.toml | 41 +- 22 files changed, 2017 insertions(+), 131 deletions(-) create mode 100644 src/evaluation/llm-byok.ts create mode 100644 src/evaluation/llm-providers.ts create mode 100644 tests/llm-byok.test.ts diff --git a/src/deploy/runner.ts b/src/deploy/runner.ts index 5f0b107b..cc0ecbaa 100644 --- a/src/deploy/runner.ts +++ b/src/deploy/runner.ts @@ -29,6 +29,7 @@ * injects its own seams: without them, testing any of the above needs a real git * remote, a real clock, and a real provider account. */ +import { type LlmProviderCatalog, llmProviderCatalog } from "../evaluation/llm-providers"; import { parsePolicyContent } from "../evaluation/policy-loader"; import type { DeployConfig, DeployRejection } from "../evaluation/types"; import { type StratumEvent, emitEvent } from "../queue/events"; @@ -353,7 +354,7 @@ async function runMergeMessage( } const files = treeResult.data; - const policy = policyFromTree(files, logger); + const policy = policyFromTree(files, logger, llmProviderCatalog(env, logger)); // A rejected entry becomes a visible failed row, never a dropped one: a deploy // the author wrote and that never runs means production silently stopped @@ -488,7 +489,7 @@ async function runDeploymentMessage( // Config comes from the pinned commit here too, so an approval granted days // later still deploys with the configuration the commit was reviewed under. - const policy = policyFromTree(treeResult.data, logger); + const policy = policyFromTree(treeResult.data, logger, llmProviderCatalog(env, logger)); const config = policy.deploys.find((entry) => entry.name === deployment.name); if (!config) { const rejection = policy.rejections.find((entry) => entry.name === deployment.name); @@ -879,14 +880,23 @@ async function readTree( * effect. A malformed file yields one named rejection rather than silence, * because a single YAML typo must not quietly stop production updating. */ -function policyFromTree(files: ReadonlyMap, logger: Logger): TreePolicy { +function policyFromTree( + files: ReadonlyMap, + logger: Logger, + /** The same allowlist the evaluation path parses with. Passed even though + * only `deploys:` is read here: a policy naming a provider the parser cannot + * resolve is *malformed*, and reading it without the catalog would report + * "no deploy configuration could be read" for a file that is perfectly + * valid. */ + providers: LlmProviderCatalog, +): TreePolicy { const decoder = new TextDecoder(); for (const { path, format } of POLICY_FILES) { const bytes = files.get(path); if (bytes === undefined) continue; - const parsed = parsePolicyContent(decoder.decode(bytes), format, logger); + const parsed = parsePolicyContent(decoder.decode(bytes), format, logger, providers); if (parsed.status === "malformed") { return { deploys: [], diff --git a/src/evaluation/index.ts b/src/evaluation/index.ts index c2e2b2f7..de25ec90 100644 --- a/src/evaluation/index.ts +++ b/src/evaluation/index.ts @@ -5,5 +5,7 @@ export * from "./composite-evaluator"; export * from "./policy-loader"; export * from "./sandbox-evaluator"; export * from "./llm-provider"; +export * from "./llm-providers"; +export * from "./llm-byok"; export * from "./llm-evaluator"; export * from "./secret-scanner"; diff --git a/src/evaluation/llm-byok.ts b/src/evaluation/llm-byok.ts new file mode 100644 index 00000000..21dd3203 --- /dev/null +++ b/src/evaluation/llm-byok.ts @@ -0,0 +1,142 @@ +/** + * Resolving the LLM provider an evaluation runs on: the operator's Workers AI + * binding, or the project's own credential against an operator-configured + * endpoint (BYOK). + * + * Two rules shape everything here, and both are PRD §6/§7: + * + * 1. **Never fall back to `env.AI`.** A BYOK misconfiguration that quietly ran + * on Workers AI would put the spend back on the operator's bill — the exact + * hole BYOK closes — and would be a gate that silently stops gating. Every + * failure below is fail-closed with a distinct reason instead. + * 2. **Skip the credential load entirely when no BYOK provider is named.** + * That is the overwhelmingly common case, and the load is a D1 read plus a + * 100k-iteration PBKDF2 derivation on the change-creation path, which + * already does two clones and a diff. + */ +import { DEPLOY_SECRET_KEY_MISSING, loadSecretValues } from "../storage/project-secrets"; +import type { Env, ProjectEntry } from "../types"; +import type { Logger } from "../utils/logger"; +import { AnthropicProvider, type LlmProvider, OpenAiCompatibleProvider } from "./llm-provider"; +import { llmProviderCatalog, providerSecretName } from "./llm-providers"; +import type { EvalPolicy, LlmEvaluatorConfig } from "./types"; + +/** Env slice the resolution needs. Named rather than `Env` so the credential + * path's dependencies are visible at a glance. */ +export type LlmByokEnv = Pick; + +export type LlmProviderSelection = + /** No BYOK provider named — the caller uses Workers AI, as it always has. */ + | { status: "platform" } + /** Run on the project's credential; cost samples from this run are `byok`. */ + | { status: "byok"; provider: LlmProvider; providerName: string } + /** The named provider cannot be used. The gate fails closed on this reason, + * and never on `env.AI`. */ + | { status: "unavailable"; reason: string }; + +/** + * The `llm` entry the evaluator will actually read. + * + * `find`, not `filter`: `LLMEvaluator` itself takes the first `llm` entry and + * ignores any others, so resolving from a later one would run the credential of + * a provider whose model and threshold are not the ones in force. + */ +export function llmConfigOf(policy: EvalPolicy): LlmEvaluatorConfig | undefined { + return policy.evaluators.find((entry): entry is LlmEvaluatorConfig => entry.type === "llm"); +} + +/** + * Resolve the provider for this policy, loading the project's credential only + * when the policy names one. + * + * The plaintext credential reaches exactly one place: an `Authorization` (or + * `x-api-key`) header on the provider request. `loadSecretValues` is the only + * read path for a secret value in the codebase and is held to the same rule as + * the deploy runner — no route returns it, renders it, or logs it. That an + * HTTP request *triggers* this resolution is not the same thing as the value + * being route-reachable. + * + * **There is deliberately no MCP surface for writing the credential this + * reads** (PRD §4c). `stratum_set_provider_key` looks like an obvious gap and + * is ruled out on purpose: a provider key must never pass through a model's + * context window, and §7's threat model rests on credentials and endpoints not + * being selectable by the wrong party. The write surface is the web UI or the + * CLI — the same rule that keeps agent tokens from submitting review verdicts. + */ +export async function resolveLlmProvider( + env: LlmByokEnv, + project: ProjectEntry, + policy: EvalPolicy, + logger: Logger, +): Promise { + const providerName = llmConfigOf(policy)?.provider; + // The skip. Nothing below this line runs for a project that never opted in. + if (providerName === undefined) return { status: "platform" }; + + const configured = llmProviderCatalog(env, logger).get(providerName); + if (!configured) { + // Unreachable through `parsePolicyContent`, which rejects an unknown name + // and fails the whole file closed. Kept because this function must be safe + // to call with a policy from anywhere — a KV-cached one written before the + // operator's allowlist changed, say — and the fallback must be a refusal. + return { + status: "unavailable", + reason: `policy names LLM provider "${providerName}", which is not configured in LLM_PROVIDERS`, + }; + } + + const secretName = providerSecretName(providerName); + const loaded = await loadSecretValues(env.DB, logger, env, { + projectId: project.id, + names: [secretName], + }); + + if (!loaded.success) { + // `DEPLOY_SECRET_KEY` unset is an operator misconfiguration with a specific + // remedy, and it is the coupling PRD Open Question 4 accepted: BYOK reuses + // the deploy secret key, so rotating that key blocks BYOK merges until the + // provider keys are re-entered. Named explicitly so the reader is not left + // guessing which of the two subsystems broke. + if (loaded.error.code === DEPLOY_SECRET_KEY_MISSING) { + return { + status: "unavailable", + reason: `project secret ${secretName} cannot be decrypted because the operator has not configured DEPLOY_SECRET_KEY`, + }; + } + return { + status: "unavailable", + reason: `project secret ${secretName} could not be read (${loaded.error.code})`, + }; + } + + if (loaded.data.missing.includes(secretName)) { + return { + status: "unavailable", + reason: `project secret ${secretName} is not set for this project; add it before selecting the "${providerName}" provider`, + }; + } + if (loaded.data.undecryptable.includes(secretName)) { + return { + status: "unavailable", + reason: `project secret ${secretName} could not be decrypted; DEPLOY_SECRET_KEY was rotated or the value was restored from another project, so it must be re-entered`, + }; + } + + const apiKey = loaded.data.values.get(secretName); + if (apiKey === undefined) { + // Belt and braces: `loadSecretValues` reports every unresolved name in one + // of the two lists above, so this is not reachable — but the alternative to + // checking is constructing a provider with `undefined` for a credential. + return { + status: "unavailable", + reason: `project secret ${secretName} did not resolve to a value`, + }; + } + + const provider = + configured.kind === "anthropic" + ? new AnthropicProvider(configured.baseUrl, apiKey) + : new OpenAiCompatibleProvider(configured.baseUrl, apiKey); + + return { status: "byok", provider, providerName }; +} diff --git a/src/evaluation/llm-evaluator.ts b/src/evaluation/llm-evaluator.ts index 42ef8068..cb9869c7 100644 --- a/src/evaluation/llm-evaluator.ts +++ b/src/evaluation/llm-evaluator.ts @@ -6,7 +6,7 @@ import { err, ok } from "../utils/result"; import type { LlmProvider } from "./llm-provider"; import { LlmProviderResponseError } from "./llm-provider"; import { sanitizePolicy } from "./sanitize-policy"; -import type { EvalPolicy, EvalResult, Evaluator } from "./types"; +import type { EvalPolicy, EvalResult, Evaluator, LlmEvaluatorConfig } from "./types"; const DEFAULT_MODEL = "@cf/meta/llama-3.1-8b-instruct"; const DEFAULT_THRESHOLD = 0.7; @@ -31,7 +31,20 @@ const SYSTEM_PROMPT = [ ].join(" "); export class LLMEvaluator implements Evaluator { - constructor(private provider: LlmProvider) {} + /** + * @param provider Who runs the model. Workers AI by default; an operator- + * configured HTTP provider on the project's own credential under BYOK. + * @param costSource Who paid for the tokens this run reports. `"byok"` marks + * every cost sample so a project running on its own key does not decrement + * the hosted allowance — the whole point of the field, and something no + * later layer can reconstruct, since the recorded sample is all it sees. + * Defaults to `"platform"`, which the sample omits entirely (absent means + * the operator paid). + */ + constructor( + private provider: LlmProvider, + private costSource: "platform" | "byok" = "platform", + ) {} async evaluate( diff: string, @@ -41,9 +54,7 @@ export class LLMEvaluator implements Evaluator { logger.debug("Starting LLM evaluation"); try { - const config = policy.evaluators.find((e) => e.type === "llm") as - | { type: "llm"; model?: string; threshold?: number; maxDiffChars?: number } - | undefined; + const config = policy.evaluators.find((e): e is LlmEvaluatorConfig => e.type === "llm"); const model = config?.model ?? DEFAULT_MODEL; const threshold = config?.threshold ?? DEFAULT_THRESHOLD; const maxDiffChars = @@ -87,8 +98,10 @@ export class LLMEvaluator implements Evaluator { if (!runResult.success) { // A response that arrived but cannot be used fails the gate closed with // a readable reason; a transport failure stays an error the caller - // surfaces. Either way the provider's own message is metadata only — - // it is copied into a user-visible reason by `runEvaluation`. + // surfaces. Either way the provider's own message is metadata only and + // becomes user-visible: this block interpolates it into the reason + // below, and `runEvaluation` copies the message of the error the other + // branch returns into a reason of its own. if (runResult.error instanceof LlmProviderResponseError) { logger.error(`LLM evaluation failed: ${runResult.error.message}`); return ok({ @@ -103,16 +116,23 @@ export class LLMEvaluator implements Evaluator { const responseText = runResult.data.text; const usage = runResult.data.usage; + // Written only for `byok`: an absent `source` already means "platform", + // so always emitting it would add a field that says nothing to every + // sample the operator pays for. + const source = this.costSource === "byok" ? { source: "byok" as const } : {}; // Real counts when the provider reports them, otherwise ~4 chars/token — - // the distinction is per provider, not per source. Workers AI reports no - // usage at all, so that path stays estimated. + // the distinction is per response, not per source. Every provider behind + // this seam reports counts when it has them, Workers AI included; a + // response that omits them, or reports ones the provider would not trust, + // is what keeps a run on the estimate. const costs: EvalResult["costs"] = usage - ? [{ kind: "llm_tokens", quantity: usage.inputTokens + usage.outputTokens }] + ? [{ kind: "llm_tokens", quantity: usage.inputTokens + usage.outputTokens, ...source }] : [ { kind: "llm_tokens", quantity: Math.ceil((promptChars + responseText.length) / 4), estimated: true, + ...source, }, ]; diff --git a/src/evaluation/llm-provider.ts b/src/evaluation/llm-provider.ts index c26c049d..2ebd8468 100644 --- a/src/evaluation/llm-provider.ts +++ b/src/evaluation/llm-provider.ts @@ -26,9 +26,10 @@ export interface LlmResponse { /** * Present only when the provider reported real token counts. Its absence is * load-bearing: the evaluator falls back to the `~4 chars/token` estimate and - * marks the cost sample `estimated`. That distinction is **per provider, not - * per source** — Workers AI reports nothing, and an OpenAI-compatible - * endpoint that omits `usage` is estimated too. + * marks the cost sample `estimated`. That distinction is **per response, not + * per source**: every provider here reports counts when it has them (Workers + * AI included), and any response that omits `usage` — or reports one this + * file will not trust, see `usageOrNothing` — is estimated. */ usage?: LlmUsage; } @@ -48,19 +49,26 @@ export interface LlmProvider { /** * The provider answered, but the answer is unusable — a non-2xx status, a body - * that is not JSON, a JSON body of the wrong shape, or a stream where a - * completion was expected. + * that is not JSON, a JSON body of the wrong shape, a verdict truncated at the + * token cap, or a stream where a completion was expected. * - * **The message is metadata only.** It reaches the user: `runEvaluation` - * (`services/change-flow.ts`) copies `error.message` straight into the recorded - * `EvalResult.reason`. A provider's error body can quote the request — which - * contains the diff — and some providers echo the credential that failed, so - * the body is never interpolated into it. Status codes and byte counts only. + * **The message is metadata only.** It reaches the user: `LLMEvaluator` + * intercepts every error of this class and interpolates the message into the + * `EvalResult.reason` it fails the gate closed with. A provider's error body + * can quote the request — which contains the diff — and some providers echo the + * credential that failed, so nothing the provider sent is ever interpolated + * into it: these messages carry an HTTP status, a fixed description of the + * shape problem, or a constant this file chose, and nothing else. */ export class LlmProviderResponseError extends AppError { - constructor(message: string) { + constructor(message: string, cause?: unknown) { super(message, "LLM_PROVIDER_RESPONSE", 502); this.name = "LlmProviderResponseError"; + // Attached rather than described: a `SyntaxError` from `JSON.parse` quotes + // the bytes that failed, which are the provider's body — exactly what the + // message must not carry. Dropping it entirely would swallow the only + // record of what went wrong. + if (cause !== undefined) this.cause = cause; } } @@ -73,23 +81,65 @@ function transportError(error: unknown): AppError { ) as AppError; } +/** How long any provider gets before the gate gives up on it. */ +const DEFAULT_TIMEOUT_MS = 30_000; + +export interface ProviderOptions { + /** Give up on the request after this many milliseconds. */ + timeoutMs?: number; +} + +/** + * Bound a call that cannot be cancelled. + * + * `AiBinding.run` takes no `AbortSignal`, so unlike the HTTP providers there is + * nothing here to abort: the inference keeps running (and is still billed) + * after this rejects. What it bounds is how long the *gate* waits, which is the + * failure that matters — a binding that never settles otherwise holds the + * evaluation open for as long as the Worker's own limits allow. + */ +function withTimeout(work: Promise, timeoutMs: number, what: string): Promise { + let timer: ReturnType | undefined; + const expiry = new Promise((_resolve, reject) => { + timer = setTimeout( + () => reject(new Error(`${what} did not respond within ${timeoutMs}ms`)), + timeoutMs, + ); + }); + return Promise.race([work, expiry]).finally(() => { + if (timer !== undefined) clearTimeout(timer); + }); +} + /** * The Workers AI binding, and the default. Behaviour here is deliberately * identical to what `LLMEvaluator` did inline before the seam existed: * a `ReadableStream` is rejected rather than consumed, and a response object - * with no `response` field yields `""` so the evaluator fails it closed as - * unparseable output rather than treating it as a distinct condition. + * with no usable `response` field yields `""` so the evaluator fails it closed + * as unparseable output rather than treating it as a distinct condition. * - * It reports no token usage, which is why the Workers AI path keeps the - * `~4 chars/token` estimate and `estimated: true`. + * It does report token counts — Cloudflare documents the synchronous output of + * a text-generation model as `response`, `tool_calls` and `usage`, and its own + * generated types name the counts `usage.prompt_tokens` / + * `usage.completion_tokens` (`AiTextGenerationOutput` in + * `@cloudflare/workers-types`). They go through the same `usageOrNothing` + * validation as the HTTP providers, so this path falls back to the + * `~4 chars/token` estimate only when a response actually omits them. */ export class WorkersAiProvider implements LlmProvider { - constructor(private ai: AiBinding) {} + constructor( + private ai: AiBinding, + private options: ProviderOptions = {}, + ) {} async run(model: string, messages: Message[]): Promise> { let raw: Awaited>; try { - raw = await this.ai.run(model, { messages }); + raw = await withTimeout( + this.ai.run(model, { messages }), + this.options.timeoutMs ?? DEFAULT_TIMEOUT_MS, + "Workers AI", + ); } catch (error) { return err(transportError(error)); } @@ -98,18 +148,16 @@ export class WorkersAiProvider implements LlmProvider { return err(new LlmProviderResponseError("unexpected stream response")); } - return ok({ text: raw.response ?? "" }); + // Checked, not asserted: the binding's own types make `response` optional, + // and the evaluator dereferences `.length` on whatever comes back. + const text = typeof raw.response === "string" ? raw.response : ""; + return ok({ + text, + ...usageOrNothing(raw.usage?.prompt_tokens, raw.usage?.completion_tokens), + }); } } -/** How long a hosted provider gets before the gate gives up on it. */ -const DEFAULT_TIMEOUT_MS = 30_000; - -export interface HttpProviderOptions { - /** Abort the request after this many milliseconds. */ - timeoutMs?: number; -} - /** * Credentials and endpoints for both HTTP providers are **constructor * arguments supplied by the caller** — never read from `env` here, and never @@ -125,7 +173,7 @@ abstract class HttpLlmProvider implements LlmProvider { constructor( protected readonly baseUrl: string, protected readonly apiKey: string, - protected readonly options: HttpProviderOptions = {}, + protected readonly options: ProviderOptions = {}, ) {} abstract run(model: string, messages: Message[]): Promise>; @@ -158,11 +206,17 @@ abstract class HttpLlmProvider implements LlmProvider { ); } + // Read the bytes first, parse them second. `response.json()` rejects both + // when the transfer drops mid-body and when a body that arrived whole is + // not JSON, and those mean opposite things to the gate: the first is a + // transport failure the caller surfaces (it falls to the outer catch), + // the second fails the evaluation closed with a user-visible reason. + const raw = await response.text(); try { - return ok((await response.json()) as unknown); - } catch { + return ok(JSON.parse(raw) as unknown); + } catch (error) { return err( - new LlmProviderResponseError("provider response was not valid JSON") as AppError, + new LlmProviderResponseError("provider response was not valid JSON", error) as AppError, ); } } catch (error) { @@ -173,6 +227,14 @@ abstract class HttpLlmProvider implements LlmProvider { } } +/** + * The output budget for one verdict, shared by both HTTP providers so the gate + * is bounded by the same number whichever one runs it. The evaluator asks for a + * single small JSON object; a response that needs more than this has stopped + * answering the question it was asked. + */ +const VERDICT_TOKEN_BUDGET = 1024; + /** * Anthropic's Messages API: `POST {baseUrl}/messages`, `x-api-key` plus the * required `anthropic-version` header, assistant text in `content[]` blocks of @@ -186,7 +248,7 @@ export class AnthropicProvider extends HttpLlmProvider { * a value the caller could vary is a shape this code has not been written for. */ static readonly API_VERSION = "2023-06-01"; /** The evaluator asks for one small JSON verdict; `max_tokens` is required. */ - static readonly MAX_TOKENS = 1024; + static readonly MAX_TOKENS = VERDICT_TOKEN_BUDGET; async run(model: string, messages: Message[]): Promise> { // Anthropic takes the system prompt as a top-level parameter, not as a @@ -214,6 +276,7 @@ export class AnthropicProvider extends HttpLlmProvider { const body = posted.data as { content?: unknown; + stop_reason?: unknown; usage?: { input_tokens?: unknown; output_tokens?: unknown }; } | null; @@ -223,6 +286,18 @@ export class AnthropicProvider extends HttpLlmProvider { ); } + // A verdict cut off at our own cap is well-formed JSON containing half a + // JSON object, so without this the evaluator fails closed blaming the + // model's formatting for a limit this file chose. Anthropic says so + // explicitly: `stop_reason` is `"max_tokens"` in exactly that case. + if (body.stop_reason === "max_tokens") { + return err( + new LlmProviderResponseError( + `provider response was truncated at the ${AnthropicProvider.MAX_TOKENS}-token limit`, + ) as AppError, + ); + } + const text = body.content .filter( (block): block is { type: "text"; text: string } => @@ -250,14 +325,24 @@ export class AnthropicProvider extends HttpLlmProvider { * https://docs.vllm.ai/en/latest/serving/openai_compatible_server/. */ export class OpenAiCompatibleProvider extends HttpLlmProvider { + /** The same budget the Anthropic sibling sends; see `VERDICT_TOKEN_BUDGET`. */ + static readonly MAX_COMPLETION_TOKENS = VERDICT_TOKEN_BUDGET; + async run(model: string, messages: Message[]): Promise> { - // No `max_tokens`: the four vendors this one endpoint serves disagree about - // it (newer OpenAI models take `max_completion_tokens` and reject - // `max_tokens`), and an unrecognised field is a 400 from the strict ones. + // `max_completion_tokens`, never `max_tokens`: OpenAI documents `max_tokens` + // as deprecated in its favour and "not compatible with o-series models", + // Groq documents the same deprecation, and OpenRouter and vLLM take either. + // Together AI documents only `max_tokens` and its schema does not forbid + // extra fields, so there the bound is still the timeout — an unenforced + // field on one vendor beats a 400 on the ones that reject `max_tokens`. const posted = await this.post( "/chat/completions", { authorization: `Bearer ${this.apiKey}` }, - { model, messages }, + { + model, + messages, + max_completion_tokens: OpenAiCompatibleProvider.MAX_COMPLETION_TOKENS, + }, ); if (!posted.success) return posted; @@ -272,8 +357,23 @@ export class OpenAiCompatibleProvider extends HttpLlmProvider { ); } - const content = (body.choices[0] as { message?: { content?: unknown } } | undefined)?.message - ?.content; + const choice = body.choices[0] as + | { message?: { content?: unknown }; finish_reason?: unknown } + | undefined; + + // The same truncation the Anthropic path reports, under this shape's name + // for it: `finish_reason: "length"` means the cap above ended the verdict, + // not the model, and the half-object left behind must not be reported as + // the model's own formatting failure. + if (choice?.finish_reason === "length") { + return err( + new LlmProviderResponseError( + `provider response was truncated at the ${OpenAiCompatibleProvider.MAX_COMPLETION_TOKENS}-token limit`, + ) as AppError, + ); + } + + const content = choice?.message?.content; if (typeof content !== "string") { return err( new LlmProviderResponseError("provider response had no assistant message text") as AppError, @@ -288,18 +388,25 @@ export class OpenAiCompatibleProvider extends HttpLlmProvider { } /** - * Report usage only when both counts are real numbers. A partial or malformed + * A count that could actually be a number of tokens: a non-negative integer. + * + * `Number.isInteger` already excludes `NaN` and both infinities, so the whole + * test is integrality plus sign. Both halves matter downstream — a negative + * quantity becomes a ledger row for negative spend, and a fractional one a + * fractional token — and neither is a count any provider can legitimately + * report, so a value failing this is a malformed `usage`, not a small error. + */ +function isTokenCount(value: unknown): value is number { + return typeof value === "number" && Number.isInteger(value) && value >= 0; +} + +/** + * Report usage only when both counts are usable. A partial or malformed * `usage` object is treated as absent rather than half-trusted: a wrong token - * count recorded as exact is worse than an estimate labelled as one. + * count recorded as exact is worse than an estimate labelled as one — the + * estimate at least carries `estimated: true` into `cost_records`. */ function usageOrNothing(input: unknown, output: unknown): { usage?: LlmUsage } { - if ( - typeof input !== "number" || - typeof output !== "number" || - !Number.isFinite(input) || - !Number.isFinite(output) - ) { - return {}; - } + if (!isTokenCount(input) || !isTokenCount(output)) return {}; return { usage: { inputTokens: input, outputTokens: output } }; } diff --git a/src/evaluation/llm-providers.ts b/src/evaluation/llm-providers.ts new file mode 100644 index 00000000..136a5541 --- /dev/null +++ b/src/evaluation/llm-providers.ts @@ -0,0 +1,281 @@ +/** + * The operator's LLM provider allowlist, parsed from the `LLM_PROVIDERS` var. + * + * This is the **only** place a provider endpoint may be named. A project's + * `.stratum/policy.yaml` may SELECT one of these by name and can never supply a + * `baseUrl` (PRD §7): a base URL a repository's own policy chose would turn the + * operator's Worker into a request-forgery primitive aimed at whatever host the + * policy names, and the prompt it would receive contains the diff. + * + * Every URL rule below is checked once, here, at parse time — not per request. + * Per-request checking would re-validate an immutable value on the hot path and, + * worse, would mean a bad entry is discovered by the first project unlucky + * enough to select it rather than by the operator who wrote it. + */ +import type { Env } from "../types"; +import type { Logger } from "../utils/logger"; + +/** + * Wire protocols the evaluator can speak, one per HTTP implementation in + * `llm-provider.ts`. A third kind needs a class before it can be a value here. + */ +export const LLM_PROVIDER_KINDS = ["anthropic", "openai-compatible"] as const; +export type LlmProviderKind = (typeof LLM_PROVIDER_KINDS)[number]; + +const KIND_SET = new Set(LLM_PROVIDER_KINDS); + +/** One operator-configured provider. Credentials are deliberately absent: they + * live per project in `project_secrets`, never in an environment variable that + * every project's policy can select. */ +export interface ConfiguredLlmProvider { + name: string; + kind: LlmProviderKind; + baseUrl: string; +} + +/** Provider name to definition. Empty means "Workers AI only" — the default. */ +export type LlmProviderCatalog = ReadonlyMap; + +/** Parsing either yields a catalog or says why the operator's config is unusable. */ +export type LlmProvidersParse = + | { status: "ok"; providers: LlmProviderCatalog } + | { status: "invalid"; reason: string }; + +/** Fields an `LLM_PROVIDERS` entry may carry. */ +const ENTRY_KEYS = new Set(["name", "kind", "baseUrl"]); + +/** + * Provider names are as narrow as deploy names, and for the same reason: a + * policy file selects one by name, and the name is echoed into log lines, + * evaluation reasons, and the derived secret name. + */ +const PROVIDER_NAME_PATTERN = /^[a-z][a-z0-9-]{0,31}$/; + +/** Enough for every provider one instance realistically fronts; bounds the parse. */ +const MAX_PROVIDERS = 16; + +const EMPTY_CATALOG: LlmProviderCatalog = new Map(); + +/** + * The `project_secrets` name holding this provider's credential. + * + * Derived rather than configured so an entry cannot name an arbitrary secret: + * `LLM_PROVIDERS` is operator config, but making the secret *name* configurable + * would let one provider entry read a credential a project stored for something + * else entirely. `PROVIDER_NAME_PATTERN` guarantees the result matches + * `SECRET_NAME_PATTERN` (`^[A-Z][A-Z0-9_]{0,63}$`) — an uppercased 32-char + * name plus the suffix is 40. + */ +export function providerSecretName(providerName: string): string { + return `${providerName.toUpperCase().replace(/-/g, "_")}_API_KEY`; +} + +/** + * Why this host must not be fetched, or null when it is fine. + * + * Loopback, link-local and RFC1918 addresses are the ones that turn an outbound + * fetch into a read of something only the Worker's own network can see. The + * WHATWG URL parser normalizes IPv4 for special schemes, so the obfuscated + * spellings (`0x7f.0.0.1`, `2130706433`) arrive here already in dotted-quad + * form and are caught by the same rules. + * + * **Known limit:** this is an address check, not a resolution check. A public + * DNS name that resolves to 127.0.0.1 defeats it, and nothing at parse time can + * see that. It is defence in depth — the actual control is that only the + * operator writes this variable, and that a policy file can never add to it. + */ +export function blockedHostReason(hostname: string): string | null { + // A trailing dot is the same name to a resolver; lowercase for the literals. + const host = hostname.toLowerCase().replace(/\.$/, ""); + + if (host === "localhost" || host.endsWith(".localhost")) { + return "'localhost' names the loopback interface"; + } + + if (host.startsWith("[") && host.endsWith("]")) { + const inner = host.slice(1, -1); + if (inner === "::1") return "[::1] is the IPv6 loopback address"; + if (inner === "::" || inner === "::0") return "[::] is the IPv6 unspecified address"; + if (inner.startsWith("fe80")) return `${host} is an IPv6 link-local address`; + // fc00::/7 — unique local, the IPv6 equivalent of RFC1918. + if (/^f[cd]/.test(inner)) return `${host} is an IPv6 unique-local address`; + // An IPv4-mapped address exists only to reach an IPv4 host through an IPv6 + // socket; an operator has no reason to write one, and each spelling + // (dotted-quad tail, hex tail) would need its own check to be safe. + if (inner.includes("::ffff:")) return `${host} is an IPv4-mapped IPv6 address`; + return null; + } + + const quad = /^(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$/.exec(host); + if (!quad) return null; + const [a, b] = [Number(quad[1]), Number(quad[2])]; + + if (a === 127) return `${host} is in the loopback range 127.0.0.0/8`; + if (a === 0) return `${host} is in the unspecified range 0.0.0.0/8`; + if (a === 169 && b === 254) return `${host} is in the link-local range 169.254.0.0/16`; + if (a === 10) return `${host} is in the private range 10.0.0.0/8`; + if (a === 172 && b >= 16 && b <= 31) return `${host} is in the private range 172.16.0.0/12`; + if (a === 192 && b === 168) return `${host} is in the private range 192.168.0.0/16`; + return null; +} + +function baseUrlError(name: string, raw: unknown): string | null { + if (typeof raw !== "string" || raw.length === 0) { + return `provider "${name}": "baseUrl" is required and must be a string`; + } + + let url: URL; + try { + url = new URL(raw); + } catch { + return `provider "${name}": "baseUrl" is not a valid absolute URL`; + } + + if (url.protocol !== "https:") { + // Not tidiness: the request carries the project's API key and the diff. + return `provider "${name}": "baseUrl" must use https, not ${url.protocol.replace(":", "")}`; + } + if (url.username !== "" || url.password !== "") { + return `provider "${name}": "baseUrl" must not embed credentials`; + } + if (url.search !== "" || url.hash !== "") { + return `provider "${name}": "baseUrl" must not carry a query string or fragment`; + } + const blocked = blockedHostReason(url.hostname); + if (blocked) { + return `provider "${name}": "baseUrl" host is not reachable from an allowlist — ${blocked}`; + } + return null; +} + +/** + * Parse `LLM_PROVIDERS`, rejecting the whole document on the first problem. + * + * All-or-nothing rather than "keep the entries that parsed", because a + * half-applied allowlist is the failure this is meant to prevent: an operator + * who typo'd one entry would get merges blocked on the projects selecting it + * with nothing anywhere saying why. One rejection with a reason is louder, and + * `llmProvidersConfigError` puts it in the logs on the first request after a + * bad deploy. + * + * An unset (or blank) value is not an error — it is the default, and it means + * Workers AI only. + */ +export function parseLlmProviders(raw: string | undefined): LlmProvidersParse { + if (raw === undefined || raw.trim() === "") return { status: "ok", providers: EMPTY_CATALOG }; + + let parsed: unknown; + try { + parsed = JSON.parse(raw); + } catch (error) { + return { + status: "invalid", + reason: `LLM_PROVIDERS is not valid JSON (${error instanceof Error ? error.message : "parse error"})`, + }; + } + + if (!Array.isArray(parsed)) { + return { status: "invalid", reason: "LLM_PROVIDERS must be a JSON array of provider objects" }; + } + if (parsed.length > MAX_PROVIDERS) { + return { + status: "invalid", + reason: `LLM_PROVIDERS declares ${parsed.length} providers; at most ${MAX_PROVIDERS} are allowed`, + }; + } + + const providers = new Map(); + + for (const [index, entry] of parsed.entries()) { + const at = `LLM_PROVIDERS[${index}]`; + if (typeof entry !== "object" || entry === null || Array.isArray(entry)) { + return { status: "invalid", reason: `${at} must be an object` }; + } + const source = entry as Record; + + const unknownKeys = Object.keys(source).filter((key) => !ENTRY_KEYS.has(key)); + if (unknownKeys.length > 0) { + // Rejected rather than ignored: a misspelled key here is silently + // dropped configuration on a security boundary, and `baseURL` for + // `baseUrl` is exactly the typo an operator would not notice. + return { + status: "invalid", + reason: `${at} has unrecognized field(s) ${unknownKeys.map((k) => `"${k}"`).join(", ")}`, + }; + } + + if (typeof source.name !== "string" || !PROVIDER_NAME_PATTERN.test(source.name)) { + return { + status: "invalid", + reason: `${at}: "name" must match ${PROVIDER_NAME_PATTERN.source} (lowercase letters, digits and dashes, max 32 chars)`, + }; + } + const name = source.name; + if (providers.has(name)) { + return { status: "invalid", reason: `${at}: duplicate provider name "${name}"` }; + } + + if (typeof source.kind !== "string" || !KIND_SET.has(source.kind)) { + return { + status: "invalid", + reason: `provider "${name}": "kind" must be one of ${LLM_PROVIDER_KINDS.join(", ")}`, + }; + } + + const urlProblem = baseUrlError(name, source.baseUrl); + if (urlProblem) return { status: "invalid", reason: urlProblem }; + + providers.set(name, { + name, + kind: source.kind as LlmProviderKind, + // Trailing slashes are stripped by the provider classes when they build a + // path; store what the operator wrote. + baseUrl: source.baseUrl as string, + }); + } + + return { status: "ok", providers }; +} + +/** + * The parse memoized for the isolate's lifetime, keyed on the raw value so a + * test (or a rebind) that changes it is not served a stale catalog. + */ +let cached: { raw: string | undefined; parse: LlmProvidersParse } | null = null; + +/** Drop the memoized parse. Exists for tests; production rebinds per isolate. */ +export function resetLlmProviderCache(): void { + cached = null; +} + +function parseCached(raw: string | undefined): LlmProvidersParse { + if (cached === null || cached.raw !== raw) cached = { raw, parse: parseLlmProviders(raw) }; + return cached.parse; +} + +/** + * The catalog for this environment, or an empty one when the config is unusable. + * + * Empty-on-invalid is the fail-closed direction, not a silent disable: an empty + * catalog means every policy naming a provider is rejected by + * `sanitizeLlmConfig` and blocks merges with a reason naming the provider, while + * projects on Workers AI — which is every project that never opted into BYOK — + * are untouched by an operator's typo. The typo itself is reported by + * `llmProvidersConfigError` and logged here on the isolate's first parse. + */ +export function llmProviderCatalog( + env: Pick, + logger?: Logger, +): LlmProviderCatalog { + const hadCache = cached !== null && cached.raw === env.LLM_PROVIDERS; + const parse = parseCached(env.LLM_PROVIDERS); + if (parse.status === "invalid") { + if (!hadCache) { + logger?.error("LLM_PROVIDERS is invalid; BYOK providers are unavailable", undefined, { + reason: parse.reason, + }); + } + return EMPTY_CATALOG; + } + return parse.providers; +} diff --git a/src/evaluation/policy-loader.ts b/src/evaluation/policy-loader.ts index 206638ee..e4996421 100644 --- a/src/evaluation/policy-loader.ts +++ b/src/evaluation/policy-loader.ts @@ -9,7 +9,14 @@ import { MIN_PHASE_TIMEOUT_MS, MIN_TOTAL_BUDGET_MS, } from "./limits"; -import type { EvalPolicy, EvaluatorConfig, MergePolicy, SandboxEvaluatorConfig } from "./types"; +import type { LlmProviderCatalog } from "./llm-providers"; +import type { + EvalPolicy, + EvaluatorConfig, + LlmEvaluatorConfig, + MergePolicy, + SandboxEvaluatorConfig, +} from "./types"; const DEFAULT_POLICY: EvalPolicy = { evaluators: [{ type: "diff" }], @@ -25,6 +32,15 @@ const DEFAULT_POLICY: EvalPolicy = { */ export const PROTECTED_CONFIG_FILES = [".stratum/policy.yaml", "stratum.config.json"] as const; +/** + * The operator's provider allowlist, as the policy parser sees it. + * + * Defaulted to empty rather than made required so a caller that forgets to pass + * one blocks BYOK policies instead of admitting them: an empty catalog resolves + * no provider name, and an unresolved name fails the file closed. + */ +const NO_PROVIDERS: LlmProviderCatalog = new Map(); + /** Does a unified diff (git-style `diff --git a/… b/…` headers) modify a * protected merge-protection config file? */ export function diffTouchesProtectedConfig(diff: string): boolean { @@ -44,6 +60,9 @@ export async function loadPolicy( token: string, logger: Logger, branch = "main", + /** Operator-configured LLM providers a policy may select from. Omitted means + * none are configured, so any policy naming one fails closed. */ + providers: LlmProviderCatalog = NO_PROVIDERS, ): Promise { const yaml = await readAndParsePolicy( remote, @@ -52,6 +71,7 @@ export async function loadPolicy( "yaml", logger, branch, + providers, ); if (yaml.status === "ok") return yaml.policy; if (yaml.status === "malformed") @@ -64,6 +84,7 @@ export async function loadPolicy( "json", logger, branch, + providers, ); if (json.status === "ok") return json.policy; if (json.status === "malformed") @@ -115,6 +136,7 @@ export function parsePolicyContent( content: string, format: "json" | "yaml", logger: Logger = defaultLogger, + providers: LlmProviderCatalog = NO_PROVIDERS, ): PolicyParse { let parsed: unknown; try { @@ -162,7 +184,7 @@ export function parsePolicyContent( // default) by mutating what it was handed. const declared = raw.evaluators as unknown[]; const evaluators = declared - .map((entry) => sanitizeEvaluator(entry, logger)) + .map((entry) => sanitizeEvaluator(entry, logger, providers)) .filter((entry): entry is EvaluatorConfig => entry !== null); // Any dropped entry fails the gate closed — not just the case where every @@ -223,6 +245,7 @@ async function readAndParsePolicy( format: "json" | "yaml", logger: Logger, branch = "main", + providers: LlmProviderCatalog = NO_PROVIDERS, ): Promise { try { const contentResult = await readFileFromRepo(remote, token, path, logger, branch); @@ -231,7 +254,7 @@ async function readAndParsePolicy( const content = contentResult.data; if (content === null || content === undefined) return { status: "absent" }; - return parsePolicyContent(content, format, logger); + return parsePolicyContent(content, format, logger, providers); } catch (e) { // A transient repo-read blip — treat as absent so it doesn't block every // merge. Sanitization failures cannot reach here; `parsePolicyContent` @@ -404,6 +427,100 @@ function sanitizeSandboxConfig( return config; } +/** Fields an `llm` evaluator entry may carry. Everything else is ignored. */ +const LLM_CONFIG_KEYS = new Set(["type", "provider", "model", "threshold", "maxDiffChars"]); + +/** + * Longest `model` a policy may name. Generous for every provider's identifiers + * (`@cf/meta/llama-3.1-8b-instruct`, `anthropic/claude-sonnet-4-5`) and short + * enough that the value cannot become a payload in its own right. + */ +const MAX_MODEL_LENGTH = 128; + +/** Model identifiers across the three providers: slugs with `/`, `.`, `:`, `@`, `-`, `_`. */ +const MODEL_PATTERN = /^[A-Za-z0-9._:@/-]+$/; + +/** + * Keep only whitelisted, well-typed fields from a user-supplied `llm` entry, + * resolving `provider` against the operator's allowlist. + * + * Returning null fails the **whole policy file** closed and blocks merges (see + * the dropped-entry check in `parsePolicyContent`). That is the opposite of + * what a rejected `deploys:` entry does two files away — `deploy/config.ts` + * says in as many words "do not fix this" — and the asymmetry is deliberate: + * a deploy runs *after* the merge, so a bad entry costs a failed deployment + * row, whereas an `llm` entry is a gate. A gate that cannot run must not pass. + * An unresolvable provider name is precisely a gate that cannot run, so it + * blocks rather than degrading to Workers AI on the operator's bill. + * + * `baseUrl` is never read here, at any level of validation. The endpoint set is + * closed and lives in `LLM_PROVIDERS`; a policy may only *select* from it by + * name. A `baseUrl` in the file is warned about and dropped like any other + * unrecognized key — see `llm-providers.ts` for why the field cannot exist. + */ +export function sanitizeLlmConfig( + source: Record, + logger: Logger, + providers: LlmProviderCatalog, +): LlmEvaluatorConfig | null { + const config: LlmEvaluatorConfig = { type: "llm" }; + + for (const key of Object.keys(source)) { + if (!LLM_CONFIG_KEYS.has(key)) { + logger.warn("Ignoring unrecognized llm evaluator field", { field: key }); + } + } + + if (source.provider !== undefined) { + if (typeof source.provider !== "string") { + logger.error("Dropping llm evaluator entry whose provider is not a string", undefined, { + submitted: forLog(source.provider), + }); + return null; + } + if (!providers.has(source.provider)) { + logger.error( + "Dropping llm evaluator entry naming an unconfigured provider — failing the policy closed", + undefined, + { provider: forLog(source.provider) }, + ); + return null; + } + config.provider = source.provider; + } + + if (source.model !== undefined) { + const model = typeof source.model === "string" ? source.model.trim() : ""; + if (model.length === 0 || model.length > MAX_MODEL_LENGTH || !MODEL_PATTERN.test(model)) { + logger.error("Dropping llm evaluator entry with an unusable model", undefined, { + submitted: forLog(source.model), + }); + return null; + } + config.model = model; + } + + // Clamped, not rejected, for the reason `clampScore` gives: a threshold + // outside [0,1] is a bounded mistake, and 1.5 (nothing passes) is the safe + // direction anyway. -5 is not — it would make every model verdict pass — so + // the clamp is load-bearing rather than cosmetic. + const threshold = clampScore(source.threshold, logger); + if (threshold !== undefined) config.threshold = threshold; + + if (typeof source.maxDiffChars === "number" && Number.isFinite(source.maxDiffChars)) { + // Left unclamped here on purpose: `LLMEvaluator` clamps it against its own + // floor and ceiling, which are the numbers that actually bound the prompt. + config.maxDiffChars = source.maxDiffChars; + } else if (source.maxDiffChars !== undefined) { + logger.warn("Ignoring policy field that is not a finite number", { + field: "llm.maxDiffChars", + submitted: forLog(source.maxDiffChars), + }); + } + + return config; +} + /** * Validate one entry of the user-supplied `evaluators` array. * @@ -417,7 +534,11 @@ function sanitizeSandboxConfig( * passed through by reference, so a returned policy never shares object * identity with the parsed input or with `DEFAULT_POLICY`. */ -function sanitizeEvaluator(raw: unknown, logger: Logger): EvaluatorConfig | null { +function sanitizeEvaluator( + raw: unknown, + logger: Logger, + providers: LlmProviderCatalog, +): EvaluatorConfig | null { if (typeof raw !== "object" || raw === null || Array.isArray(raw)) { logger.warn("Dropping non-object evaluator entry", { submitted: forLog(raw) }); return null; @@ -430,6 +551,11 @@ function sanitizeEvaluator(raw: unknown, logger: Logger): EvaluatorConfig | null if (source.type === "sandbox") return sanitizeSandboxConfig(source, logger); + // Replaces the copy-through below for `llm`, which used to hand the parsed + // object's fields to the evaluator unvalidated — including any `baseUrl` or + // inline credential the file carried. + if (source.type === "llm") return sanitizeLlmConfig(source, logger, providers); + if (source.type === "webhook") { // `url` is required by the type, so check it rather than asserting it. The // evaluator does fail closed on a missing URL, but telling the compiler a diff --git a/src/evaluation/sanitize-policy.ts b/src/evaluation/sanitize-policy.ts index c0e78886..cf7a469f 100644 --- a/src/evaluation/sanitize-policy.ts +++ b/src/evaluation/sanitize-policy.ts @@ -1,19 +1,78 @@ import type { EvalPolicy, EvaluatorConfig } from "./types"; +/** + * Fields whose *name* says the value is a credential. + * + * Matched on the key rather than the value because there is no reliable way to + * recognise a secret by its bytes, and this runs on repository content: whatever + * a user wrote in the policy file, under whatever key they chose. + */ +const CREDENTIAL_WORDS = new Set([ + "secret", + "secrets", + "token", + "password", + "passwd", + "credential", + "credentials", + "authorization", + "bearer", + "key", + "apikey", + "signature", +]); + +/** + * Split a field name into words on camelCase and separator boundaries, so + * `apiKey`, `api_key` and `API-KEY` are all recognised while `monkeys` — which + * merely contains "key" — is not. + */ +function words(key: string): string[] { + return key + .split(/[_\-.\s]+|(?<=[a-z0-9])(?=[A-Z])/) + .map((word) => word.toLowerCase()) + .filter((word) => word.length > 0); +} + +/** Keys stripped from every evaluator entry regardless of shape. */ +function isStrippedKey(key: string): boolean { + // The provider *name* is not a credential, but it names the operator's + // infrastructure to a third party the operator did not choose — and the model + // has no use for it. It costs nothing to withhold, and withholding it is what + // makes adding a provider to a policy free against + // `MAX_POLICY_CONTEXT_CHARS`: the serialized policy does not grow, so a + // project that fits today cannot start failing its gate closed by opting into + // BYOK. + if (key === "provider") return true; + return words(key).some((word) => CREDENTIAL_WORDS.has(word)); +} + /** * Strip credentials from a policy before it leaves the Worker (serialized into - * an LLM prompt or a webhook request body). The webhook secret is used to sign - * outgoing requests; it must never appear in a payload itself. + * an LLM prompt, or a webhook request body). + * + * This is the last thing that runs before the policy is handed to a third + * party, so it strips by key shape rather than by known field. `webhook.secret` + * — the one field this used to handle — is used to sign outgoing requests and + * must never appear in a payload itself. The general rule exists because + * `sanitizeEvaluator` copies an entry of an unmodelled `type` through whole + * (a policy naming a future evaluator), so a credential-shaped field on such an + * entry would otherwise reach the model verbatim. `sanitizeLlmConfig`'s + * whitelist already keeps one off an `llm` entry; this is the backstop for + * every entry it does not model. + * + * Note what this is *not*: a defence against the model seeing secrets in the + * diff. The diff is the thing being reviewed, and `redactSecrets` + * (`deploy/redact.ts`) is explicit that literal-substring redaction is a + * backstop. The structural control is upstream — provider output never reaches + * `EvalResult.reason` (`llm-evaluator.ts`). */ export function sanitizePolicy(policy: EvalPolicy): EvalPolicy { return { ...policy, evaluators: policy.evaluators.map((cfg: EvaluatorConfig) => { - if (cfg.type === "webhook") { - const { secret: _secret, ...rest } = cfg; - return rest; - } - return cfg; + const entries = Object.entries(cfg).filter(([key]) => !isStrippedKey(key)); + return Object.fromEntries(entries) as EvaluatorConfig; }), }; } diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index e5e2ebca..e03b3259 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -206,4 +206,31 @@ export type EvaluatorConfig = } | { type: "webhook"; url: string; secret?: string; timeoutMs?: number } | SandboxEvaluatorConfig - | { type: "llm"; model?: string; threshold?: number; maxDiffChars?: number }; + | LlmEvaluatorConfig; + +/** + * The `llm` evaluator's slice of `.stratum/policy.yaml`. + * + * Named rather than inlined for the reason `SandboxEvaluatorConfig` is: two + * places narrow this shape out of a policy (`LLMEvaluator`, and the BYOK + * provider resolution), and an inline re-declaration in either would drift. + * + * Produced only by `sanitizeLlmConfig` (`policy-loader.ts`), which is a + * whitelist: these four fields are the whole surface a repository's policy file + * may set. There is deliberately **no `baseUrl`** — see `llm-providers.ts`. + */ +export interface LlmEvaluatorConfig { + type: "llm"; + /** + * Selects one of the operator's `LLM_PROVIDERS` entries by name, running the + * evaluation on the project's own credential (BYOK). Absent means Workers AI, + * on the operator's bill. A name the operator has not configured is not a + * fallback: it fails the policy file closed. + */ + provider?: string; + model?: string; + /** Score at or above which the model's verdict is allowed to pass. */ + threshold?: number; + /** Diff characters sent to the model; clamped again by the evaluator. */ + maxDiffChars?: number; +} diff --git a/src/middleware/config-guard.ts b/src/middleware/config-guard.ts index fae4ccb5..ba95f889 100644 --- a/src/middleware/config-guard.ts +++ b/src/middleware/config-guard.ts @@ -1,4 +1,5 @@ import type { MiddlewareHandler } from "hono"; +import { parseLlmProviders } from "../evaluation/llm-providers"; import type { Env } from "../types"; import { createLogger } from "../utils/logger"; @@ -57,20 +58,44 @@ export function entitlementsConfigError( return null; } +/** + * Returns a human-readable error when `LLM_PROVIDERS` is set but unusable, else + * null. + * + * Loud on purpose, and the reason Open Question 3 settled on a typed parse: a + * malformed allowlist silently disables BYOK, and the projects that opted in + * discover it as blocked merges naming a provider the operator believes is + * configured. Unset is not a problem — it is the default, and it means Workers + * AI only. + */ +export function llmProvidersConfigError(env: Pick): string | null { + const parse = parseLlmProviders(env.LLM_PROVIDERS); + if (parse.status === "invalid") { + // One template literal rather than concatenation: `useTemplate` rejects + // mixing an interpolated piece with `+`, and the sibling above can use plain + // strings only because it interpolates nothing. + return `LLM_PROVIDERS is set but could not be parsed (${parse.reason}) — no BYOK provider is available, so every policy selecting one blocks merges until this is fixed. Correct the value or unset it to run on Workers AI.`; + } + return null; +} + // Log the config problems at most once per isolate rather than on every request. let hasLoggedConfigError = false; /** - * Surfaces an incoherent RepoDO/R2 or entitlements configuration loudly in - * Workers Logs on the first request after a bad deploy, instead of only when a - * later merge throws — or, for entitlements, never. Non-fatal: reads and the UI - * still work in both cases, so we log rather than reject the request. + * Surfaces an incoherent RepoDO/R2, entitlements or LLM-provider configuration + * loudly in Workers Logs on the first request after a bad deploy, instead of + * only when a later merge throws — or, for entitlements, never. Non-fatal: + * reads and the UI still work in every case, so we log rather than reject the + * request. */ export const configGuardMiddleware: MiddlewareHandler<{ Bindings: Env }> = async (c, next) => { if (!hasLoggedConfigError) { - const problems = [repoDoConfigError(c.env), entitlementsConfigError(c.env)].filter( - (problem): problem is string => problem !== null, - ); + const problems = [ + repoDoConfigError(c.env), + entitlementsConfigError(c.env), + llmProvidersConfigError(c.env), + ].filter((problem): problem is string => problem !== null); if (problems.length > 0) { hasLoggedConfigError = true; const logger = createLogger({ component: "config-guard" }); diff --git a/src/routes/changes.ts b/src/routes/changes.ts index 33e80157..a29fdd39 100644 --- a/src/routes/changes.ts +++ b/src/routes/changes.ts @@ -1,5 +1,6 @@ import { type Context, Hono } from "hono"; import { diffTouchesProtectedConfig, loadPolicy } from "../evaluation"; +import { llmProviderCatalog } from "../evaluation/llm-providers"; import type { EvalPolicy } from "../evaluation/types"; import { GitHubClient } from "../github/client"; import { buildEvaluationReport, reportEvaluationToGitHub } from "../github/sync"; @@ -135,6 +136,7 @@ async function loadMergePolicyCached( tok.data, logger, projectDefaultBranch(project), + llmProviderCatalog(env, logger), ); if (cacheable) { policyCache.set(project.id, { policy: loaded, expires: Date.now() + POLICY_CACHE_TTL_MS }); @@ -1383,7 +1385,13 @@ app.post("/changes/:id/evaluate", async (c) => { if (!workspaceReadToken.success) return internalError(workspaceReadToken.error.message); const branch = projectDefaultBranch(project); - const policy = await loadPolicy(project.remote, projectReadToken.data, logger, branch); + const policy = await loadPolicy( + project.remote, + projectReadToken.data, + logger, + branch, + llmProviderCatalog(c.env, logger), + ); const diffResult = await getDiffBetweenRepos( project.remote, @@ -1407,7 +1415,7 @@ app.post("/changes/:id/evaluate", async (c) => { baseOid, } = diffResult.data; - const evaluators = buildEvaluators(c.env, policy, project, logger, { + const evaluators = await buildEvaluators(c.env, policy, project, logger, { remote: workspace.remote, token: workspaceReadToken.data, ref: evaluatedSha, diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index 925e6e22..82c41d9d 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -1,4 +1,5 @@ import { Hono } from "hono"; +import { llmProviderCatalog } from "../evaluation/llm-providers"; import { loadPolicy } from "../evaluation/policy-loader"; import { scanContentForSecrets } from "../evaluation/secret-scanner"; import { checkResolutionMergeProtection } from "../merge/protection"; @@ -656,13 +657,19 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { } const { diff, baseSha } = diffResult.data; - const policy = await loadPolicy(project.remote, projectToken.data, logger, branch); + const policy = await loadPolicy( + project.remote, + projectToken.data, + logger, + branch, + llmProviderCatalog(c.env, logger), + ); // No workspace repo access is passed for the sandbox evaluator: the content // being judged here has no commit of its own yet — that's the point, it must // pass BEFORE resolveConflict creates one — so there is no ref a sandbox // could check out. A policy naming `sandbox` fails closed via // UnavailableEvaluator, same as any other missing prerequisite. - const evaluators = buildEvaluators(c.env, policy, project, logger); + const evaluators = await buildEvaluators(c.env, policy, project, logger); // `buildManualResolutionDiff` resolved this base from the clone it built the // diff on, so it names the tree the resolution actually applies to (#274). // This path runs the LLM evaluator too, so it needs a payer as much as diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index af4daec7..1b6c6427 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -9,6 +9,8 @@ import { diffTouchesProtectedConfig, loadPolicy, } from "../evaluation"; +import { resolveLlmProvider } from "../evaluation/llm-byok"; +import { llmProviderCatalog } from "../evaluation/llm-providers"; import type { SandboxRepoAccess } from "../evaluation/sandbox-evaluator"; import type { BillingContext, @@ -153,13 +155,19 @@ export function billingContextFor(project: ProjectEntry): BillingContext | undef * ops decision) — a policy naming `sandbox` fails closed with a reason that * says exactly which prerequisite is missing. */ -export function buildEvaluators( +export async function buildEvaluators( env: Env, policy: EvalPolicy, project: ProjectEntry, logger: Logger, workspaceRepo?: SandboxRepoAccess, -): Array<{ type: string; evaluator: Evaluator }> { +): Promise> { + // Async only because of this: a policy that selects a BYOK provider needs the + // project's credential resolved before the evaluator exists. It is a no-op — + // no D1 read, no key derivation — for every policy that names no provider, + // which is every policy that has not opted in. + const llmProvider = await resolveLlmProvider(env, project, policy, logger); + // Guarded, because a legacy entry can carry no namespace — see // `projectDisplayName`, and this file's own `resolveProjectHead`. const projectName = projectDisplayName(project); @@ -174,7 +182,25 @@ export function buildEvaluators( return [{ type: "diff", evaluator: new DiffEvaluator() }]; case "webhook": return [{ type: "webhook", evaluator: new WebhookEvaluator() }]; - case "llm": + case "llm": { + // Every failure to resolve a project-supplied credential lands here, + // and NONE of them falls back to `env.AI`. Falling back would move + // the spend to the operator's bill — the hole BYOK closes — and would + // turn a misconfigured gate into one that quietly passes on someone + // else's budget. Fail closed with the reason instead. + if (llmProvider.status === "unavailable") { + return [ + { type: "llm", evaluator: new UnavailableEvaluator("llm", llmProvider.reason) }, + ]; + } + if (llmProvider.status === "byok") { + return [ + { + type: "llm", + evaluator: new LLMEvaluator(llmProvider.provider, "byok"), + }, + ]; + } if (env.AI) return [{ type: "llm", evaluator: new LLMEvaluator(new WorkersAiProvider(env.AI)) }]; return [ @@ -183,6 +209,7 @@ export function buildEvaluators( evaluator: new UnavailableEvaluator("llm", "AI binding is not configured"), }, ]; + } case "sandbox": if (!env.SANDBOX) { // Fail closed with an actionable reason: the [[sandboxes]] binding @@ -421,7 +448,13 @@ export async function createChangeWithEvaluation( } const branch = projectDefaultBranch(project); - const policy = await loadPolicy(project.remote, projectReadToken.data, logger, branch); + const policy = await loadPolicy( + project.remote, + projectReadToken.data, + logger, + branch, + llmProviderCatalog(env, logger), + ); const diffResult = await getDiffBetweenRepos( project.remote, @@ -447,7 +480,7 @@ export async function createChangeWithEvaluation( baseOid, } = diffResult.data; - const evaluators = buildEvaluators(env, policy, project, logger, { + const evaluators = await buildEvaluators(env, policy, project, logger, { remote: workspaceRemote, token: workspaceReadToken.data, ref: evaluatedSha, diff --git a/src/storage/costs.ts b/src/storage/costs.ts index 777474bd..f2f8f410 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -186,9 +186,18 @@ export async function recordCosts( // evaluation. Filtering here keeps one bad number from taking the ledger with // it, and is why `upsertUsage`'s own guard is a backstop rather than the // first line of defence. - const usable = samples.filter((sample) => Number.isFinite(sample.quantity)); + // + // Negative quantities are dropped for a different reason: they write fine. + // A row for negative spend is never right — it silently refunds an owner's + // month through `accumulateUsage` — and no meter this records can go + // backwards. The providers validate their own token counts (`usageOrNothing` + // in `evaluation/llm-provider.ts`); this is the defence that does not depend + // on every future sample source remembering to. + const usable = samples.filter( + (sample) => Number.isFinite(sample.quantity) && sample.quantity >= 0, + ); if (usable.length !== samples.length) { - logger.warn("Cost samples dropped: quantity is not finite", { + logger.warn("Cost samples dropped: quantity is not a finite, non-negative number", { project: opts.project, dropped: samples.length - usable.length, }); diff --git a/src/storage/project-secrets.ts b/src/storage/project-secrets.ts index f9df1f56..21fcd038 100644 --- a/src/storage/project-secrets.ts +++ b/src/storage/project-secrets.ts @@ -127,12 +127,29 @@ async function resolveKey(env: SecretKeyEnv): Promise } | null = null; + /** * Creates or replaces a project's secret, returning only its metadata. * diff --git a/src/types.ts b/src/types.ts index 81b9ee9b..3d3ae430 100644 --- a/src/types.ts +++ b/src/types.ts @@ -73,7 +73,23 @@ export interface AiBinding { messages?: Array<{ role: string; content: string }>; prompt?: string; }, - ): Promise<{ response?: string } | ReadableStream>; + ): Promise< + | { + response?: string; + /** + * Token counts Workers AI actually reported. Declared because it is + * genuinely there and this repo was throwing it away: Cloudflare's + * generated types give a text-generation output as `response`, + * `tool_calls` and `usage`, with the counts named `prompt_tokens` / + * `completion_tokens` (`AiTextGenerationOutput` and `UsageTags` in + * `@cloudflare/workers-types`). Optional field by field, unlike those + * types, because a binding that omits them must be a cost sample marked + * `estimated`, not a type error — `usageOrNothing` decides which. + */ + usage?: { prompt_tokens?: number; completion_tokens?: number }; + } + | ReadableStream + >; } export interface SandboxBinding { @@ -166,6 +182,15 @@ export interface Env { * Setting it without BILLING_SERVICE_URL is enforcement that looks on and does * nothing — `entitlementsConfigError` flags exactly that. */ ENTITLEMENTS_ENFORCE?: string; + /** JSON array of the operator's LLM providers, e.g. + * `[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]`. + * The ONLY place a provider endpoint may be named: a project's policy file + * selects one of these by name and can never supply a `baseUrl`. Unset (the + * default) means Workers AI only. Credentials are NOT here — they live per + * project in `project_secrets`. Parsed by `parseLlmProviders`, which rejects + * a non-https, credential-bearing, loopback, link-local or private-range + * `baseUrl` at parse time. */ + LLM_PROVIDERS?: string; ANALYTICS?: AnalyticsEngineDataset; SANDBOX?: SandboxBinding; AI?: AiBinding; diff --git a/tests/cost-attribution.test.ts b/tests/cost-attribution.test.ts index c63149a1..2d9bd069 100644 --- a/tests/cost-attribution.test.ts +++ b/tests/cost-attribution.test.ts @@ -283,6 +283,27 @@ describe("recordCosts attribution columns", () => { expect(kept[0]?.kind).toBe("git_ops"); }); + it("drops a negative quantity, which would write fine and refund a month", async () => { + // Unlike NaN this one costs nothing to store, which is the problem: it + // becomes a ledger row for negative spend and decrements the owner's usage + // period. No meter recorded here can run backwards. + const { db, raw } = makeSqliteD1(); + const result = await recordCosts( + db, + logger, + { project: "p", changeId: "chg_1", ownerId: "user_alice", ownerType: "user" }, + [ + { kind: "llm_tokens", quantity: -1e12 }, + { kind: "git_ops", quantity: 2 }, + ], + ); + + expect(result.success).toBe(true); + const kept = rows(raw); + expect(kept).toHaveLength(1); + expect(kept[0]?.kind).toBe("git_ops"); + }); + it("serves an owner-scoped, newest-first history from the new index", async () => { // What idx_costs_owner exists for. `created_at` is written only by // recordCosts as an ISO 8601 string, so ordering it in SQL is safe here — diff --git a/tests/evaluator-billing-context.test.ts b/tests/evaluator-billing-context.test.ts index 3b4a5a8b..29b2c5ab 100644 --- a/tests/evaluator-billing-context.test.ts +++ b/tests/evaluator-billing-context.test.ts @@ -286,8 +286,9 @@ describe("billingContextFor", () => { describe("buildEvaluators — ProjectEntry replaces the display name", () => { const project = projectEntry(); - function evaluatorFor(policy: EvalPolicy, type: string, env: Env = makeEnv()) { - const entry = buildEvaluators(env, policy, project, logger).find((e) => e.type === type); + async function evaluatorFor(policy: EvalPolicy, type: string, env: Env = makeEnv()) { + const built = await buildEvaluators(env, policy, project, logger); + const entry = built.find((e) => e.type === type); expect(entry).toBeDefined(); return entry?.evaluator; } @@ -296,30 +297,32 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { // file substitutes above, not the shipped classes. The assertions still // separate the wired branch from `UnavailableEvaluator`, which is what they // are for — they just don't prove the real class was constructed. - it("always prepends the blocking secret scan", () => { - const built = buildEvaluators(makeEnv(), { evaluators: [] }, project, logger); + it("always prepends the blocking secret scan", async () => { + const built = await buildEvaluators(makeEnv(), { evaluators: [] }, project, logger); expect(built.map((e) => e.type)).toEqual(["secret_scan"]); expect(built[0]?.evaluator).toBeInstanceOf(SecretScanEvaluator); }); - it("builds a real DiffEvaluator for a diff entry", () => { - expect(evaluatorFor({ evaluators: [{ type: "diff" }] }, "diff")).toBeInstanceOf(DiffEvaluator); + it("builds a real DiffEvaluator for a diff entry", async () => { + expect(await evaluatorFor({ evaluators: [{ type: "diff" }] }, "diff")).toBeInstanceOf( + DiffEvaluator, + ); }); - it("builds a real WebhookEvaluator for a webhook entry", () => { + it("builds a real WebhookEvaluator for a webhook entry", async () => { const policy: EvalPolicy = { evaluators: [{ type: "webhook", url: "https://hook.test" }] }; - expect(evaluatorFor(policy, "webhook")).toBeInstanceOf(WebhookEvaluator); + expect(await evaluatorFor(policy, "webhook")).toBeInstanceOf(WebhookEvaluator); }); - it("wires the llm evaluator, not UnavailableEvaluator, when the AI binding is present", () => { - const evaluator = evaluatorFor({ evaluators: [{ type: "llm" }] }, "llm"); + it("wires the llm evaluator, not UnavailableEvaluator, when the AI binding is present", async () => { + const evaluator = await evaluatorFor({ evaluators: [{ type: "llm" }] }, "llm"); expect(evaluator).toBeInstanceOf(LLMEvaluator); expect(evaluator).not.toBeInstanceOf(UnavailableEvaluator); }); it("still fails the llm evaluator closed when the AI binding is missing", async () => { const policy: EvalPolicy = { evaluators: [{ type: "llm" }] }; - const evaluator = evaluatorFor(policy, "llm", {} as Env); + const evaluator = await evaluatorFor(policy, "llm", {} as Env); expect(evaluator).toBeInstanceOf(UnavailableEvaluator); const result = await evaluator?.evaluate("", policy, logger); expect(result?.success).toBe(true); @@ -331,7 +334,7 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { it("still fails the sandbox evaluator closed when the binding is missing", async () => { const policy: EvalPolicy = { evaluators: [{ type: "sandbox" }] }; - const evaluator = evaluatorFor(policy, "sandbox"); + const evaluator = await evaluatorFor(policy, "sandbox"); expect(evaluator).toBeInstanceOf(UnavailableEvaluator); const result = await evaluator?.evaluate("", policy, logger); expect(result?.success).toBe(true); @@ -339,12 +342,12 @@ describe("buildEvaluators — ProjectEntry replaces the display name", () => { expect(result.data.reason).toContain("SANDBOX binding is not configured"); }); - it("drops an unknown evaluator type and names the project in the warning", () => { + it("drops an unknown evaluator type and names the project in the warning", async () => { const warn = vi.fn(); const warningLogger = { ...logger, warn } as unknown as Logger; const policy = { evaluators: [{ type: "quantum" }] } as unknown as EvalPolicy; - const built = buildEvaluators(makeEnv(), policy, project, warningLogger); + const built = await buildEvaluators(makeEnv(), policy, project, warningLogger); expect(built.map((e) => e.type)).toEqual(["secret_scan"]); expect(warn).toHaveBeenCalledTimes(1); diff --git a/tests/llm-byok.test.ts b/tests/llm-byok.test.ts new file mode 100644 index 00000000..68c23fa8 --- /dev/null +++ b/tests/llm-byok.test.ts @@ -0,0 +1,730 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { resolveLlmProvider } from "../src/evaluation/llm-byok"; +import { AnthropicProvider, OpenAiCompatibleProvider } from "../src/evaluation/llm-provider"; +import { + blockedHostReason, + llmProviderCatalog, + parseLlmProviders, + providerSecretName, + resetLlmProviderCache, +} from "../src/evaluation/llm-providers"; +import { parsePolicyContent, sanitizeLlmConfig } from "../src/evaluation/policy-loader"; +import { sanitizePolicy } from "../src/evaluation/sanitize-policy"; +import type { EvalPolicy } from "../src/evaluation/types"; +import { llmProvidersConfigError } from "../src/middleware/config-guard"; +import { buildEvaluators, runEvaluation } from "../src/services/change-flow"; +import { putSecret } from "../src/storage/project-secrets"; +import type { AiBinding, Env, ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { makeSqliteD1, makeThrowingD1 } from "./helpers/sqlite-d1"; + +/** + * Task 7 — BYOK configuration, credentials, and the threat controls. + * + * The threat model these test (PRD §7): the reachable attacker is a project + * OWNER using the operator's Worker as an SSRF/exfiltration proxy, or a writer + * whose policy edit lands at merge. Policy is loaded from the DEFAULT BRANCH, so + * a proposed edit never governs its own evaluation — which is why the assertions + * below are about what a *merged* policy can reach, not about live escalation. + */ + +const logger: Logger = { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + child: vi.fn(() => logger), +}; + +/** The one entry every test's operator config uses unless it is the thing under test. */ +const ANTHROPIC_ENTRY = { + name: "anthropic", + kind: "anthropic", + baseUrl: "https://api.anthropic.com", +}; + +const PROVIDERS_JSON = JSON.stringify([ANTHROPIC_ENTRY]); + +function catalogFrom(json: string | undefined) { + resetLlmProviderCache(); + return llmProviderCatalog({ LLM_PROVIDERS: json }, logger); +} + +function project(id = "proj_byok"): ProjectEntry { + return { + id, + name: "repo", + slug: "repo", + namespace: "@alice", + ownerId: "user_alice", + ownerType: "user", + remote: "https://artifacts.example.com/repos/repo", + createdAt: "2026-01-01T00:00:00.000Z", + }; +} + +beforeEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + vi.clearAllMocks(); + resetLlmProviderCache(); +}); + +// --------------------------------------------------------------------------- +// 7.1 — LLM_PROVIDERS is parsed and validated once, at parse time +// --------------------------------------------------------------------------- + +describe("parseLlmProviders — the allowlist is the only place an endpoint is named", () => { + it("treats unset and blank as 'Workers AI only' rather than an error", () => { + for (const raw of [undefined, "", " "]) { + const parse = parseLlmProviders(raw); + expect(parse.status).toBe("ok"); + if (parse.status !== "ok") return; + expect(parse.providers.size).toBe(0); + } + }); + + it("parses a well-formed entry into a name-keyed catalog", () => { + const parse = parseLlmProviders(PROVIDERS_JSON); + expect(parse.status).toBe("ok"); + if (parse.status !== "ok") return; + expect(parse.providers.get("anthropic")).toEqual({ + name: "anthropic", + kind: "anthropic", + baseUrl: "https://api.anthropic.com", + }); + }); + + it("rejects a non-https baseUrl", () => { + const parse = parseLlmProviders( + JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl: "http://api.anthropic.com" }]), + ); + expect(parse.status).toBe("invalid"); + if (parse.status !== "invalid") return; + expect(parse.reason).toContain("must use https"); + }); + + // The whole point of the allowlist: an operator (or a compromised config) must + // not be able to point the evaluator at something only this Worker can reach. + it.each([ + ["https://localhost/v1", "loopback"], + ["https://sub.localhost/v1", "loopback"], + ["https://127.0.0.1/v1", "127.0.0.0/8"], + ["https://127.9.9.9/v1", "127.0.0.0/8"], + // The URL parser normalizes these to 127.0.0.1 before the check sees them. + ["https://2130706433/v1", "127.0.0.0/8"], + ["https://0x7f.0.0.1/v1", "127.0.0.0/8"], + ["https://[::1]/v1", "loopback"], + ["https://0.0.0.0/v1", "0.0.0.0/8"], + // The cloud metadata endpoint — the single most valuable SSRF target there is. + ["https://169.254.169.254/latest", "169.254.0.0/16"], + ["https://10.1.2.3/v1", "10.0.0.0/8"], + ["https://172.16.0.1/v1", "172.16.0.0/12"], + ["https://172.31.255.254/v1", "172.16.0.0/12"], + ["https://192.168.1.1/v1", "192.168.0.0/16"], + ["https://[fe80::1]/v1", "link-local"], + ["https://[fd00::1]/v1", "unique-local"], + ["https://[::ffff:127.0.0.1]/v1", "IPv4-mapped"], + ])("rejects %s at parse time", (baseUrl, expected) => { + const parse = parseLlmProviders(JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl }])); + expect(parse.status, `${baseUrl} was accepted`).toBe("invalid"); + if (parse.status !== "invalid") return; + expect(parse.reason).toContain(expected); + }); + + it("does not over-reject public addresses that merely look private", () => { + // 172.32/16 is outside the 172.16/12 block, and 10.x only matches the first + // octet — a check written as a string prefix would fail both of these. + for (const baseUrl of ["https://172.32.0.1/v1", "https://110.0.0.1/v1"]) { + expect(blockedHostReason(new URL(baseUrl).hostname)).toBeNull(); + } + }); + + it("rejects credentials embedded in the URL", () => { + const parse = parseLlmProviders( + JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl: "https://user:pass@api.anthropic.com" }]), + ); + expect(parse.status).toBe("invalid"); + if (parse.status !== "invalid") return; + expect(parse.reason).toContain("must not embed credentials"); + }); + + it.each([ + [[{ ...ANTHROPIC_ENTRY, baseURL: "https://x.example" }], "unrecognized field"], + [[{ ...ANTHROPIC_ENTRY, kind: "ollama" }], '"kind" must be one of'], + [[{ ...ANTHROPIC_ENTRY, name: "Anthropic" }], '"name" must match'], + [[ANTHROPIC_ENTRY, { ...ANTHROPIC_ENTRY }], "duplicate provider name"], + [[{ kind: "anthropic", name: "x" }], '"baseUrl" is required'], + [["anthropic"], "must be an object"], + ])("rejects a structurally wrong entry (%#)", (entries, expected) => { + const parse = parseLlmProviders(JSON.stringify(entries)); + expect(parse.status).toBe("invalid"); + if (parse.status !== "invalid") return; + expect(parse.reason).toContain(expected); + }); + + it("rejects a non-array and unparseable JSON", () => { + expect(parseLlmProviders('{"name":"x"}').status).toBe("invalid"); + expect(parseLlmProviders("not json").status).toBe("invalid"); + }); + + it("fails loudly rather than silently disabling BYOK", () => { + // Loud is three things at once: the parse says why, the config guard + // surfaces it on the first request after the deploy, and the catalog is + // empty so every policy naming a provider blocks instead of quietly + // running on the operator's Workers AI bill. + const bad = '[{"name":"anthropic","kind":"anthropic","baseURL":"https://api.anthropic.com"}]'; + + expect(llmProvidersConfigError({ LLM_PROVIDERS: bad })).toContain("unrecognized field"); + expect(llmProvidersConfigError({ LLM_PROVIDERS: PROVIDERS_JSON })).toBeNull(); + expect(llmProvidersConfigError({})).toBeNull(); + + expect(catalogFrom(bad).size).toBe(0); + expect(logger.error).toHaveBeenCalledWith( + "LLM_PROVIDERS is invalid; BYOK providers are unavailable", + undefined, + expect.objectContaining({ reason: expect.stringContaining("unrecognized field") }), + ); + }); + + it("derives a secret name that satisfies SECRET_NAME_PATTERN", () => { + expect(providerSecretName("anthropic")).toBe("ANTHROPIC_API_KEY"); + expect(providerSecretName("my-gateway")).toBe("MY_GATEWAY_API_KEY"); + // The 32-char maximum a provider name may reach, plus the suffix, still + // fits the 64-char store limit. + const longest = `a${"b".repeat(31)}`; + expect(/^[A-Z][A-Z0-9_]{0,63}$/.test(providerSecretName(longest))).toBe(true); + }); + + it("memoizes the parse per isolate but re-parses when the value changes", () => { + resetLlmProviderCache(); + const first = llmProviderCatalog({ LLM_PROVIDERS: PROVIDERS_JSON }, logger); + const second = llmProviderCatalog({ LLM_PROVIDERS: PROVIDERS_JSON }, logger); + expect(second).toBe(first); + const other = llmProviderCatalog( + { LLM_PROVIDERS: JSON.stringify([{ ...ANTHROPIC_ENTRY, name: "other" }]) }, + logger, + ); + expect(other.has("other")).toBe(true); + }); +}); + +// --------------------------------------------------------------------------- +// 7.2 — sanitizeLlmConfig: a whitelist, and a name the operator configured +// --------------------------------------------------------------------------- + +describe("sanitizeLlmConfig — the policy may select a provider, never describe one", () => { + const providers = () => { + const parse = parseLlmProviders(PROVIDERS_JSON); + if (parse.status !== "ok") throw new Error("fixture"); + return parse.providers; + }; + + function parsePolicy(policy: unknown, json = PROVIDERS_JSON) { + return parsePolicyContent(JSON.stringify(policy), "json", logger, catalogFrom(json)); + } + + it("keeps a provider the operator configured", () => { + const parsed = parsePolicy({ evaluators: [{ type: "llm", provider: "anthropic" }] }); + expect(parsed.status).toBe("ok"); + if (parsed.status !== "ok") return; + expect(parsed.policy.evaluators[0]).toEqual({ type: "llm", provider: "anthropic" }); + expect(parsed.policy.configError).toBeUndefined(); + }); + + it("fails the whole policy CLOSED on a provider name the operator has not configured", () => { + // The evaluator-side convention (`policy-loader.ts`'s dropped-entry check), + // deliberately opposite to a rejected `deploys:` entry: an unresolvable + // provider means the gate cannot run, and a gate that cannot run must not + // pass. + const parsed = parsePolicy({ evaluators: [{ type: "llm", provider: "evil-corp" }] }); + expect(parsed.status).toBe("malformed"); + if (parsed.status !== "malformed") return; + expect(parsed.reason).toContain("unusable entr"); + }); + + it("blocks merges rather than degrading to Workers AI", () => { + // What "fails closed" has to mean end to end: the loaded policy carries a + // configError, which the merge gate treats as blocking. A policy that + // silently dropped the entry would let the change through on its remaining + // evaluators. + const parsed = parsePolicy({ + evaluators: [{ type: "llm", provider: "evil-corp" }, { type: "diff" }], + }); + expect(parsed.status).toBe("malformed"); + }); + + it("fails closed when NO providers are configured at all", () => { + // Not a special case, and it must not be: an operator who never set + // LLM_PROVIDERS has an empty catalog, which resolves no name. + const parsed = parsePolicyContent( + JSON.stringify({ evaluators: [{ type: "llm", provider: "anthropic" }] }), + "json", + logger, + catalogFrom(undefined), + ); + expect(parsed.status).toBe("malformed"); + }); + + it("ignores a policy-supplied baseUrl entirely", () => { + // The SSRF control. The field is not validated, not sanitized and not + // preferred-over — it does not exist on the way through. + const config = sanitizeLlmConfig( + { + type: "llm", + provider: "anthropic", + baseUrl: "https://evil.example.com", + base_url: "https://evil.example.com", + endpoint: "https://evil.example.com", + }, + logger, + providers(), + ); + expect(config).toEqual({ type: "llm", provider: "anthropic" }); + expect(JSON.stringify(config)).not.toContain("evil.example.com"); + }); + + it("drops an inline credential a policy tried to smuggle in", () => { + const config = sanitizeLlmConfig( + { type: "llm", provider: "anthropic", apiKey: "sk-ant-inline", secret: "s" }, + logger, + providers(), + ); + expect(JSON.stringify(config)).not.toContain("sk-ant-inline"); + }); + + it("fails closed on a non-string provider", () => { + expect(sanitizeLlmConfig({ type: "llm", provider: 7 }, logger, providers())).toBeNull(); + expect(sanitizeLlmConfig({ type: "llm", provider: {} }, logger, providers())).toBeNull(); + }); + + it("keeps the three non-provider fields and clamps the threshold", () => { + const config = sanitizeLlmConfig( + { type: "llm", model: "claude-sonnet-4-5", threshold: 0.9, maxDiffChars: 2000 }, + logger, + providers(), + ); + expect(config).toEqual({ + type: "llm", + model: "claude-sonnet-4-5", + threshold: 0.9, + maxDiffChars: 2000, + }); + + // A negative threshold would make every model verdict pass — the gate + // disabled by a single character, which is why this clamps rather than + // being copied through as it used to be. + const hostile = sanitizeLlmConfig({ type: "llm", threshold: -5 }, logger, providers()); + expect(hostile?.threshold).toBe(0); + }); + + it("fails closed on a model that is not a usable identifier", () => { + for (const model of [42, "", "a".repeat(129), "model\nInjected: header"]) { + expect(sanitizeLlmConfig({ type: "llm", model }, logger, providers())).toBeNull(); + } + }); + + it("keeps the deploy asymmetry intact: a bad deploy entry does NOT block merges", () => { + // The two conventions sit two files apart and look contradictory. This + // pins both at once so neither is "fixed" into the other. + const parsed = parsePolicy({ + evaluators: [{ type: "diff" }], + deploys: [{ name: "NOT A SLUG", target: "vercel" }], + }); + expect(parsed.status).toBe("ok"); + if (parsed.status !== "ok") return; + expect(parsed.policy.configError).toBeUndefined(); + expect(parsed.policy.deployRejections).toHaveLength(1); + }); +}); + +// --------------------------------------------------------------------------- +// 7.3 / 7.4 — the credential, and failing closed on every way of not having it +// --------------------------------------------------------------------------- + +describe("resolveLlmProvider — the credential path", () => { + const KEY = "test-deploy-secret-key-0123456789"; + + function makeEnv(db: D1Database, overrides: Partial = {}): Env { + return { + DB: db, + DEPLOY_SECRET_KEY: KEY, + LLM_PROVIDERS: PROVIDERS_JSON, + ...overrides, + } as unknown as Env; + } + + async function withStoredKey(value = "sk-ant-real-credential", masterKey = KEY) { + const { db } = makeSqliteD1(); + const stored = await putSecret( + db, + logger, + { DEPLOY_SECRET_KEY: masterKey }, + { projectId: project().id, name: "ANTHROPIC_API_KEY", value, actorId: "user_alice" }, + ); + expect(stored.success).toBe(true); + return db; + } + + const byokPolicy: EvalPolicy = { evaluators: [{ type: "llm", provider: "anthropic" }] }; + + it("skips the secret load entirely when the policy names no provider", async () => { + // The skip is the performance requirement, so it is asserted the only way + // that cannot be faked: a database that throws on any use at all. This is + // the overwhelmingly common case — every project that never opted in. + const env = makeEnv(makeThrowingD1("D1 must not be touched")); + const selection = await resolveLlmProvider( + env, + project(), + { evaluators: [{ type: "llm" }, { type: "diff" }] }, + logger, + ); + expect(selection.status).toBe("platform"); + }); + + it("builds the provider matching the configured kind", async () => { + const db = await withStoredKey(); + const anthropic = await resolveLlmProvider(makeEnv(db), project(), byokPolicy, logger); + expect(anthropic.status).toBe("byok"); + if (anthropic.status !== "byok") return; + expect(anthropic.provider).toBeInstanceOf(AnthropicProvider); + + const openAiEnv = makeEnv(db, { + LLM_PROVIDERS: JSON.stringify([ + { name: "anthropic", kind: "openai-compatible", baseUrl: "https://api.openai.com/v1" }, + ]), + }); + resetLlmProviderCache(); + const openai = await resolveLlmProvider(openAiEnv, project(), byokPolicy, logger); + expect(openai.status).toBe("byok"); + if (openai.status !== "byok") return; + expect(openai.provider).toBeInstanceOf(OpenAiCompatibleProvider); + }); + + it("derives the PBKDF2 key once per isolate across evaluations", async () => { + // 100k iterations, now on the change-creation path, which already does two + // clones and a diff. Three operations against the same master key — the + // write, then two evaluations — must pay for it exactly once. A master key + // no earlier test used, so the isolate's cache starts cold for it. + const coldKey = "a-master-key-this-test-alone-uses"; + const deriveKey = vi.spyOn(crypto.subtle, "deriveKey"); + + const db = await withStoredKey("sk-ant-real-credential", coldKey); + const env = makeEnv(db, { DEPLOY_SECRET_KEY: coldKey }); + expect((await resolveLlmProvider(env, project(), byokPolicy, logger)).status).toBe("byok"); + expect((await resolveLlmProvider(env, project(), byokPolicy, logger)).status).toBe("byok"); + + expect(deriveKey).toHaveBeenCalledTimes(1); + }); + + it("fails closed with a distinct reason for each of the three secret outcomes", async () => { + const stored = await withStoredKey(); + const empty = makeSqliteD1().db; + + const noMasterKey = await resolveLlmProvider( + makeEnv(stored, { DEPLOY_SECRET_KEY: undefined }), + project(), + byokPolicy, + logger, + ); + const missing = await resolveLlmProvider(makeEnv(empty), project(), byokPolicy, logger); + // A row that exists but authenticates under a different master key: exactly + // what a DEPLOY_SECRET_KEY rotation leaves behind. + const rotated = await resolveLlmProvider( + makeEnv(stored, { DEPLOY_SECRET_KEY: "a-different-master-key-9876543210" }), + project(), + byokPolicy, + logger, + ); + + const reasons = [noMasterKey, missing, rotated].map((selection) => { + expect(selection.status).toBe("unavailable"); + return selection.status === "unavailable" ? selection.reason : ""; + }); + + expect(reasons[0]).toContain("DEPLOY_SECRET_KEY"); + expect(reasons[1]).toContain("is not set for this project"); + expect(reasons[2]).toContain("could not be decrypted"); + expect(new Set(reasons).size).toBe(3); + // Each reason names the secret so an operator knows what to re-enter. + for (const reason of reasons) expect(reason).toContain("ANTHROPIC_API_KEY"); + }); + + it("fails closed when the database itself is unavailable", async () => { + const selection = await resolveLlmProvider( + makeEnv(makeThrowingD1()), + project(), + byokPolicy, + logger, + ); + expect(selection.status).toBe("unavailable"); + }); + + it("refuses a provider name that is no longer in the allowlist", async () => { + // Reachable through the KV-cached policy: the operator removed the entry + // after the policy was parsed. The fallback is a refusal, never Workers AI. + const db = await withStoredKey(); + const selection = await resolveLlmProvider( + makeEnv(db, { LLM_PROVIDERS: undefined }), + project(), + byokPolicy, + logger, + ); + expect(selection.status).toBe("unavailable"); + if (selection.status !== "unavailable") return; + expect(selection.reason).toContain("not configured in LLM_PROVIDERS"); + }); + + it("never leaks the stored credential into the reason it returns", async () => { + const db = await withStoredKey("sk-ant-super-secret"); + const selection = await resolveLlmProvider( + makeEnv(db, { DEPLOY_SECRET_KEY: "rotated-master-key-aaaaaaaaaaaaaaa" }), + project(), + byokPolicy, + logger, + ); + if (selection.status !== "unavailable") throw new Error("expected unavailable"); + expect(selection.reason).not.toContain("sk-ant"); + }); +}); + +// --------------------------------------------------------------------------- +// 7.4 — the wiring: a credential failure NEVER reaches env.AI +// --------------------------------------------------------------------------- + +describe("buildEvaluators — BYOK never falls back to the operator's Workers AI", () => { + const KEY = "test-deploy-secret-key-0123456789"; + const byokPolicy: EvalPolicy = { evaluators: [{ type: "llm", provider: "anthropic" }] }; + + function envWith(db: D1Database, ai: AiBinding, overrides: Partial = {}): Env { + return { + DB: db, + AI: ai, + DEPLOY_SECRET_KEY: KEY, + LLM_PROVIDERS: PROVIDERS_JSON, + ...overrides, + } as unknown as Env; + } + + /** Runs only the `llm` evaluator, since that is the one BYOK changes. + * Cost samples ride on the eval RUN — `aggregate` does not carry them. */ + async function evaluateWith(env: Env, policy: EvalPolicy = byokPolicy) { + const built = await buildEvaluators(env, policy, project(), logger); + const llm = built.find((entry) => entry.type === "llm"); + expect(llm).toBeDefined(); + if (!llm) throw new Error("no llm evaluator"); + const { evalRuns, evalResult } = await runEvaluation( + [llm], + "diff --git a/x b/x", + policy, + logger, + ); + return { evalResult, costs: evalRuns[0]?.result.costs }; + } + + it.each([ + ["the secret is missing", {} as Partial], + ["DEPLOY_SECRET_KEY is unset", { DEPLOY_SECRET_KEY: undefined }], + ["the allowlist no longer has the provider", { LLM_PROVIDERS: undefined }], + ])("fails the gate closed and does not call env.AI when %s", async (_case, overrides) => { + const ai: AiBinding = { run: vi.fn(async () => ({ response: "{}" })) } as unknown as AiBinding; + const { db } = makeSqliteD1(); + resetLlmProviderCache(); + + const { evalResult, costs } = await evaluateWith(envWith(db, ai, overrides)); + + expect(evalResult.passed).toBe(false); + expect(evalResult.score).toBe(0); + expect(evalResult.reason).toContain("llm unavailable"); + // The whole point: the spend does not silently move back to the operator. + expect(ai.run).not.toHaveBeenCalled(); + // …and nothing was billed to anyone, because nothing ran. + expect(costs ?? []).toHaveLength(0); + }); + + it("marks the cost samples of a BYOK run as byok", async () => { + // Tasks 2/3 built `source` and the entire allowance design depends on it: a + // BYOK sample must not decrement the hosted allowance. + const { db } = makeSqliteD1(); + const stored = await putSecret( + db, + logger, + { DEPLOY_SECRET_KEY: KEY }, + { + projectId: project().id, + name: "ANTHROPIC_API_KEY", + value: "sk-ant-real", + actorId: "user_alice", + }, + ); + expect(stored.success).toBe(true); + + vi.stubGlobal( + "fetch", + vi.fn(async () => ({ + ok: true, + status: 200, + text: async () => + JSON.stringify({ + content: [{ type: "text", text: '{"score":1,"passed":true,"reason":"ok"}' }], + usage: { input_tokens: 900, output_tokens: 100 }, + }), + })), + ); + + const ai: AiBinding = { run: vi.fn() } as unknown as AiBinding; + const { evalResult, costs } = await evaluateWith(envWith(db, ai)); + + expect(evalResult.passed).toBe(true); + expect(costs).toEqual([{ kind: "llm_tokens", quantity: 1000, source: "byok" }]); + expect(ai.run).not.toHaveBeenCalled(); + }); + + it("leaves a Workers AI run unmarked, because absent already means platform", async () => { + const ai: AiBinding = { + run: vi.fn(async () => ({ response: '{"score":1,"passed":true,"reason":"ok"}' })), + } as unknown as AiBinding; + const { db } = makeSqliteD1(); + + const { costs } = await evaluateWith(envWith(db, ai), { evaluators: [{ type: "llm" }] }); + + expect(ai.run).toHaveBeenCalled(); + expect(costs?.[0]).not.toHaveProperty("source"); + expect(costs?.[0]).toMatchObject({ kind: "llm_tokens", estimated: true }); + }); + + it("keeps a provider error body — including an echoed API key — out of the result and the logs", async () => { + // Some providers echo the credential that failed. `llm-provider.ts` returns + // the status and never the body, and this is the assertion that keeps it + // that way: the reason is user-visible (`LLMEvaluator` interpolates the + // message into the EvalResult it fails closed with) and the logs are + // operator-visible. + const apiKey = "sk-ant-leaky-credential-value"; + const { db } = makeSqliteD1(); + await putSecret( + db, + logger, + { DEPLOY_SECRET_KEY: KEY }, + { projectId: project().id, name: "ANTHROPIC_API_KEY", value: apiKey, actorId: "user_alice" }, + ); + + vi.stubGlobal( + "fetch", + vi.fn(async () => ({ + ok: false, + status: 401, + text: async () => + JSON.stringify({ + error: { message: `invalid x-api-key: ${apiKey}`, type: "authentication_error" }, + }), + })), + ); + + const ai: AiBinding = { run: vi.fn() } as unknown as AiBinding; + const { evalResult } = await evaluateWith(envWith(db, ai)); + + expect(evalResult.passed).toBe(false); + expect(evalResult.reason).toContain("401"); + expect(evalResult.reason).not.toContain(apiKey); + expect(JSON.stringify(evalResult)).not.toContain(apiKey); + + const logged = JSON.stringify( + (["trace", "debug", "info", "warn", "error", "fatal"] as const).flatMap( + (level) => (logger[level] as unknown as { mock: { calls: unknown[] } }).mock.calls, + ), + ); + expect(logged).not.toContain(apiKey); + }); +}); + +// --------------------------------------------------------------------------- +// 7.5 — what is serialized into the prompt, and how big it is allowed to get +// --------------------------------------------------------------------------- + +describe("sanitizePolicy — the strip list, and MAX_POLICY_CONTEXT_CHARS", () => { + it("still strips the webhook secret it was written for", () => { + const policy: EvalPolicy = { + evaluators: [{ type: "webhook", url: "https://hook.test", secret: "shhh" }], + }; + expect(JSON.stringify(sanitizePolicy(policy))).not.toContain("shhh"); + expect(sanitizePolicy(policy).evaluators[0]).toEqual({ + type: "webhook", + url: "https://hook.test", + }); + }); + + it("strips the provider name and every credential-shaped field", () => { + // `sanitizeEvaluator` copies an entry of an unmodelled type through whole, + // so this is the only thing standing between such a field and the provider. + const policy = { + evaluators: [ + { type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }, + { + type: "future_evaluator", + apiKey: "sk-leak-1", + api_key: "sk-leak-2", + AUTHORIZATION: "Bearer sk-leak-3", + token: "sk-leak-4", + privateKey: "sk-leak-5", + url: "https://kept.example", + }, + ], + } as unknown as EvalPolicy; + + const serialized = JSON.stringify(sanitizePolicy(policy)); + for (let i = 1; i <= 5; i++) expect(serialized).not.toContain(`sk-leak-${i}`); + expect(serialized).not.toContain("anthropic"); + // Non-credential fields survive — the model still needs the policy. + expect(serialized).toContain("claude-sonnet-4-5"); + expect(serialized).toContain("https://kept.example"); + }); + + it("does not strip a field that merely contains a credential word", () => { + const policy = { + evaluators: [{ type: "diff", monkeys: 3, keystone: "kept" }], + } as unknown as EvalPolicy; + expect(JSON.stringify(sanitizePolicy(policy))).toContain("keystone"); + }); + + it("makes opting into BYOK free against the 8000-char policy context", () => { + // The interaction that would otherwise block merges on a project that works + // today: the serialized policy shares the model's input budget and an + // oversize one fails the gate CLOSED before any model call. Because + // `provider` is stripped, adding it costs zero characters. + const large: EvalPolicy = { + evaluators: [{ type: "diff", forbiddenPatterns: [`${"x".repeat(7_500)}`] }, { type: "llm" }], + }; + const withProvider: EvalPolicy = { + evaluators: [ + { type: "diff", forbiddenPatterns: [`${"x".repeat(7_500)}`] }, + { type: "llm", provider: "anthropic" }, + ], + }; + + const before = JSON.stringify(sanitizePolicy(large)).length; + const after = JSON.stringify(sanitizePolicy(withProvider)).length; + expect(after).toBe(before); + expect(after).toBeLessThan(8_000); + }); + + it("still fails the gate closed for a policy that is genuinely too large", async () => { + const { LLMEvaluator } = await import("../src/evaluation/llm-evaluator"); + const ai = { run: vi.fn() }; + const { WorkersAiProvider } = await import("../src/evaluation/llm-provider"); + const policy: EvalPolicy = { + evaluators: [{ type: "diff", forbiddenPatterns: [`${"x".repeat(9_000)}`] }, { type: "llm" }], + }; + + const result = await new LLMEvaluator( + new WorkersAiProvider(ai as unknown as AiBinding), + ).evaluate("diff", policy, logger); + + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.passed).toBe(false); + expect(result.data.reason).toContain("policy context is"); + expect(ai.run).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/llm-provider.test.ts b/tests/llm-provider.test.ts index 61b8322f..34eb6ef8 100644 --- a/tests/llm-provider.test.ts +++ b/tests/llm-provider.test.ts @@ -30,17 +30,26 @@ function jsonResponse(status: number, body: unknown) { return { ok: status >= 200 && status < 300, status, - json: async () => body, + text: async () => JSON.stringify(body), }; } -/** A response whose body is not JSON — `json()` rejects, as the real one does. */ +/** A response whose body arrived whole and is not JSON — an HTML error page. */ function nonJsonResponse(status = 200) { return { ok: status >= 200 && status < 300, status, - json: async () => { - throw new SyntaxError("Unexpected token < in JSON at position 0"); + text: async () => "502 Bad Gateway", + }; +} + +/** A response whose headers arrived but whose body transfer then died. */ +function bodyReadFailureResponse(status = 200) { + return { + ok: status >= 200 && status < 300, + status, + text: async () => { + throw new TypeError("network connection lost"); }, }; } @@ -56,10 +65,14 @@ function lastFetchCall() { beforeEach(() => { vi.restoreAllMocks(); + // `restoreAllMocks` does not undo `stubGlobal`: without this a `fetch` stub + // outlives its test, and a later test that forgot to stub one would pass on + // the previous test's provider rather than failing. + vi.unstubAllGlobals(); }); describe("WorkersAiProvider", () => { - it("returns the binding's response text and reports no usage", async () => { + it("returns the binding's response text, and no usage when it reported none", async () => { const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: "hello" }) }; const result = await new WorkersAiProvider(ai).run("@cf/model", [ { role: "user", content: "hi" }, @@ -67,10 +80,53 @@ describe("WorkersAiProvider", () => { expect(result.success).toBe(true); if (!result.success) return; expect(result.data.text).toBe("hello"); - // Absent usage is what keeps the Workers AI path on the char estimate. + // Absent usage is what puts this response on the char estimate. expect(result.data.usage).toBeUndefined(); }); + it("reads the token counts Workers AI does report", async () => { + const ai: AiBinding = { + run: vi.fn().mockResolvedValue({ + response: "hello", + usage: { prompt_tokens: 31, completion_tokens: 7 }, + }), + }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toEqual({ inputTokens: 31, outputTokens: 7 }); + }); + + it("a negative token count is treated as absent, never recorded as exact", async () => { + const ai: AiBinding = { + run: vi.fn().mockResolvedValue({ + response: "hello", + usage: { prompt_tokens: -1e12, completion_tokens: 7 }, + }), + }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a non-string `response` yields empty text rather than reaching the evaluator", async () => { + const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: 42 }) }; + const result = await new WorkersAiProvider(ai).run("@cf/model", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe(""); + }); + + it("a binding that never settles becomes a transport error, not a hung gate", async () => { + const ai: AiBinding = { run: vi.fn(() => new Promise(() => {})) }; + const result = await new WorkersAiProvider(ai, { timeoutMs: 5 }).run("@cf/model", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("did not respond within 5ms"); + }); + it("forwards the model and messages to the binding unchanged", async () => { const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: "{}" }) }; const messages = [ @@ -143,7 +199,11 @@ describe("AnthropicProvider — request shape", () => { expect(init.headers["anthropic-version"]).toBe("2023-06-01"); expect(init.headers["content-type"]).toBe("application/json"); expect(body.model).toBe("claude-sonnet-4-5"); - expect(typeof body.max_tokens).toBe("number"); + // The value, not just the type: Anthropic accepts `max_tokens: 0` and + // generates nothing for it, so "is a number" is not the assertion that + // catches a bound going wrong. + expect(body.max_tokens).toBe(1024); + expect(body.max_tokens).toBe(AnthropicProvider.MAX_TOKENS); // The system prompt is a top-level parameter; a system-role message is rejected. expect(body.system).toBe("be strict"); expect(body.messages).toEqual([{ role: "user", content: "review this" }]); @@ -288,6 +348,93 @@ describe("AnthropicProvider — response handling", () => { expect(result.data.usage).toBeUndefined(); }); + it("a negative token count is treated as absent, never billed as exact", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "x" }], + usage: { input_tokens: -1e12, output_tokens: 5 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a fractional token count is treated as absent", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "x" }], + usage: { input_tokens: 10.5, output_tokens: 5 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a verdict truncated at max_tokens fails with that reason, not a parse error", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: '{"score": 0.9, "passed": tr' }], + stop_reason: "max_tokens", + usage: { input_tokens: 100, output_tokens: 1024 }, + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("provider response was truncated at the 1024-token limit"); + }); + + it("a normal `stop_reason` is not mistaken for truncation", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: "verdict" }], + stop_reason: "end_turn", + }), + ), + ); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.text).toBe("verdict"); + }); + + it("a body that dies mid-transfer is a transport error, not a malformed body", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(bodyReadFailureResponse())); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + // The distinction is load-bearing: a response error fails the gate closed + // with a user-visible reason, a transport error surfaces to the caller. + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("network connection lost"); + }); + + it("keeps the parse failure as the response error's cause rather than dropping it", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(nonJsonResponse())); + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.cause).toBeInstanceOf(SyntaxError); + // ...and it stays off the message, which reaches the user. + expect(result.error.message).toBe("provider response was not valid JSON"); + }); + it("a network failure becomes a transport error, not a response error", async () => { vi.stubGlobal("fetch", vi.fn().mockRejectedValue(new TypeError("fetch failed"))); const result = await new AnthropicProvider("https://api.anthropic.com/v1", "k").run("m", []); @@ -345,6 +492,10 @@ describe("OpenAiCompatibleProvider — request shape", () => { expect(body.model).toBe("gpt-4o-mini"); // Unlike Anthropic, the system turn stays in the messages array. expect(body.messages).toEqual(messages); + // `max_completion_tokens`, never the deprecated `max_tokens`: OpenAI's + // reasoning models reject the latter outright. + expect(body.max_completion_tokens).toBe(1024); + expect(body).not.toHaveProperty("max_tokens"); }); it("works against a self-hosted base URL with a trailing slash", async () => { @@ -448,6 +599,56 @@ describe("OpenAiCompatibleProvider — response handling", () => { expect(result.error.message).toBe("provider response had no assistant message text"); }); + it("a negative token count is treated as absent, never billed as exact", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + choices: [{ message: { content: "x" } }], + usage: { prompt_tokens: -1e12, completion_tokens: 3 }, + }), + ), + ); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.usage).toBeUndefined(); + }); + + it("a verdict truncated at the token bound fails with that reason", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + choices: [{ message: { content: '{"score": 0.9, "pass' }, finish_reason: "length" }], + }), + ), + ); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("provider response was truncated at the 1024-token limit"); + }); + + it("a body that dies mid-transfer is a transport error, not a malformed body", async () => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(bodyReadFailureResponse())); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).not.toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toContain("network connection lost"); + }); + it("a response with no usage reports none, keeping the caller on an estimate", async () => { vi.stubGlobal( "fetch", @@ -522,7 +723,7 @@ describe("LLMEvaluator — cost accounting is per provider, not per source", () expect(result.data.costs?.[0]?.quantity).toBeGreaterThan(0); }); - it("Workers AI, which reports nothing, is on the estimate", async () => { + it("Workers AI without a usage object is on the estimate", async () => { const ai: AiBinding = { run: vi.fn().mockResolvedValue({ response: VERDICT }) }; const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); const result = await evaluator.evaluate("diff content", POLICY, mockLogger); @@ -530,6 +731,20 @@ describe("LLMEvaluator — cost accounting is per provider, not per source", () if (!result.success) return; expect(result.data.costs?.[0]?.estimated).toBe(true); }); + + it("Workers AI reporting usage is billed on exact counts, not the estimate", async () => { + const ai: AiBinding = { + run: vi.fn().mockResolvedValue({ + response: VERDICT, + usage: { prompt_tokens: 900, completion_tokens: 100 }, + }), + }; + const evaluator = new LLMEvaluator(new WorkersAiProvider(ai)); + const result = await evaluator.evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs).toEqual([{ kind: "llm_tokens", quantity: 1000 }]); + }); }); describe("LLMEvaluator — provider failures fail the gate closed", () => { diff --git a/tests/sandbox-evaluator.test.ts b/tests/sandbox-evaluator.test.ts index c589f15f..982072b4 100644 --- a/tests/sandbox-evaluator.test.ts +++ b/tests/sandbox-evaluator.test.ts @@ -797,7 +797,7 @@ describe("buildEvaluators — sandbox wiring", () => { createdAt: "2026-01-01T00:00:00.000Z", }; - function findSandbox(evaluators: ReturnType) { + function findSandbox(evaluators: Awaited>) { const entry = evaluators.find((e) => e.type === "sandbox"); expect(entry).toBeDefined(); // biome-ignore lint/style/noNonNullAssertion: asserted above @@ -805,7 +805,7 @@ describe("buildEvaluators — sandbox wiring", () => { } it("no SANDBOX binding → fails closed with an actionable wrangler.toml reason", async () => { - const evaluators = buildEvaluators({} as Env, sandboxPolicy, project, mockLogger, repo); + const evaluators = await buildEvaluators({} as Env, sandboxPolicy, project, mockLogger, repo); const result = await findSandbox(evaluators).evaluate("", sandboxPolicy, mockLogger); expect(result.success).toBe(true); if (!result.success) return; @@ -817,7 +817,7 @@ describe("buildEvaluators — sandbox wiring", () => { it("SANDBOX binding without workspace repo access → fails closed", async () => { const { binding } = makeMockSandbox({ exitCode: 0 }); - const evaluators = buildEvaluators( + const evaluators = await buildEvaluators( { SANDBOX: binding } as Env, sandboxPolicy, project, @@ -831,9 +831,9 @@ describe("buildEvaluators — sandbox wiring", () => { expect(binding.create).not.toHaveBeenCalled(); }); - it("SANDBOX binding + workspace repo access → a real SandboxEvaluator", () => { + it("SANDBOX binding + workspace repo access → a real SandboxEvaluator", async () => { const { binding } = makeMockSandbox({ exitCode: 0 }); - const evaluators = buildEvaluators( + const evaluators = await buildEvaluators( { SANDBOX: binding } as Env, sandboxPolicy, project, diff --git a/wrangler.toml b/wrangler.toml index 21c2b663..e09c2825 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -231,17 +231,30 @@ OAUTH_REDIRECT_URI = "http://localhost:8787/auth/github/callback" # `entitlementsConfigError` (src/middleware/config-guard.ts) logs exactly that. ENTITLEMENTS_ENFORCE = "0" # -# NOT YET IMPLEMENTED — reserved, and documented here so the shape is settled -# before anything reads it. Setting it today does nothing: no code parses it. +# The evaluator's provider allowlist, and the ONLY place a provider endpoint may +# be named. A project's .stratum/policy.yaml SELECTS one of these by name +# (`evaluators: [{type: llm, provider: anthropic}]`) and can never supply a +# `baseUrl`, so a project owner cannot aim the operator's Worker at a host of +# their choosing. An unrecognized name fails the policy closed and blocks that +# project's merges — a merge gate that cannot run must not pass — and never +# falls back to Workers AI, which would move the spend back onto this bill. # -# When Task 7 lands it becomes the evaluator's provider allowlist, and the ONLY -# place a provider endpoint may be named: a project's .stratum/policy.yaml will -# SELECT one of these by name and will never be able to supply a `baseUrl`, so a -# project owner cannot aim the operator's Worker at a host of their choosing. An -# unrecognized name will fail the policy closed — a merge gate that cannot run -# must not pass. Unset (the default) means Workers AI (`env.AI`) only, which is -# both today's behavior and the behavior until that work ships. Credentials are -# NOT here: they live per project in `project_secrets`. Intended shape: +# Parsed once per isolate by `parseLlmProviders` (src/evaluation/llm-providers.ts) +# and validated there, not per request: `kind` must be "anthropic" or +# "openai-compatible", `baseUrl` must be https with no embedded credentials, and +# a loopback, link-local or RFC1918 host is rejected. One bad entry rejects the +# whole value — every project selecting a provider then blocks, projects on +# Workers AI are unaffected, and `llmProvidersConfigError` logs the reason on the +# first request after the deploy. +# +# Unset (the default) means Workers AI (`env.AI`) only. Credentials are NOT here: +# each project stores its own key in `project_secrets` under the name derived +# from the provider — "anthropic" reads ANTHROPIC_API_KEY — encrypted under +# DEPLOY_SECRET_KEY. That coupling is deliberate (PRD Open Question 4: one key, +# not two) and it WIDENS that key's blast radius — rotating it already makes +# every deploy secret undecryptable, and now blocks the merge gate on every BYOK +# project until each re-enters its provider key. See the rotation warning in +# docs/user-guide/deployments.md. # LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' # Instance-wide telemetry kill switch. Not a secret — a plain var, listed here @@ -425,7 +438,10 @@ REFERRAL_SERVICE_URL = "https://referral.usestratum.dev" # Observe-only: the limits are measured for a period before they gate anything # (PRD Goal 3 — metering is always on, enforcement is a separate switch). ENTITLEMENTS_ENFORCE = "0" -# Operator-configured LLM providers; unset = Workers AI only. +# Operator-configured LLM provider allowlist; unset = Workers AI only, which is +# what production runs on today. A project opts in by naming one of these in its +# policy and storing its own key in that project's secrets. See the top-level +# [vars] block for the validation rules and the DEPLOY_SECRET_KEY coupling. # LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' # ============================================================================= @@ -584,4 +600,7 @@ REFERRAL_SERVICE_URL = "https://referral.usestratum.dev" # here — but only once the billing service is live and a period of observation # has produced real numbers. Off until then. ENTITLEMENTS_ENFORCE = "0" +# Operator-configured LLM provider allowlist; unset = Workers AI only. Staging is +# where a new provider entry gets exercised end-to-end before production, since a +# bad one is only visible as blocked merges on the projects selecting it. # LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' From 71b4242c8b700ac07293c6123c3f80836c5f06e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 12:06:59 +0000 Subject: [PATCH 13/26] Enforce allowances observe-only, and close the BYOK exfiltration paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 8, plus a security review of the BYOK work that found six real holes — two of which reopened fail-opens this series claimed to have closed. Enforcement charges the ACTOR, not the recorded owner, so an allowance follows the person: ten orgs do not help and being added to someone else's does not either. Pooling requires an org positively known to be paid, and that is the one check here that does not fail open — org entitlements are never warmed and the provider is a cached read that never fetches, so an org's plan is always a miss on first contact, and failing open there would leave every org unlimited forever. Everything else is inert unless both a billing service is configured and ENTITLEMENTS_ENFORCE is set, so a self-hoster cannot be refused anything. The security findings, worst first. The provider fetch followed redirects and re-validated nothing, so the parse-time allowlist bound the configured URL and not the destination. Any allowlisted host serving a 3xx sent the Worker somewhere nobody checked; on 307/308 the diff went with it, and x-api-key is not a header the Fetch spec strips cross-origin, so the project's own credential followed. The webhook evaluator in the same directory already refused to follow redirects for exactly this reason. A stored credential could reach a world-readable change page. putSecret validated byte length only, so a CR or LF in a value made Headers throw a TypeError embedding it, and that message flowed into the evaluation reason, the change record, and the GitHub report. Both ends are closed: control characters are rejected on the way in, and headers are built outside the try so a construction failure maps to a constant. sanitizePolicy stripped one level, and only inside evaluator entries — while parsePolicyContent spreads unknown root keys onto the policy. So a credential at the top of the policy file reached both the model prompt and the policy-supplied webhook URL, which is an attacker-chosen host. The strip is now recursive over the whole policy with depth and size bounds, and the word list covers the plurals and `auth` it was missing. Two silent fail-opens: with two llm entries, BYOK resolution read only the first while an evaluator was built for each, so a policy naming a provider ran on the operator's account with no error and no log. And the shipped LLM_PROVIDERS example omitted /v1, so anyone uncommenting the documented line got a 404 on every BYOK call — the test that should have caught it ignored the request URL entirely. The host filter was also weaker than the one already in the tree — but not in the way the review assumed. Reusing validateWebhookUrl would have fixed two gaps and left two, because the webhook path had them too: [::127.0.0.1] normalizes to [::7f00:1], which its digit-stripping test read as "7f1" rather than loopback, and fe80::/10 was matched as a literal prefix so feb0:: passed. The two filters are now one, and it classifies bracketed literals by mask instead of by string shape — which closes a pre-existing SSRF gap on the webhook evaluator as well. Typecheck clean; lint clean; 4060 tests pass, up 66. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- CHANGELOG.md | 46 ++ docs/user-guide/getting-started.md | 9 +- src/billing/enforcement.ts | 481 ++++++++++++ src/billing/entitlements.ts | 57 +- src/billing/usage-notifications.ts | 239 ++++++ src/deploy/runner.ts | 48 ++ src/email/templates.ts | 78 ++ src/evaluation/llm-byok.ts | 40 +- src/evaluation/llm-evaluator.ts | 197 ++++- src/evaluation/llm-provider.ts | 61 +- src/evaluation/llm-providers.ts | 108 ++- src/evaluation/policy-loader.ts | 105 ++- src/evaluation/sanitize-policy.ts | 115 ++- src/evaluation/types.ts | 15 + src/merge/post-merge.ts | 5 + src/middleware/rate-limit.ts | 53 +- src/queue/usage-meter.ts | 14 + src/routes/changes.ts | 22 +- src/routes/deployments.ts | 10 +- src/routes/projects.ts | 35 + src/routes/sync-management.ts | 15 +- src/services/change-flow.ts | 61 +- src/storage/costs.ts | 54 +- src/storage/project-secrets.ts | 63 +- src/storage/usage.ts | 101 ++- src/types.ts | 6 +- src/utils/validation.ts | 175 ++++- tests/enforcement.test.ts | 683 ++++++++++++++++++ tests/entitlements.test.ts | 13 +- tests/evaluator-billing-context.test.ts | 28 +- tests/llm-byok.test.ts | 277 ++++++- tests/llm-provider.test.ts | 139 +++- tests/project-secrets.test.ts | 34 + .../content/docs/guides/getting-started.md | 9 +- wrangler.toml | 17 +- 35 files changed, 3187 insertions(+), 226 deletions(-) create mode 100644 src/billing/enforcement.ts create mode 100644 src/billing/usage-notifications.ts create mode 100644 tests/enforcement.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 0da1e909..de000d59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,6 +67,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **A policy file is bounded, and an unusable `llm` entry says why.** + `evaluators:` now takes at most 16 entries, the cap `deploys:` has always had + and for the same reason — one merge must not amplify into an unbounded number + of external calls. An entry that names a provider must also name a `model`, + since the default is a Workers AI model id that fails every call against any + other provider. The merge-blocking reason now names the provider (or the field) + that made the entry unusable instead of only counting entries. A verdict + truncated at the provider's token cap records the tokens the provider reported + for it, rather than recording a charged call as free. +- **The commented `LLM_PROVIDERS` example carries the API version path.** The + provider appends `/messages` (or `/chat/completions`) to `baseUrl`, so the + `https://api.anthropic.com` shown in `wrangler.toml` would have 404'd on every + call; it is `https://api.anthropic.com/v1` now. A `baseUrl` is also stored + normalized, so a trailing `#` or `?` — both of which pass the "no fragment, no + query" check — cannot swallow the path the provider appends. - **`api_request` reported the wrong route for most pages.** `routePath(c, -1)` returns the last *registered* route matching a path, not the one that answered it, and `uiRouter`'s `/:namespace/:slug` catch-all is mounted last — @@ -77,6 +92,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 events are not corrected. ### Security +- **A project secret can no longer escape through an error message.** A stored + value containing CR, LF or NUL made the outbound provider request's header + construction throw a `TypeError` that quotes the offending value; that message + became an evaluation reason, which is persisted on the change and rendered — + world-readably, for a public project. Both ends are closed: the store rejects + control characters in a value, and the provider builds its headers where a + failure maps to a constant that interpolates nothing. +- **The LLM provider request no longer follows redirects.** The allowlist + validates the host a request is *sent* to; a 3xx from an allowlisted host used + to move the prompt (the diff and the policy) to an unvalidated one, re-sending + the body on 307/308 with the project's `x-api-key` attached, which the Fetch + spec does not strip cross-origin. A redirect is now a failed evaluation. +- **Credential stripping covers the whole policy, recursively.** Fields whose + name says "credential" were removed only from the top level of each evaluator + entry, so one written at the top of `.stratum/policy.yaml`, or nested inside an + entry of an unmodelled type, reached the review model and the body POSTed to a + policy-supplied webhook URL. The word list also missed common names — + `auth`, plurals such as `tokens` and `apiKeys`, `pwd`, `pat`, `hmac`, `sig` — + while keeping innocuous ones like `keystone`. +- **One host filter, not two.** The LLM provider allowlist had its own copy of + the private-address check, and it had drifted: CGNAT (`100.64/10`), the + `.internal`/`.local` suffixes (GCP's metadata endpoint is + `metadata.google.internal`), everything in `fe80::/10` not spelled `fe80`, and + IPv4-compatible IPv6 (`[::127.0.0.1]`, which the URL parser rewrites to + `[::7f00:1]`) all got through. Both callers now share one filter, which + expands an IPv6 literal rather than matching its spelling. +- **A policy declaring two `llm` entries is refused.** The BYOK provider was + resolved from the first entry while an evaluator was built for every entry, so + `[{llm}, {llm, provider: …}]` silently ran twice on the operator's Workers AI + bill — the fail-open this work exists to prevent. Both the parser and the + provider resolution now refuse rather than choosing one. - **GitHub sign-in no longer trusts an unverified email.** The callback picked the primary address, verified or not, and fell back to *any* address on the account. That address then matched (and linked to) an existing Stratum diff --git a/docs/user-guide/getting-started.md b/docs/user-guide/getting-started.md index 3dbda34c..868bbb1c 100644 --- a/docs/user-guide/getting-started.md +++ b/docs/user-guide/getting-started.md @@ -194,7 +194,14 @@ downgrade your governance. minimum passing score (0.0–1.0); `maxDiffChars` bounds how much diff is sent (default 24,000, max 100,000). An unavailable model or unparseable verdict fails closed. Token usage is recorded on the change as a cost - record. + record. **A policy may declare at most one `llm` entry** — a second one is a + merge-blocking policy error, because the two entries cannot both be the + configuration in force. + +A policy may declare **at most 16 `evaluators:` entries**; more than that is a +policy error and blocks merges, the same as any other unusable entry. (The +`deploys:` list has the same cap, for the same reason: one merge must not turn +into an unbounded number of external calls.) Two top-level knobs sit alongside `evaluators`: `requireAll` (default `true`) makes the aggregate verdict demand every evaluator pass — set it to `false` to diff --git a/src/billing/enforcement.ts b/src/billing/enforcement.ts new file mode 100644 index 00000000..db8db76e --- /dev/null +++ b/src/billing/enforcement.ts @@ -0,0 +1,481 @@ +/** + * Consulting entitlements: the one place a limit turns into a decision. + * + * `entitlements.ts` defines and fetches allowances and enforces nothing. This + * module is the other half — it evaluates a decision, records it, and hands the + * call site an answer. Every metered call site goes through {@link checkMeter} + * or {@link checkGauge} rather than reading a limit itself, so "what does the + * plan allow", "what has been used", "who is this checked against" and "does + * the answer bind" are decided once instead of five times. + * + * ## Everything here ships observe-only + * + * A decision is always *evaluated* and *recorded*; it is *binding* only when + * `ENTITLEMENTS_ENFORCE=1`. Under observe-only a refusal still reserves, still + * logs, and still admits — that is the month of measurement PRD §8 asks for + * before anything blocks. `admitted` is what a caller acts on; `refused` is what + * the decision was, and the two differ exactly when enforcement is off. + * + * ## Inert with `BILLING_SERVICE_URL` unset + * + * The first line of every entry point is `entitlementsEnabled`. With the billing + * vars unset — every self-hoster — nothing here reads entitlements, nothing + * touches the `USAGE_METER` binding, and no refusal can be produced. + * `tests/enforcement.test.ts` asserts the binding is never touched in that + * configuration, because "inert by default" is a claim that decays silently. + * + * ## The subject (PRD §4a) + * + * The subject a limit is CHECKED against is the acting user, not the subject the + * spend is RECORDED against. See {@link resolveEnforcementSubject}. + */ +import { + type RateMeterKey, + type UsageMeterKey, + rateWindowSeconds, + usageMeterName, +} from "../queue/usage-meter"; +import { resolveBillingSubject } from "../storage/costs"; +import { type MeterKey, getOwnerMeterTotals, usagePeriod } from "../storage/usage"; +import type { Env } from "../types"; +import type { Logger } from "../utils/logger"; +import { + type CountKey, + type Entitlements, + RemoteEntitlements, + UNLIMITED, + UnlimitedEntitlements, + entitlementsEnabled, + warmEntitlements, +} from "./entitlements"; + +/** Who a limit is checked against. Not necessarily who the spend is billed to. */ +export interface EnforcementSubject { + ownerId: string; + ownerType: "user" | "org"; + /** + * True only when this subject is an org that is *positively known* to pool + * its members' usage on a paid plan. Never inferred; see + * {@link resolveEnforcementSubject}. + */ + pooled: boolean; +} + +/** What one check decided, and whether that decision binds. */ +export interface EnforcementDecision { + /** + * Whether the caller may proceed. Under observe-only this is `true` even for + * a refusal — gate on this, and report {@link EnforcementDecision.refused}. + */ + admitted: boolean; + /** What the limit says, regardless of whether it binds. */ + refused: boolean; + /** Whether `ENTITLEMENTS_ENFORCE=1` made this decision binding. */ + enforcing: boolean; + /** Whether a limit was consulted at all (false when the seam is inert). */ + checked: boolean; + /** Whether a reservation was taken and therefore needs settling. */ + reserved: boolean; + /** The limit consulted: `-1` unlimited, `0` a hard block, else a ceiling. */ + limit: number; + /** The subject's counter after this check, as the meter reported it. */ + count: number; + /** + * Present when refused: the user- and agent-facing copy naming what is + * exhausted, when it resets and both remedies (PRD §4c). Never a bare error + * string — this is the highest-intent moment in the funnel. + */ + reason?: string; +} + +/** Every entry point short-circuits to this when the seam is off. */ +function inert(): EnforcementDecision { + return { + admitted: true, + refused: false, + enforcing: false, + checked: false, + reserved: false, + limit: UNLIMITED, + count: 0, + }; +} + +/** + * Whether a decision, once evaluated, is allowed to refuse anything. + * + * Two switches, not one (PRD Goal 3): metering is always on, enforcement is + * opt-in on top of a configured billing service. `ENTITLEMENTS_ENFORCE=1` with + * no service is a misconfiguration `entitlementsConfigError` reports; here it + * simply cannot bind, because there is no limit to bind to. + */ +export function enforcementBinding(env: Env): boolean { + return entitlementsEnabled(env) && env.ENTITLEMENTS_ENFORCE === "1"; +} + +/** + * The subject a limit is checked against (PRD §4a). + * + * **The acting user, not the project's owner.** A project in an org namespace is + * recorded against the org, and an org costs nothing to create and needs only + * *write* access to put a project in — so an allowance keyed on the recorded + * owner is one an exhausted user resets by making an org, or by being added to + * someone else's. Keying the check on the person who ran the evaluation makes + * the allowance follow the person: ten orgs do not help, and nobody can drain an + * allowance that is not theirs. + * + * **An org pools only when positively known to be paid**, which is the one place + * this PRD does not fail open, and the direction matters. `forOwner` is a cached + * read that never fetches and nothing warms an org's entitlements before the + * first project of that org is evaluated, so an org's plan is ALWAYS a miss on + * first contact. Reading "unknown" as "own subject" would therefore make every + * org its own permanently-unlimited allowance — the hole, reopened by the + * safety mechanism. Unknown or uncached means free, and the actor is charged. + * The cost of that direction is that during a billing outage a paying org's + * members fall back to their personal limits: a degradation, and recoverable, + * where the alternative is permanent. + * + * The org's entitlements are warmed here via `waitUntil` — the auth middleware + * knows the caller but not the project, so it cannot do it — which is why the + * first check for a given org resolves to the actor and a later one can pool. + * + * @param opts.actorUserId - The user who ran the operation, when the site knows one + * @param opts.owner - The project's recorded owner (`agent` walks to its user) + * @param opts.waitUntil - Used to warm an org's plan off the response path + * @returns The subject, or `null` when nobody can be named + */ +export async function resolveEnforcementSubject( + env: Env, + logger: Logger, + opts: { + actorUserId?: string; + owner?: { ownerId?: string; ownerType?: "user" | "org" | "agent" }; + waitUntil?: (promise: Promise) => void; + }, +): Promise { + const { actorUserId, owner, waitUntil } = opts; + + if (owner?.ownerType === "org" && owner.ownerId) { + const orgId = owner.ownerId; + if (entitlementsEnabled(env)) { + const resolved = await new RemoteEntitlements(env, logger).forOwner(orgId, "org"); + // `source: "default"` is "we have never heard about this org", not "this + // org is on the default plan" — pooling on it would pool on ignorance. + const pooled = + resolved.success && resolved.data.source !== "default" && resolved.data.entitlements.pooled; + if (pooled) { + return { ownerId: orgId, ownerType: "org", pooled: true }; + } + // Warmed for NEXT time, never awaited: a paid org pools from its second + // evaluation onward, and the first one is charged to the actor. + warmEntitlements(env, waitUntil, { ownerId: orgId, ownerType: "org" }, logger); + } + } + + if (actorUserId) return { ownerId: actorUserId, ownerType: "user", pooled: false }; + + // No actor: a queue consumer, or a path that does not carry one. The recorded + // subject is then the only nameable one — weaker than §4a wants for an + // org-owned project, and still strictly better than no check at all, which is + // the only other option here. `resolveBillingSubject` also performs the + // agent -> owner walk, so an agent-owned project resolves to a person. + if (!owner?.ownerId) return null; + const recorded = await resolveBillingSubject(env.DB, logger, owner); + if (!recorded) return null; + return { ...recorded, pooled: recorded.ownerType === "org" }; +} + +/** + * The subject's allowances, from the cached read only. + * + * Never fetches: this runs on the request path, where a cold billing service + * would put a five-second dependency in front of a merge. A miss resolves to + * `UnlimitedEntitlements`, which admits — the fail-open rule of the module it + * reads from. + */ +async function entitlementsFor( + env: Env, + logger: Logger, + subject: EnforcementSubject, +): Promise { + const resolved = await new RemoteEntitlements(env, logger).forOwner( + subject.ownerId, + subject.ownerType, + ); + return resolved.success ? resolved.data.entitlements : UnlimitedEntitlements; +} + +/** + * Subjects whose D1 aggregate has already been reconciled into their counter in + * THIS isolate, keyed `subjectType:subjectId:period`. + * + * The counter is authoritative for enforcement but it is not the ledger: a + * subject that spent before the meter existed, or whose object was purged, or + * whose spend was recorded by a path that never reserved, has a `usage_periods` + * total the object has never seen. `setFloor` raises the counter to it (never + * lowers — see the object's doc comment), and doing that once per isolate keeps + * the D1 read off every single evaluation. + */ +const reconciled = new Set(); + +/** + * Raise the subject's counters to the D1 aggregate, at most once per isolate. + * + * `source = "platform"` is not optional here and not a detail: migration 049 + * keys usage on `(owner_id, period, meter, source)` exactly so a project paying + * its own provider bill is not charged against the hosted allowance, and a sum + * that omitted the filter would hand that distinction straight back (Goal 5). + */ +async function reconcileFloor( + env: Env, + logger: Logger, + meterNamespace: NonNullable, + subject: EnforcementSubject, + period: string, +): Promise { + const key = `${subject.ownerType}:${subject.ownerId}:${period}`; + if (reconciled.has(key)) return; + reconciled.add(key); + if (!env.DB) return; + const totals = await getOwnerMeterTotals(env.DB, logger, subject.ownerId, period, "platform"); + if (!totals.success) return; + const stub = meterNamespace.get( + meterNamespace.idFromName(usageMeterName(subject.ownerType, subject.ownerId)), + ); + for (const [meter, quantity] of Object.entries(totals.data)) { + if (typeof quantity === "number") await stub.setFloor(meter as MeterKey, quantity, period); + } +} + +export interface MeterCheckOptions { + subject: EnforcementSubject; + meter: UsageMeterKey; + /** Upper bound on what this operation will consume. Settled afterwards. */ + estimate: number; + /** Caller's clock, so one operation's counters cannot straddle a boundary. */ + nowMs?: number; + /** 'YYYY-MM' UTC; defaults to the month `nowMs` falls in. */ + period?: string; + /** What is being gated, for the refusal copy (e.g. "AI review"). */ + what: string; +} + +/** + * Evaluate one metered limit for one subject: reserve an upper bound, record the + * decision, and admit unless enforcement is on and the bound does not fit. + * + * The reservation is the point. An LLM call's true cost is known only from its + * response, so checking a counter and then spending would be check-then-act + * against a shared total; `reserve` admits or refuses the whole bound + * atomically, and {@link settleMeter} hands back the difference. + * + * A refusal takes no reservation (the meter does not count what it refuses), so + * only an admitted decision reports `reserved`. + */ +export async function checkMeter( + env: Env, + logger: Logger, + opts: MeterCheckOptions, +): Promise { + if (!entitlementsEnabled(env)) return inert(); + const meterNamespace = env.USAGE_METER; + if (!meterNamespace) { + // Configured to enforce with nothing to count against. Logged rather than + // refusing: `entitlementsConfigError` is where this gets fixed, and a + // missing binding must not become a merge outage. + logger.warn("Entitlement check skipped: USAGE_METER binding is not configured", { + meter: opts.meter, + }); + return inert(); + } + + const nowMs = opts.nowMs ?? Date.now(); + const period = opts.period ?? usagePeriod(new Date(nowMs)); + const entitlements = await entitlementsFor(env, logger, opts.subject); + const limit = limitFor(entitlements, opts.meter); + + if (isMonthlyMeter(opts.meter)) { + await reconcileFloor(env, logger, meterNamespace, opts.subject, period); + } + + const stub = meterNamespace.get( + meterNamespace.idFromName(usageMeterName(opts.subject.ownerType, opts.subject.ownerId)), + ); + const outcome = await stub.reserve(opts.meter, opts.estimate, limit, period, nowMs); + const enforcing = enforcementBinding(env); + const refused = !outcome.admitted; + + const decision: EnforcementDecision = { + admitted: outcome.admitted || !enforcing, + refused, + enforcing, + checked: true, + reserved: outcome.admitted, + limit, + count: outcome.count, + }; + if (refused) { + decision.reason = meterRefusalReason(opts.what, opts.meter, limit, period, nowMs); + } + record(logger, opts.subject, opts.meter, decision, entitlements.plan); + return decision; +} + +/** + * Apply the true cost of a spend a {@link checkMeter} reservation covered. + * + * Call it in a `finally`. A provider call that throws has still consumed the + * reservation, and leaving it standing would charge the subject the upper bound + * for a call that never produced a token — a few of those and the allowance is + * gone with nothing to show for it. + */ +export async function settleMeter( + env: Env, + // Kept for signature symmetry with `checkMeter`, which every call site pairs + // it with; settling reports nothing of its own. + _logger: Logger, + opts: { + subject: EnforcementSubject; + meter: UsageMeterKey; + /** Signed difference between the true cost and the reserved bound. */ + delta: number; + nowMs?: number; + period?: string; + }, +): Promise { + if (!entitlementsEnabled(env)) return; + const meterNamespace = env.USAGE_METER; + if (!meterNamespace) return; + if (!Number.isFinite(opts.delta) || opts.delta === 0) return; + const nowMs = opts.nowMs ?? Date.now(); + const period = opts.period ?? usagePeriod(new Date(nowMs)); + const stub = meterNamespace.get( + meterNamespace.idFromName(usageMeterName(opts.subject.ownerType, opts.subject.ownerId)), + ); + await stub.settle(opts.meter, opts.delta, period, nowMs); +} + +/** + * Evaluate a counted gauge — a thing that exists rather than a thing that is + * spent, so there is nothing to reserve and nothing to settle. + * + * @param opts.current - How many the subject already has, counted by the caller + */ +export async function checkGauge( + env: Env, + logger: Logger, + opts: { subject: EnforcementSubject; count: CountKey; current: number; what: string }, +): Promise { + if (!entitlementsEnabled(env)) return inert(); + const entitlements = await entitlementsFor(env, logger, opts.subject); + const limit = entitlements.counts[opts.count]; + const enforcing = enforcementBinding(env); + // Same reading of a limit as the meter: `-1` unlimited, `0` a hard block, + // anything not a usable limit admits rather than blocking on a bad payload. + const usable = Number.isInteger(limit) && limit >= 0; + const refused = usable && opts.current + 1 > limit; + + const decision: EnforcementDecision = { + admitted: !refused || !enforcing, + refused, + enforcing, + checked: true, + reserved: false, + limit, + count: opts.current, + }; + if (refused) decision.reason = gaugeRefusalReason(opts.what, limit); + record(logger, opts.subject, opts.count, decision, entitlements.plan); + return decision; +} + +/** + * Every decision is recorded, admitted or not. + * + * This is what makes observe-only worth a month: a log line per decision, with + * the plan and whether the decision bound, is the difference between measuring + * limits and guessing at them. + */ +function record( + logger: Logger, + subject: EnforcementSubject, + meter: string, + decision: EnforcementDecision, + plan: string, +): void { + logger.info("Entitlement decision", { + subjectId: subject.ownerId, + subjectType: subject.ownerType, + pooled: subject.pooled, + plan, + meter, + limit: decision.limit, + count: decision.count, + refused: decision.refused, + enforcing: decision.enforcing, + admitted: decision.admitted, + }); +} + +function isMonthlyMeter(meter: UsageMeterKey): meter is MeterKey { + return meter.endsWith("_month"); +} + +function limitFor(entitlements: Entitlements, meter: UsageMeterKey): number { + return isMonthlyMeter(meter) + ? entitlements.meters[meter] + : entitlements.rates[meter as RateMeterKey]; +} + +/** First instant of the month after `period`, ISO 8601 — when a monthly meter frees up. */ +export function periodResetsAt(period: string): string { + const match = /^(\d{4})-(\d{2})$/.exec(period); + if (!match) return "the start of next month"; + const year = Number(match[1]); + const month = Number(match[2]); + return new Date(Date.UTC(year, month, 1)).toISOString(); +} + +/** + * The copy a blocked caller reads (PRD §4c, §8). + * + * It is product copy, not an error message: this is the moment somebody learns + * a limit exists, and an agent over `/mcp` gets only this string. So it names + * what is exhausted, when it resets, and BOTH ways out — and for the rate meter + * it says plainly that a project's own key is not one of them, because that is + * the remedy an agent would otherwise reach for first. + */ +function meterRefusalReason( + what: string, + meter: UsageMeterKey, + limit: number, + period: string, + nowMs: number, +): string { + if (meter === "evaluations_per_hour") { + const resetsAt = new Date(nowMs + rateWindowSeconds(meter) * 1000).toISOString(); + return `${what} could not run: this account has used its hourly evaluation allowance (${limit} per hour). Capacity frees up as the window rolls, and in full by ${resetsAt}. Two ways forward: wait for the window, or raise this account's plan limits. Bringing your own provider key does NOT lift this one — it bounds how often an evaluation runs, not whose tokens it spends.`; + } + const remedy = + meter === "llm_tokens_month" + ? "set `provider:` on the llm evaluator in .stratum/policy.yaml to run reviews on this project's own API key" + : "reduce this month's usage"; + return `${what} could not run: this account's ${meterLabel(meter)} for ${period} is used up (limit ${limit}). It resets at ${periodResetsAt(period)}. Two ways forward: ${remedy}, or raise this account's plan limits.`; +} + +function gaugeRefusalReason(what: string, limit: number): string { + return `${what} is not allowed: this account's plan allows ${limit} of them. Two ways forward: remove or make public one you already have, or raise this account's plan limits.`; +} + +function meterLabel(meter: UsageMeterKey): string { + switch (meter) { + case "llm_tokens_month": + return "monthly LLM token allowance"; + case "sandbox_ms_month": + return "monthly sandbox time allowance"; + case "deploys_month": + return "monthly deploy allowance"; + default: + return meter; + } +} diff --git a/src/billing/entitlements.ts b/src/billing/entitlements.ts index 851ae300..d237f1a1 100644 --- a/src/billing/entitlements.ts +++ b/src/billing/entitlements.ts @@ -10,7 +10,8 @@ * Three key sets, deliberately not one union, because they are not one kind of * thing and do not live in one place: a `MeterKey` is a monthly *flow* accumulated * in `usage_periods`, a `CountKey` is a *gauge* counted on demand, and a `RateKey` - * is a per-minute *rate* enforced out of KV. Collapsing them invites a limit to be + * is a *rate* over a short window, enforced out of KV or off the usage meter but + * never out of the monthly aggregate. Collapsing them invites a limit to be * looked up from the wrong store. * * ## Fail OPEN — deliberately the opposite of the beta gate @@ -48,8 +49,16 @@ import { type Result, err, ok } from "../utils/result"; /** A gauge counted on demand, not a monthly flow — there is no `private_projects` meter. */ export type CountKey = "private_projects"; -/** A rate, enforced from KV by the request limiter rather than from a spend counter. */ -export type RateKey = "requests_per_minute"; +/** + * A bound on how often something may happen, rather than on how much it spends. + * + * The two members are enforced in different places on purpose, and the split is + * the one PRD §4b argues: `requests_per_minute` stays on the KV request limiter, + * where the bucket rolls as fast as KV goes stale, and `evaluations_per_hour` + * lives on the `UsageMeter` Durable Object, because an hour-long window read + * through a per-colo cache that is a minute stale is not a limit at all. + */ +export type RateKey = "requests_per_minute" | "evaluations_per_hour"; /** Every meter a limit may be set on. Mirrors `MeterKey`; see `parseEntitlements`. */ export const METER_KEYS: readonly MeterKey[] = [ @@ -60,7 +69,7 @@ export const METER_KEYS: readonly MeterKey[] = [ export const COUNT_KEYS: readonly CountKey[] = ["private_projects"]; -export const RATE_KEYS: readonly RateKey[] = ["requests_per_minute"]; +export const RATE_KEYS: readonly RateKey[] = ["requests_per_minute", "evaluations_per_hour"]; /** * One billing subject's allowances. @@ -72,6 +81,24 @@ export const RATE_KEYS: readonly RateKey[] = ["requests_per_minute"]; export interface Entitlements { /** Opaque plan identifier as the billing service names it; for display and logs. */ plan: string; + /** + * Whether an ORG on this plan pools its members' usage onto the org itself + * (PRD §4a). + * + * Explicit rather than derived from {@link Entitlements.plan}, which is an + * opaque display string: pooling decides WHOSE allowance a spend is checked + * against, and inferring that from a name the billing service is free to + * change is how a rename becomes a free-usage hole. + * + * It defaults to `false`, and that default is load bearing — it is the one + * place this file does not fail open. `forOwner` never fetches and org + * entitlements are never warmed by the auth path, so an org's plan is always + * a cache miss on FIRST contact; reading an unknown plan as "pooled" would + * make every org its own unlimited subject permanently, which is exactly the + * hole the enforcement subject exists to close. Unknown therefore means free, + * and the acting user is charged instead. + */ + pooled: boolean; meters: Record; counts: Record; rates: Record; @@ -125,7 +152,11 @@ export const UnlimitedEntitlements: Entitlements = Object.freeze({ deploys_month: UNLIMITED, }), counts: Object.freeze({ private_projects: UNLIMITED }), - rates: Object.freeze({ requests_per_minute: UNLIMITED }), + rates: Object.freeze({ requests_per_minute: UNLIMITED, evaluations_per_hour: UNLIMITED }), + // Never pooled by default: "we do not know this org's plan" must not read as + // "this org pools", or an unknown org becomes a fresh allowance. See the + // field's own doc comment. + pooled: false, }); /** @@ -156,8 +187,17 @@ const CACHE_TTL_SECONDS = 300; */ const FAILURE_TTL_SECONDS = 60; -/** Versioned so a shape change cannot be read back as the old shape. */ -const KEY_PREFIX = "entitlements:v1:"; +/** + * Versioned so a shape change cannot be read back as the old shape. + * + * Bumped to v2 by the enforcement work: `pooled` and `evaluations_per_hour` + * joined the payload, and a v1 entry read back through the new parser would + * silently take the defaults for both — which for `pooled` is the safe answer + * and for `evaluations_per_hour` is "unlimited". Neither is a value anybody + * chose, so the old entries are abandoned rather than reinterpreted. The + * billing service's response contract changes with it (private cloud layer). + */ +const KEY_PREFIX = "entitlements:v2:"; function cacheKey(ownerType: string, ownerId: string): string { return `${KEY_PREFIX}${ownerType}:${ownerId}`; @@ -221,6 +261,9 @@ function parseEntitlements(payload: unknown, fallback: Entitlements): Entitlemen const plan = typeof record.plan === "string" && record.plan.trim() ? record.plan.trim() : null; return { plan: plan ?? fallback.plan, + // Only a real boolean pools. A truthy string or a 1 is a payload we do not + // understand, and the safe reading of that is "not positively known". + pooled: typeof record.pooled === "boolean" ? record.pooled : fallback.pooled, meters: readGroup(record.meters, METER_KEYS, fallback.meters), counts: readGroup(record.counts, COUNT_KEYS, fallback.counts), rates: readGroup(record.rates, RATE_KEYS, fallback.rates), diff --git a/src/billing/usage-notifications.ts b/src/billing/usage-notifications.ts new file mode 100644 index 00000000..f4b8bff8 --- /dev/null +++ b/src/billing/usage-notifications.ts @@ -0,0 +1,239 @@ +/** + * "You are approaching a limit" — the missing middle between a usage page + * nobody visits and a merge that suddenly blocks (PRD §8, Goal 10). + * + * ## Edge-triggered, not state-triggered + * + * The notice fires on the CROSSING of 80%, which is why `upsertUsage` returns + * the post-write total: the totals either side of one write are what make "just + * went past" distinguishable from "is past", without a second read and without + * a job that scans everyone. + * + * ## The receipt is in KV, and the limit is part of its key + * + * A column on `usage_periods` would be at the wrong grain (per `source`) and on + * the wrong subject (the recorded one), and a new table would pull in + * `tests/backup-coverage.test.ts`, which requires every migration-created table + * to be classified — a notification receipt is neither billing data nor worth + * restoring. The limit is IN the key on purpose: raising someone's limit + * mid-period drops them back under 80%, and a receipt keyed without the limit + * would silence the warning for the second crossing, which is the case the + * dedupe was supposed to make safe. + * + * ## Not on the events queue + * + * `events.project` is NOT NULL and a usage threshold is owner-scoped, and every + * domain event fans out through `webhookHandler` — which would publish a + * customer's billing state to whatever third-party endpoint a project happens to + * have configured. If this ever goes on the queue it must be an internal event + * type excluded from fan-out. + * + * Gated on `entitlementsEnabled` throughout: a self-hoster must never be nagged + * about a limit they do not have. + */ +import { getUsageThresholdEmail } from "../email/templates"; +import type { BillingSubject } from "../storage/costs"; +import type { UsageWriteTotal } from "../storage/usage"; +import { getUser } from "../storage/users"; +import type { Env } from "../types"; +import type { Logger } from "../utils/logger"; +import { periodResetsAt, resolveEnforcementSubject } from "./enforcement"; +import { RemoteEntitlements, UnlimitedEntitlements, entitlementsEnabled } from "./entitlements"; + +/** Fraction of a limit that earns a warning. One threshold, so one receipt shape. */ +export const USAGE_NOTICE_THRESHOLD = 0.8; + +/** Percent form, used in the receipt key and the copy. */ +const THRESHOLD_PERCENT = Math.round(USAGE_NOTICE_THRESHOLD * 100); + +/** Versioned like the entitlements cache, so a key-shape change cannot be misread. */ +const RECEIPT_PREFIX = "usage-notice:v1:"; + +/** + * Long enough to outlive the period the crossing belongs to (a month plus the + * grace the meter itself uses), short enough that receipts do not accumulate + * forever. KV expires them; nothing sweeps. + */ +const RECEIPT_TTL_SECONDS = 40 * 24 * 60 * 60; + +/** What a recording site hands over so a crossing can be noticed. */ +export interface UsageNoticeInput { + /** Whose `usage_periods` rows were just written. */ + recorded: BillingSubject; + /** The user who ran the operation, when the site knows one (PRD §4a). */ + actorUserId?: string; + period: string; + /** Post-write totals from `upsertUsage`, including what this write added. */ + totals: readonly UsageWriteTotal[]; + /** Present on a request path; absent in the queue consumers, which run inline. */ + waitUntil?: (promise: Promise) => void; +} + +/** + * Notice any 80% crossing in a just-written batch of usage and email the actor. + * + * Best-effort by contract, like the recording it hangs off: every failure is + * logged and swallowed, because a notification must never fail a merge. + * + * Delivery is `waitUntil` where a request context exists and inline where it + * does not — the deploy runner and the post-merge check are queue consumers + * with no request to hang work off. + */ +export function noticeUsageThresholds(env: Env, logger: Logger, input: UsageNoticeInput): void { + if (!entitlementsEnabled(env)) return; + if (input.totals.length === 0) return; + const run = deliver(env, logger, input).catch((error) => { + logger.warn("Usage threshold notification failed", { + ownerId: input.recorded.ownerId, + error: error instanceof Error ? error.message : String(error), + }); + }); + if (input.waitUntil) input.waitUntil(run); + // Deliberately NOT awaited when there is no `waitUntil`: the two queue + // consumers that land here are mid-deploy and mid-post-merge, and an email + // send is not something either should block a terminal write on. The promise + // is best-effort and its failure path is above. +} + +async function deliver(env: Env, logger: Logger, input: UsageNoticeInput): Promise { + // BYOK spend is not under the allowance at all (migration 049, Goal 5), so a + // `byok` row can never approach a hosted limit and must not warn about one. + const platform = input.totals.filter((total) => total.source === "platform" && total.added > 0); + if (platform.length === 0) return; + + const subject = await resolveEnforcementSubject(env, logger, { + ...(input.actorUserId !== undefined ? { actorUserId: input.actorUserId } : {}), + owner: input.recorded, + }); + if (!subject) return; + + // The totals in hand are the RECORDED subject's rows. When the enforcement + // subject is somebody else — an actor charged for a spend recorded against a + // free org (§4a) — those totals are the wrong denominator for their limit, and + // warning off them would be a number about someone else's month. That case is + // left to the 100% refusal, which names the limit at the moment it binds. + if (subject.ownerId !== input.recorded.ownerId) { + logger.debug("Usage threshold check skipped: recorded subject is not the enforcement subject", { + recordedId: input.recorded.ownerId, + subjectId: subject.ownerId, + }); + return; + } + + const resolved = await new RemoteEntitlements(env, logger).forOwner( + subject.ownerId, + subject.ownerType, + ); + const entitlements = resolved.success ? resolved.data.entitlements : UnlimitedEntitlements; + + for (const total of platform) { + const limit = entitlements.meters[total.meter]; + // Unlimited (-1) has no 80%, and a hard block (0) is not something anyone + // can approach — they hit it on the first unit, where the refusal speaks. + if (!Number.isInteger(limit) || limit <= 0) continue; + const threshold = limit * USAGE_NOTICE_THRESHOLD; + const before = total.quantity - total.added; + if (!(before < threshold && total.quantity >= threshold)) continue; + + const key = + `${RECEIPT_PREFIX}${subject.ownerType}:${subject.ownerId}:${input.period}:` + + `${total.meter}:${THRESHOLD_PERCENT}:${limit}`; + if (await alreadySent(env, logger, key)) continue; + + await send(env, logger, { + subject, + actorUserId: input.actorUserId, + meter: total.meter, + used: total.quantity, + limit, + period: input.period, + }); + await markSent(env, logger, key); + } +} + +async function alreadySent(env: Env, logger: Logger, key: string): Promise { + if (!env.STATE) return true; + try { + return (await env.STATE.get(key)) !== null; + } catch (error) { + // A receipt we cannot read is treated as SENT: a duplicate warning is a + // worse failure than a missing one, and the 100% refusal still speaks. + logger.warn("Usage threshold receipt unreadable; treating as already sent", { + error: error instanceof Error ? error.message : String(error), + }); + return true; + } +} + +async function markSent(env: Env, logger: Logger, key: string): Promise { + if (!env.STATE) return; + try { + await env.STATE.put(key, String(Date.now()), { expirationTtl: RECEIPT_TTL_SECONDS }); + } catch (error) { + logger.warn("Usage threshold receipt could not be written", { + error: error instanceof Error ? error.message : String(error), + }); + } +} + +/** + * Mail the person whose allowance it is. + * + * An org has no email column (`migrations/003_orgs.sql`) and there is no org UI + * to render a banner in, so under §4a the addressee is the actor — the person + * whose allowance is being spent. `users.email` is `UNIQUE NOT NULL` and there + * is no email-change route in the tree, so a user can only ever mail themselves: + * this is not an enumeration or spam surface. + */ +async function send( + env: Env, + logger: Logger, + opts: { + subject: { ownerId: string; ownerType: "user" | "org" }; + actorUserId?: string; + meter: string; + used: number; + limit: number; + period: string; + }, +): Promise { + const recipientId = + opts.actorUserId ?? (opts.subject.ownerType === "user" ? opts.subject.ownerId : undefined); + if (!recipientId) return; + const fromAddress = env.EMAIL_FROM_ADDRESS; + // Both halves are required and one of them was declared for staging only until + // this work; without either, the send path returns early and the feature would + // ship silently never sending. See wrangler.toml's [env.production.vars]. + if (!env.EMAIL || !fromAddress || !env.DB) return; + + const user = await getUser(env.DB, recipientId, logger); + if (!user.success) { + logger.warn("Usage threshold notice not sent: recipient could not be resolved", { + userId: recipientId, + }); + return; + } + + const content = getUsageThresholdEmail({ + meter: opts.meter, + used: opts.used, + limit: opts.limit, + percent: THRESHOLD_PERCENT, + period: opts.period, + resetsAt: periodResetsAt(opts.period), + }); + await env.EMAIL.send({ + to: user.data.email, + from: { email: fromAddress, name: "Stratum" }, + subject: content.subject, + text: content.text, + html: content.html, + }); + logger.info("Usage threshold notice sent", { + userId: recipientId, + meter: opts.meter, + limit: opts.limit, + period: opts.period, + }); +} diff --git a/src/deploy/runner.ts b/src/deploy/runner.ts index cc0ecbaa..6e7d2f71 100644 --- a/src/deploy/runner.ts +++ b/src/deploy/runner.ts @@ -29,6 +29,7 @@ * injects its own seams: without them, testing any of the above needs a real git * remote, a real clock, and a real provider account. */ +import { checkMeter, resolveEnforcementSubject } from "../billing/enforcement"; import { type LlmProviderCatalog, llmProviderCatalog } from "../evaluation/llm-providers"; import { parsePolicyContent } from "../evaluation/policy-loader"; import type { DeployConfig, DeployRejection } from "../evaluation/types"; @@ -573,6 +574,17 @@ async function runOneDeployment( }); } + // Volume, and only volume (PRD §8). The queue is FIFO with + // `max_concurrency = 1` and no priority, so this cannot make anything fairer; + // what it bounds is how many deploys a month one subject may run. It goes + // BEFORE the claim so a refusal never takes the lease it would then have to + // release, and a refusal writes a persisted `failed` row with a named reason + // and acks — the message is not left to burn its two retries into the DLQ + // over a limit that redelivery cannot change, and it is never silently acked, + // because a deploy that did not happen must be visible. + const volume = await checkDeployVolume(ctx, deployment); + if (volume !== null) return await failWithoutRunning(ctx, deployment, volume); + const claim = await claimDeployment(env.DB, logger, { projectId: project.id, deploymentId: deployment.id, @@ -694,6 +706,37 @@ async function runOneDeployment( return record; } +/** + * The reason this deploy may not run against the subject's allowance, or `null` + * when it may. + * + * The enforcement subject is the deploy's requester when a human asked for it + * (an approve or a retry), and otherwise the project's own billing subject: a + * merge-triggered deploy names no acting user at all, and a check against the + * recorded owner is weaker than PRD §4a wants but strictly better than the only + * alternative here, which is no check. + * + * Inert with the billing service unconfigured — `checkMeter` returns before it + * reads anything — so a self-hoster's deploys are untouched. + */ +async function checkDeployVolume(ctx: RunContext, deployment: Deployment): Promise { + const { env, logger, project } = ctx; + const actorUserId = deployment.requestedByType === "user" ? deployment.requestedById : undefined; + const subject = await resolveEnforcementSubject(env, logger, { + ...(actorUserId !== undefined ? { actorUserId } : {}), + owner: { ownerId: project.ownerId, ownerType: project.ownerType }, + }); + if (!subject) return null; + const decision = await checkMeter(env, logger, { + subject, + meter: "deploys_month", + estimate: 1, + nowMs: ctx.now(), + what: "This deployment", + }); + return decision.admitted ? null : (decision.reason ?? "Deploy allowance exhausted"); +} + /** * Resolve secrets and hand the tree to the provider. Returns what should be * written; never writes anything itself, so its caller's `finally` owns the row. @@ -864,6 +907,11 @@ async function readTree( projectId: project.id, ...(changeId !== undefined ? { changeId } : {}), ...(subject ?? {}), + // Best-effort inline, with no `waitUntil`: this is a queue consumer and + // there is no request to schedule against (PRD §8). `git_ops` maps to no + // meter today, so this notices nothing until `deploys_month` gains a D1 + // writer — wired anyway so that writer does not have to remember. + notify: { env }, }, [{ kind: "git_ops", quantity: 1 }], ); diff --git a/src/email/templates.ts b/src/email/templates.ts index bb8cdca4..f35fa01c 100644 --- a/src/email/templates.ts +++ b/src/email/templates.ts @@ -294,6 +294,84 @@ Your code management platform`; return { subject, text, html: wrapEmail({ title: subject, body }) }; } +export interface UsageThresholdEmailData { + /** The meter key, e.g. `llm_tokens_month`. */ + meter: string; + used: number; + limit: number; + /** The threshold that was crossed, as a percentage. */ + percent: number; + /** 'YYYY-MM' UTC. */ + period: string; + /** ISO 8601 instant the period rolls over and the allowance refills. */ + resetsAt: string; +} + +/** Human wording for a meter key. Falls back to the key so a new meter still sends. */ +function meterTitle(meter: string): string { + switch (meter) { + case "llm_tokens_month": + return "AI review tokens"; + case "sandbox_ms_month": + return "sandbox time"; + case "deploys_month": + return "deployments"; + default: + return meter; + } +} + +/** + * "You have used 80% of X" — sent once per (subject, period, meter, threshold, + * limit) when a usage meter crosses its warning threshold. + * + * Written as a heads-up rather than a warning, because it arrives while nothing + * is blocked yet: it names what is running out, when it refills, and the two + * ways to avoid meeting the wall. Sent only when the cloud entitlements service + * is configured — a self-hoster has no limits to be warned about. + */ +export function getUsageThresholdEmail(data: UsageThresholdEmailData): { + subject: string; + text: string; + html: string; +} { + const { used, limit, percent, period, resetsAt } = data; + const title = meterTitle(data.meter); + const subject = `You've used ${percent}% of your Stratum ${title}`; + + const text = `Heads up — this month's ${title} on your Stratum account is ${percent}% used. + +Used: ${used.toLocaleString("en-US")} of ${limit.toLocaleString("en-US")} (${period}) +Refills: ${resetsAt} + +Nothing is blocked yet. When the allowance runs out, evaluations that need it +stop passing until it refills, so if you expect to need more this month you can +either raise your plan limits or point the llm evaluator at your own provider +key with \`provider:\` in .stratum/policy.yaml. + +-- +Stratum +Your code management platform`; + + const body = ` +
+ ${escapeHtml(String(percent))}% of your ${escapeHtml(title)} is used +
+
+ You've used ${escapeHtml(used.toLocaleString("en-US"))} + of ${escapeHtml(limit.toLocaleString("en-US"))} + for ${escapeHtml(period)}. The allowance refills at ${escapeHtml(resetsAt)}. +
+
+ Nothing is blocked yet. If you expect to need more this month, raise your plan + limits, or point the llm evaluator at your own + provider key with provider: in + .stratum/policy.yaml. +
`; + + return { subject, text, html: wrapEmail({ title: subject, body }) }; +} + /** * Generic email template wrapper with Stratum branding * TODO: Use this for future email templates (notifications, invites, etc.) diff --git a/src/evaluation/llm-byok.ts b/src/evaluation/llm-byok.ts index 21dd3203..12ceebfb 100644 --- a/src/evaluation/llm-byok.ts +++ b/src/evaluation/llm-byok.ts @@ -35,27 +35,41 @@ export type LlmProviderSelection = | { status: "unavailable"; reason: string }; /** - * The `llm` entry the evaluator will actually read. + * Every `llm` entry in the policy. * - * `find`, not `filter`: `LLMEvaluator` itself takes the first `llm` entry and - * ignores any others, so resolving from a later one would run the credential of - * a provider whose model and threshold are not the ones in force. + * `filter`, not `find`, and the difference is a fail-open. `LLMEvaluator` reads + * the FIRST `llm` entry, while `buildEvaluators` constructs an evaluator for + * EVERY one — so resolving a provider from the first entry alone meant + * `[{llm}, {llm, provider: anthropic}]` produced two evaluators on `env.AI`, + * with no error and no log: a policy that asked to run on its own key silently + * running on the operator's. `sanitizeLlmConfig` now rejects a second entry at + * parse time, and {@link resolveLlmProvider} refuses to choose between them for + * any policy that reaches it another way (a KV-cached one, say). */ -export function llmConfigOf(policy: EvalPolicy): LlmEvaluatorConfig | undefined { - return policy.evaluators.find((entry): entry is LlmEvaluatorConfig => entry.type === "llm"); +export function llmConfigsOf(policy: EvalPolicy): LlmEvaluatorConfig[] { + return policy.evaluators.filter((entry): entry is LlmEvaluatorConfig => entry.type === "llm"); } /** * Resolve the provider for this policy, loading the project's credential only * when the policy names one. * - * The plaintext credential reaches exactly one place: an `Authorization` (or + * The plaintext credential is used in exactly one place: an `Authorization` (or * `x-api-key`) header on the provider request. `loadSecretValues` is the only * read path for a secret value in the codebase and is held to the same rule as * the deploy runner — no route returns it, renders it, or logs it. That an * HTTP request *triggers* this resolution is not the same thing as the value * being route-reachable. * + * "Used in one place" is not the same as "reaches one place", and the + * difference was a live leak: `new Headers()` throws a `TypeError` that quotes + * the offending value, so a stored credential containing CR, LF or NUL became + * an error message, then an `EvalResult.reason`, then a rendered line on a + * world-readable change page. Both ends are now closed — the store rejects + * control characters (`storage/project-secrets.ts`) and the provider builds its + * headers outside the try, mapping a failure to a constant that interpolates + * nothing (`llm-provider.ts`). + * * **There is deliberately no MCP surface for writing the credential this * reads** (PRD §4c). `stratum_set_provider_key` looks like an obvious gap and * is ruled out on purpose: a provider key must never pass through a model's @@ -69,7 +83,17 @@ export async function resolveLlmProvider( policy: EvalPolicy, logger: Logger, ): Promise { - const providerName = llmConfigOf(policy)?.provider; + const configs = llmConfigsOf(policy); + if (configs.length > 1) { + // Never "pick the first one". Which entry wins decides whose bill the run + // lands on, and a policy that cannot say is a policy this must refuse. + return { + status: "unavailable", + reason: `policy declares ${configs.length} "llm" evaluator entries; at most one is allowed, so no provider could be resolved`, + }; + } + + const providerName = configs[0]?.provider; // The skip. Nothing below this line runs for a project that never opted in. if (providerName === undefined) return { status: "platform" }; diff --git a/src/evaluation/llm-evaluator.ts b/src/evaluation/llm-evaluator.ts index cb9869c7..de62c9ca 100644 --- a/src/evaluation/llm-evaluator.ts +++ b/src/evaluation/llm-evaluator.ts @@ -1,3 +1,11 @@ +import { + type EnforcementSubject, + checkMeter, + resolveEnforcementSubject, + settleMeter, +} from "../billing/enforcement"; +import { usagePeriod } from "../storage/usage"; +import type { Env } from "../types"; import type { AppError } from "../utils/errors"; import { ExternalServiceError } from "../utils/errors"; import type { Logger } from "../utils/logger"; @@ -6,7 +14,13 @@ import { err, ok } from "../utils/result"; import type { LlmProvider } from "./llm-provider"; import { LlmProviderResponseError } from "./llm-provider"; import { sanitizePolicy } from "./sanitize-policy"; -import type { EvalPolicy, EvalResult, Evaluator, LlmEvaluatorConfig } from "./types"; +import type { + EvalPolicy, + EvalResult, + EvaluationContext, + Evaluator, + LlmEvaluatorConfig, +} from "./types"; const DEFAULT_MODEL = "@cf/meta/llama-3.1-8b-instruct"; const DEFAULT_THRESHOLD = 0.7; @@ -20,6 +34,23 @@ const MAX_DIFF_CHARS_FLOOR = 1_000; // pathological policy must not blow the context (or starve the diff), so an // oversize one fails closed before any model call. const MAX_POLICY_CONTEXT_CHARS = 8_000; +/** + * Upper bound on the verdict this evaluator asks for, in tokens. + * + * Reserved alongside the input estimate because a reservation has to cover the + * whole spend before any of it happens, and the response is part of it. The + * verdict is a small JSON object — a score, a boolean, one sentence, a short + * issue list — so this is generous by a wide margin on purpose: over-reserving + * is corrected by the settle a few seconds later, while under-reserving is a + * limit that admits more than it says. + */ +const MAX_OUTPUT_TOKENS = 1_000; + +/** + * What the refusals call the thing they are refusing. One string so the two + * metered gates below cannot describe the same evaluator differently. + */ +const GATE_LABEL = "AI review"; const SYSTEM_PROMPT = [ "You are a rigorous code reviewer acting as an automated merge gate.", @@ -40,19 +71,40 @@ export class LLMEvaluator implements Evaluator { * later layer can reconstruct, since the recorded sample is all it sees. * Defaults to `"platform"`, which the sample omits entirely (absent means * the operator paid). + * @param enforcement What this evaluator needs to consult an allowance: the + * `Env` holding the entitlements configuration and the `USAGE_METER` + * binding, and a `waitUntil` for warming an org's plan off the response + * path (PRD §4a). Omitted where a caller has no environment to give — + * every check is then skipped, exactly as it is for a self-hoster with no + * billing service configured. */ constructor( private provider: LlmProvider, private costSource: "platform" | "byok" = "platform", + private enforcement?: { + env: Env; + waitUntil?: (promise: Promise) => void; + }, ) {} async evaluate( diff: string, policy: EvalPolicy, logger: Logger, + context?: EvaluationContext, ): Promise> { logger.debug("Starting LLM evaluation"); + const nowMs = Date.now(); + const period = usagePeriod(new Date(nowMs)); + /** Set once a reservation is taken, so the `finally` knows what to settle. */ + let subject: EnforcementSubject | null = null; + let tokensReserved: number | null = null; + let rateReserved = false; + let reachedProvider = false; + /** True token cost of the run, once known; settled against the reservation. */ + let tokensSpent = 0; + try { const config = policy.evaluators.find((e): e is LlmEvaluatorConfig => e.type === "llm"); const model = config?.model ?? DEFAULT_MODEL; @@ -93,6 +145,77 @@ export class LLMEvaluator implements Evaluator { ]; const promptChars = messages.reduce((sum, m) => sum + m.content.length, 0); + + // Allowances are consulted HERE, between building the request and making + // it, and a refusal below returns a FAILING gate rather than an absent + // one. That is the whole shape of this check: a policy that requires AI + // review must not stop requiring it because someone ran out of credit, + // which is what dropping the evaluator (or passing it) would do. It + // mirrors the oversize-policy branch above — `ok()` with `passed: false` + // and a reason a human and an agent can both act on. + if (this.enforcement && context?.billing) { + const { env, waitUntil } = this.enforcement; + subject = await resolveEnforcementSubject(env, logger, { + ...(context.billing.actorUserId !== undefined + ? { actorUserId: context.billing.actorUserId } + : {}), + owner: { ownerId: context.billing.ownerId, ownerType: context.billing.ownerType }, + ...(waitUntil !== undefined ? { waitUntil } : {}), + }); + + if (subject) { + // The burst bound first, and it applies to BYOK too (PRD §4b): the + // operator still pays for the subrequest and the Worker wall time, + // whoever owns the tokens. One evaluation, no settle — a rate check + // is a reserve that is never given back. + const rate = await checkMeter(env, logger, { + subject, + meter: "evaluations_per_hour", + estimate: 1, + nowMs, + period, + what: GATE_LABEL, + }); + rateReserved = rate.reserved; + if (!rate.admitted) { + return ok({ + score: 0, + passed: false, + reason: rate.reason ?? `${GATE_LABEL} could not run: evaluation rate limit reached`, + }); + } + + // The token allowance, and the one thing BYOK does lift: a project + // paying its own provider bill is not spending the hosted allowance, + // so charging it against one would be charging twice. + if (this.costSource === "platform") { + const estimate = Math.ceil(maxDiffChars / 4) + MAX_OUTPUT_TOKENS; + const tokens = await checkMeter(env, logger, { + subject, + meter: "llm_tokens_month", + estimate, + nowMs, + period, + what: GATE_LABEL, + }); + if (tokens.reserved) tokensReserved = estimate; + if (!tokens.admitted) { + return ok({ + score: 0, + passed: false, + reason: tokens.reason ?? `${GATE_LABEL} could not run: token allowance is used up`, + }); + } + } + } + } + + // Written only for `byok`: an absent `source` already means "platform", + // so always emitting it would add a field that says nothing to every + // sample the operator pays for. + const source = this.costSource === "byok" ? { source: "byok" as const } : {}; + + reachedProvider = true; const runResult = await this.provider.run(model, messages); if (!runResult.success) { @@ -103,11 +226,43 @@ export class LLMEvaluator implements Evaluator { // below, and `runEvaluation` copies the message of the error the other // branch returns into a reason of its own. if (runResult.error instanceof LlmProviderResponseError) { - logger.error(`LLM evaluation failed: ${runResult.error.message}`); + // The `cause` is logged, not described. It is the only record of what + // the provider actually sent — a `SyntaxError` from `JSON.parse` + // quotes the bytes that failed — and attaching it would be pointless + // if nothing ever read it. It goes to the OPERATOR's log and never + // into the reason below, which is rendered to anyone who can read the + // change; the bound keeps a large body out of a log line. + const cause = runResult.error.cause; + logger.error( + `LLM evaluation failed: ${runResult.error.message}`, + runResult.error, + cause === undefined + ? undefined + : { cause: String(cause instanceof Error ? cause.message : cause).slice(0, 200) }, + ); + // A verdict truncated at the provider's token cap was generated, and + // billed. `failClosed` further down records the cost of a response it + // could not read, and this is the same case one layer up: without it a + // charged call records zero. Only counts the provider actually + // REPORTED are recorded — a failure that produced no tokens (a 401, a + // refused redirect) carries none and bills nothing. + const billed = runResult.error.usage; + if (billed) tokensSpent = billed.inputTokens + billed.outputTokens; return ok({ score: 0, passed: false, reason: `LLM evaluator error: ${runResult.error.message}`, + ...(billed + ? { + costs: [ + { + kind: "llm_tokens" as const, + quantity: billed.inputTokens + billed.outputTokens, + ...source, + }, + ], + } + : {}), }); } logger.error("LLM evaluation failed", runResult.error); @@ -116,10 +271,6 @@ export class LLMEvaluator implements Evaluator { const responseText = runResult.data.text; const usage = runResult.data.usage; - // Written only for `byok`: an absent `source` already means "platform", - // so always emitting it would add a field that says nothing to every - // sample the operator pays for. - const source = this.costSource === "byok" ? { source: "byok" as const } : {}; // Real counts when the provider reports them, otherwise ~4 chars/token — // the distinction is per response, not per source. Every provider behind // this seam reports counts when it has them, Workers AI included; a @@ -135,6 +286,10 @@ export class LLMEvaluator implements Evaluator { ...source, }, ]; + // What the settle in the `finally` corrects the reservation by. Read off + // the same sample the ledger records, so the counter and the ledger + // cannot disagree about what one run cost. + tokensSpent = costs[0]?.quantity ?? 0; // A gate whose verdict can't be read has not passed. Never infer a score // from prose ("LGTM") — that let unparseable output half-approve a merge. @@ -202,6 +357,36 @@ export class LLMEvaluator implements Evaluator { error instanceof Error ? error : undefined, ) as AppError, ); + } finally { + // In a `finally` because a provider call that throws has still consumed + // the reservation. Leaving it standing would charge the subject the whole + // upper bound for a run that produced nothing, and a handful of those + // walls somebody off at a fraction of an allowance they never spent. + if (subject && this.enforcement) { + const { env } = this.enforcement; + if (tokensReserved !== null) { + await settleMeter(env, logger, { + subject, + meter: "llm_tokens_month", + delta: tokensSpent - tokensReserved, + nowMs, + period, + }); + } + // A rate slot is given back only when the evaluation never reached the + // provider — a token refusal after the rate check admitted. Once the + // request is out the door the operator has paid for the subrequest and + // the wall time it bounds, so a failed run still costs its slot. + if (rateReserved && !reachedProvider) { + await settleMeter(env, logger, { + subject, + meter: "evaluations_per_hour", + delta: -1, + nowMs, + period, + }); + } + } } } } diff --git a/src/evaluation/llm-provider.ts b/src/evaluation/llm-provider.ts index 2ebd8468..fb1fae20 100644 --- a/src/evaluation/llm-provider.ts +++ b/src/evaluation/llm-provider.ts @@ -61,14 +61,27 @@ export interface LlmProvider { * shape problem, or a constant this file chose, and nothing else. */ export class LlmProviderResponseError extends AppError { - constructor(message: string, cause?: unknown) { + /** + * Token counts the provider reported for a response that cannot be used. + * + * Present only where the provider *billed* the call anyway — a verdict + * truncated at the token cap is a completed, charged generation. The + * evaluator records it as a cost sample, so a run the provider charged for is + * not recorded as zero. Absent for every failure that produced no tokens (a + * non-2xx status, an unreadable body), which must not be billed. + */ + readonly usage?: LlmUsage; + + constructor(message: string, cause?: unknown, usage?: LlmUsage) { super(message, "LLM_PROVIDER_RESPONSE", 502); this.name = "LlmProviderResponseError"; // Attached rather than described: a `SyntaxError` from `JSON.parse` quotes // the bytes that failed, which are the provider's body — exactly what the // message must not carry. Dropping it entirely would swallow the only - // record of what went wrong. + // record of what went wrong, and `LLMEvaluator` logs it (operator-visible) + // without ever putting it in the user-visible reason. if (cause !== undefined) this.cause = cause; + if (usage !== undefined) this.usage = usage; } } @@ -84,6 +97,15 @@ function transportError(error: unknown): AppError { /** How long any provider gets before the gate gives up on it. */ const DEFAULT_TIMEOUT_MS = 30_000; +/** + * The two constant messages the request path can fail with. Constants rather + * than interpolations because both are reached with attacker-influenced values + * in hand — a redirect `Location`, or a header value that could not be built — + * and every message from this file becomes a user-visible evaluation reason. + */ +const REDIRECT_REFUSED = "provider redirected the request, which is not followed"; +const HEADER_CONSTRUCTION_FAILED = "provider request headers could not be constructed"; + export interface ProviderOptions { /** Give up on the request after this many milliseconds. */ timeoutMs?: number; @@ -184,6 +206,22 @@ abstract class HttpLlmProvider implements LlmProvider { * its own shape, because a body that parsed is not a body that can be read. */ protected async post(path: string, headers: Record, body: unknown) { + // Built HERE, outside the try, and never from inside it. `new Headers` throws + // a `TypeError` whose message QUOTES the offending value, and the value here + // is the project's API key: inside the try that message would become a + // transport error, then an `ExternalServiceError`, then `EvalResult.reason`, + // which is persisted on the change and rendered on a page that is + // world-readable for a public project. The store rejects control characters + // in a secret value (`storage/project-secrets.ts`) so this should be + // unreachable; it is the second of the two ends, and it maps to a constant + // that interpolates nothing. + let requestHeaders: Headers; + try { + requestHeaders = new Headers({ "content-type": "application/json", ...headers }); + } catch { + return err(new LlmProviderResponseError(HEADER_CONSTRUCTION_FAILED) as AppError); + } + const controller = new AbortController(); const timer = setTimeout( () => controller.abort(), @@ -193,11 +231,21 @@ abstract class HttpLlmProvider implements LlmProvider { try { const response = await fetch(`${this.baseUrl.replace(/\/+$/, "")}${path}`, { method: "POST", - headers: { "content-type": "application/json", ...headers }, + headers: requestHeaders, body: JSON.stringify(body), signal: controller.signal, + // Never follow a redirect: the allowlist binds the host this request was + // *sent* to, and a followed 3xx would re-send the prompt (the diff and + // the policy) to a host nobody validated — with `x-api-key` still + // attached, since the Fetch spec strips only `Authorization` + // cross-origin. The sibling `webhook-evaluator.ts` does the same. + redirect: "manual", }); + if (response.status >= 300 && response.status < 400) { + return err(new LlmProviderResponseError(REDIRECT_REFUSED) as AppError); + } + if (!response.ok) { // The status, never the body: an error body can echo the API key that // failed, and this message becomes a user-visible evaluation reason. @@ -291,9 +339,13 @@ export class AnthropicProvider extends HttpLlmProvider { // model's formatting for a limit this file chose. Anthropic says so // explicitly: `stop_reason` is `"max_tokens"` in exactly that case. if (body.stop_reason === "max_tokens") { + // With the counts: the generation ran to the cap and the provider billed + // it, so an unusable verdict here still costs real tokens. return err( new LlmProviderResponseError( `provider response was truncated at the ${AnthropicProvider.MAX_TOKENS}-token limit`, + undefined, + usageOrNothing(body.usage?.input_tokens, body.usage?.output_tokens).usage, ) as AppError, ); } @@ -366,9 +418,12 @@ export class OpenAiCompatibleProvider extends HttpLlmProvider { // not the model, and the half-object left behind must not be reported as // the model's own formatting failure. if (choice?.finish_reason === "length") { + // Billed, like the Anthropic sibling — see the note there. return err( new LlmProviderResponseError( `provider response was truncated at the ${OpenAiCompatibleProvider.MAX_COMPLETION_TOKENS}-token limit`, + undefined, + usageOrNothing(body.usage?.prompt_tokens, body.usage?.completion_tokens).usage, ) as AppError, ); } diff --git a/src/evaluation/llm-providers.ts b/src/evaluation/llm-providers.ts index 136a5541..aa3766ec 100644 --- a/src/evaluation/llm-providers.ts +++ b/src/evaluation/llm-providers.ts @@ -14,6 +14,7 @@ */ import type { Env } from "../types"; import type { Logger } from "../utils/logger"; +import { privateHostReason } from "../utils/validation"; /** * Wire protocols the evaluator can speak, one per HTTP implementation in @@ -73,79 +74,69 @@ export function providerSecretName(providerName: string): string { /** * Why this host must not be fetched, or null when it is fine. * - * Loopback, link-local and RFC1918 addresses are the ones that turn an outbound - * fetch into a read of something only the Worker's own network can see. The - * WHATWG URL parser normalizes IPv4 for special schemes, so the obfuscated - * spellings (`0x7f.0.0.1`, `2130706433`) arrive here already in dotted-quad - * form and are caught by the same rules. + * A thin alias for `privateHostReason` (`utils/validation.ts`), which is the + * one host filter in the codebase and is also what `validateWebhookUrl` uses. + * This used to be a second implementation, and the two drifted exactly as a + * duplicated security filter always does: this copy missed CGNAT, the + * `.internal` suffix (GCP's metadata endpoint is `metadata.google.internal`), + * everything in `fe80::/10` not literally spelled `fe80`, and the + * IPv4-compatible IPv6 form `[::127.0.0.1]`, which the URL parser hands on as + * `[::7f00:1]`. * - * **Known limit:** this is an address check, not a resolution check. A public - * DNS name that resolves to 127.0.0.1 defeats it, and nothing at parse time can - * see that. It is defence in depth — the actual control is that only the - * operator writes this variable, and that a policy file can never add to it. + * Kept as a named export because the reason is interpolated into the operator's + * config error, and because the allowlist is where this rule is *enforced*: + * every URL rule is checked once, at parse time, not per request. */ export function blockedHostReason(hostname: string): string | null { - // A trailing dot is the same name to a resolver; lowercase for the literals. - const host = hostname.toLowerCase().replace(/\.$/, ""); - - if (host === "localhost" || host.endsWith(".localhost")) { - return "'localhost' names the loopback interface"; - } - - if (host.startsWith("[") && host.endsWith("]")) { - const inner = host.slice(1, -1); - if (inner === "::1") return "[::1] is the IPv6 loopback address"; - if (inner === "::" || inner === "::0") return "[::] is the IPv6 unspecified address"; - if (inner.startsWith("fe80")) return `${host} is an IPv6 link-local address`; - // fc00::/7 — unique local, the IPv6 equivalent of RFC1918. - if (/^f[cd]/.test(inner)) return `${host} is an IPv6 unique-local address`; - // An IPv4-mapped address exists only to reach an IPv4 host through an IPv6 - // socket; an operator has no reason to write one, and each spelling - // (dotted-quad tail, hex tail) would need its own check to be safe. - if (inner.includes("::ffff:")) return `${host} is an IPv4-mapped IPv6 address`; - return null; - } - - const quad = /^(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$/.exec(host); - if (!quad) return null; - const [a, b] = [Number(quad[1]), Number(quad[2])]; - - if (a === 127) return `${host} is in the loopback range 127.0.0.0/8`; - if (a === 0) return `${host} is in the unspecified range 0.0.0.0/8`; - if (a === 169 && b === 254) return `${host} is in the link-local range 169.254.0.0/16`; - if (a === 10) return `${host} is in the private range 10.0.0.0/8`; - if (a === 172 && b >= 16 && b <= 31) return `${host} is in the private range 172.16.0.0/12`; - if (a === 192 && b === 168) return `${host} is in the private range 192.168.0.0/16`; - return null; + return privateHostReason(hostname); } -function baseUrlError(name: string, raw: unknown): string | null { +/** A validated `baseUrl`, normalized, or why the entry cannot be used. */ +type BaseUrlCheck = { ok: true; baseUrl: string } | { ok: false; reason: string }; + +function checkBaseUrl(name: string, raw: unknown): BaseUrlCheck { if (typeof raw !== "string" || raw.length === 0) { - return `provider "${name}": "baseUrl" is required and must be a string`; + return { ok: false, reason: `provider "${name}": "baseUrl" is required and must be a string` }; } let url: URL; try { url = new URL(raw); } catch { - return `provider "${name}": "baseUrl" is not a valid absolute URL`; + return { ok: false, reason: `provider "${name}": "baseUrl" is not a valid absolute URL` }; } if (url.protocol !== "https:") { // Not tidiness: the request carries the project's API key and the diff. - return `provider "${name}": "baseUrl" must use https, not ${url.protocol.replace(":", "")}`; + return { + ok: false, + reason: `provider "${name}": "baseUrl" must use https, not ${url.protocol.replace(":", "")}`, + }; } if (url.username !== "" || url.password !== "") { - return `provider "${name}": "baseUrl" must not embed credentials`; + return { ok: false, reason: `provider "${name}": "baseUrl" must not embed credentials` }; } if (url.search !== "" || url.hash !== "") { - return `provider "${name}": "baseUrl" must not carry a query string or fragment`; + return { + ok: false, + reason: `provider "${name}": "baseUrl" must not carry a query string or fragment`, + }; } const blocked = blockedHostReason(url.hostname); if (blocked) { - return `provider "${name}": "baseUrl" host is not reachable from an allowlist — ${blocked}`; + return { + ok: false, + reason: `provider "${name}": "baseUrl" host is not reachable from an allowlist — ${blocked}`, + }; } - return null; + + // Normalized, never the raw string. `https://host#` and `https://host?` both + // pass the two emptiness checks above — the parser reports an empty hash and + // an empty search for them — while the bytes the operator wrote would put the + // provider's `/messages` path inside a fragment or a query. Origin plus path + // is what `href` would give minus exactly those two, and the trailing slash + // goes because the provider classes append an absolute path. + return { ok: true, baseUrl: `${url.origin}${url.pathname}`.replace(/\/+$/, "") }; } /** @@ -222,15 +213,13 @@ export function parseLlmProviders(raw: string | undefined): LlmProvidersParse { }; } - const urlProblem = baseUrlError(name, source.baseUrl); - if (urlProblem) return { status: "invalid", reason: urlProblem }; + const url = checkBaseUrl(name, source.baseUrl); + if (!url.ok) return { status: "invalid", reason: url.reason }; providers.set(name, { name, kind: source.kind as LlmProviderKind, - // Trailing slashes are stripped by the provider classes when they build a - // path; store what the operator wrote. - baseUrl: source.baseUrl as string, + baseUrl: url.baseUrl, }); } @@ -258,10 +247,13 @@ function parseCached(raw: string | undefined): LlmProvidersParse { * * Empty-on-invalid is the fail-closed direction, not a silent disable: an empty * catalog means every policy naming a provider is rejected by - * `sanitizeLlmConfig` and blocks merges with a reason naming the provider, while - * projects on Workers AI — which is every project that never opted into BYOK — - * are untouched by an operator's typo. The typo itself is reported by - * `llmProvidersConfigError` and logged here on the isolate's first parse. + * `sanitizeLlmConfig`, which fails the policy file closed and blocks merges with + * a reason that NAMES the unconfigured provider — the rejection it records is + * carried into `parsePolicyContent`'s malformed reason, so the person fixing the + * file is not left with a bare count of unusable entries. Projects on Workers AI + * — which is every project that never opted into BYOK — are untouched by an + * operator's typo. The typo itself is reported by `llmProvidersConfigError` and + * logged here on the isolate's first parse. */ export function llmProviderCatalog( env: Pick, diff --git a/src/evaluation/policy-loader.ts b/src/evaluation/policy-loader.ts index e4996421..6eadb47b 100644 --- a/src/evaluation/policy-loader.ts +++ b/src/evaluation/policy-loader.ts @@ -32,6 +32,33 @@ const DEFAULT_POLICY: EvalPolicy = { */ export const PROTECTED_CONFIG_FILES = [".stratum/policy.yaml", "stratum.config.json"] as const; +/** + * Most `evaluators:` entries one policy file may declare. + * + * The `deploys:` sibling has had `MAX_DEPLOY_ENTRIES` from the start and for + * the same reason, which applies at least as strongly here: every entry runs + * inside one Worker invocation, sharing its CPU, memory and subrequest budget, + * and an `llm` or `webhook` entry is an outbound call. Without a cap a policy + * file amplifies one merge into an unbounded number of provider requests. + * + * Unlike a rejected deploy entry this fails the whole file closed, the way + * every other unusable evaluator entry does: an evaluator is a gate, and + * quietly running the first sixteen of a hundred gates is not "most of the + * policy", it is a policy nobody wrote. + */ +export const MAX_EVALUATOR_ENTRIES = 16; + +/** + * Most `llm` entries one policy file may declare. + * + * One, because two cannot be resolved: `LLMEvaluator` reads the first entry's + * model and threshold while `buildEvaluators` builds an evaluator for every + * entry, so `[{llm}, {llm, provider: x}]` used to mean two runs on the + * operator's Workers AI bill from a policy that asked for one on its own key. + * Rejecting the file is the only answer that does not silently pick one. + */ +const MAX_LLM_ENTRIES = 1; + /** * The operator's provider allowlist, as the policy parser sees it. * @@ -183,8 +210,20 @@ export function parsePolicyContent( // the parsed input (or, via `DEFAULT_POLICY`'s own array, into the module // default) by mutating what it was handed. const declared = raw.evaluators as unknown[]; + if (declared.length > MAX_EVALUATOR_ENTRIES) { + return { + status: "malformed", + reason: `'evaluators' declares ${declared.length} entries; at most ${MAX_EVALUATOR_ENTRIES} are allowed`, + }; + } + + // Why each dropped entry was dropped, in the file's own terms. Collected + // rather than only logged because the count alone ("2 unusable entries") + // is what the merge gate shows the person who has to fix the file, and a + // count does not name the provider that is not configured. + const rejections: string[] = []; const evaluators = declared - .map((entry) => sanitizeEvaluator(entry, logger, providers)) + .map((entry, index) => sanitizeEvaluator(entry, logger, providers, rejections, index)) .filter((entry): entry is EvaluatorConfig => entry !== null); // Any dropped entry fails the gate closed — not just the case where every @@ -200,11 +239,24 @@ export function parsePolicyContent( // `buildEvaluators`, so a policy naming a future evaluator type does not // trip this. if (evaluators.length < declared.length) { + const dropped = declared.length - evaluators.length; return { status: "malformed", - reason: `${declared.length - evaluators.length} unusable entr${ - declared.length - evaluators.length === 1 ? "y" : "ies" - } in 'evaluators'`, + reason: `${dropped} unusable entr${dropped === 1 ? "y" : "ies"} in 'evaluators'${ + rejections.length > 0 ? `: ${rejections.join("; ")}` : "" + }`, + }; + } + + // One `llm` entry, checked here rather than inside `sanitizeLlmConfig` + // because it is a property of the LIST, not of an entry — and it fails the + // file closed rather than dropping the extra, so nothing has to decide + // which of two entries the project meant. + const llmEntries = evaluators.filter((entry) => entry.type === "llm").length; + if (llmEntries > MAX_LLM_ENTRIES) { + return { + status: "malformed", + reason: `'evaluators' declares ${llmEntries} 'llm' entries; at most ${MAX_LLM_ENTRIES} is allowed, because a second one would run a second model call on a configuration nothing can choose between`, }; } policy.evaluators = evaluators; @@ -457,13 +509,22 @@ const MODEL_PATTERN = /^[A-Za-z0-9._:@/-]+$/; * closed and lives in `LLM_PROVIDERS`; a policy may only *select* from it by * name. A `baseUrl` in the file is warned about and dropped like any other * unrecognized key — see `llm-providers.ts` for why the field cannot exist. + * + * @param rejections Collects why an entry was dropped, so the merge-blocking + * reason can name the provider instead of only counting entries. Optional: + * a caller that only wants the config need not care. */ export function sanitizeLlmConfig( source: Record, logger: Logger, providers: LlmProviderCatalog, + rejections?: string[], ): LlmEvaluatorConfig | null { const config: LlmEvaluatorConfig = { type: "llm" }; + const reject = (reason: string): null => { + rejections?.push(reason); + return null; + }; for (const key of Object.keys(source)) { if (!LLM_CONFIG_KEYS.has(key)) { @@ -476,7 +537,7 @@ export function sanitizeLlmConfig( logger.error("Dropping llm evaluator entry whose provider is not a string", undefined, { submitted: forLog(source.provider), }); - return null; + return reject("llm evaluator: 'provider' must be a string naming a configured provider"); } if (!providers.has(source.provider)) { logger.error( @@ -484,7 +545,9 @@ export function sanitizeLlmConfig( undefined, { provider: forLog(source.provider) }, ); - return null; + return reject( + `llm evaluator names provider "${forLog(source.provider)}", which this instance has not configured in LLM_PROVIDERS`, + ); } config.provider = source.provider; } @@ -495,11 +558,26 @@ export function sanitizeLlmConfig( logger.error("Dropping llm evaluator entry with an unusable model", undefined, { submitted: forLog(source.model), }); - return null; + return reject("llm evaluator: 'model' is not a usable model identifier"); } config.model = model; } + // A BYOK entry must name its model. The default (`@cf/meta/llama-3.1-8b-instruct`) + // is a Workers AI model id, and posting it to Anthropic or an OpenAI-compatible + // endpoint fails every single call — a gate that never runs, discovered one + // merge at a time. Requiring the field is clearer than inventing a default per + // provider kind, which would silently pick a model (and a price) for the + // project. Workers AI keeps the default: there the id is the right one. + if (config.provider !== undefined && config.model === undefined) { + logger.error("Dropping llm evaluator entry that names a provider but no model", undefined, { + provider: config.provider, + }); + return reject( + `llm evaluator names provider "${config.provider}" but no "model"; a provider entry must name the model to run, since the default is a Workers AI model id`, + ); + } + // Clamped, not rejected, for the reason `clampScore` gives: a threshold // outside [0,1] is a bounded mistake, and 1.5 (nothing passes) is the safe // direction anyway. -5 is not — it would make every model verdict pass — so @@ -538,14 +616,19 @@ function sanitizeEvaluator( raw: unknown, logger: Logger, providers: LlmProviderCatalog, + rejections: string[], + index: number, ): EvaluatorConfig | null { + const at = `evaluators[${index}]`; if (typeof raw !== "object" || raw === null || Array.isArray(raw)) { logger.warn("Dropping non-object evaluator entry", { submitted: forLog(raw) }); + rejections.push(`${at} is not an evaluator object`); return null; } const source = raw as Record; if (typeof source.type !== "string") { logger.warn("Dropping evaluator entry without a string type", { submitted: forLog(raw) }); + rejections.push(`${at} has no string 'type'`); return null; } @@ -554,7 +637,12 @@ function sanitizeEvaluator( // Replaces the copy-through below for `llm`, which used to hand the parsed // object's fields to the evaluator unvalidated — including any `baseUrl` or // inline credential the file carried. - if (source.type === "llm") return sanitizeLlmConfig(source, logger, providers); + if (source.type === "llm") { + const entryRejections: string[] = []; + const llm = sanitizeLlmConfig(source, logger, providers, entryRejections); + for (const reason of entryRejections) rejections.push(`${at}: ${reason}`); + return llm; + } if (source.type === "webhook") { // `url` is required by the type, so check it rather than asserting it. The @@ -565,6 +653,7 @@ function sanitizeEvaluator( logger.warn("Dropping webhook evaluator entry without a url", { submitted: forLog(raw), }); + rejections.push(`${at}: webhook evaluator has no 'url'`); return null; } // The same unbounded-timeout defect the sandbox entry had, one array diff --git a/src/evaluation/sanitize-policy.ts b/src/evaluation/sanitize-policy.ts index cf7a469f..90eaea5b 100644 --- a/src/evaluation/sanitize-policy.ts +++ b/src/evaluation/sanitize-policy.ts @@ -6,20 +6,33 @@ import type { EvalPolicy, EvaluatorConfig } from "./types"; * Matched on the key rather than the value because there is no reliable way to * recognise a secret by its bytes, and this runs on repository content: whatever * a user wrote in the policy file, under whatever key they chose. + * + * Stems, not spellings. Plurals are normalized by {@link isCredentialWord}, so + * `token` covers `tokens` and `apiKey` covers `apiKeys` and `access_tokens` + * without either having to be listed. `auth` is here on its own because it is + * one of the most common credential key names there is, and listing only + * `authorization` let it through. */ const CREDENTIAL_WORDS = new Set([ "secret", - "secrets", "token", "password", "passwd", + "pwd", + "passphrase", "credential", - "credentials", + "cred", + "auth", "authorization", + "authorisation", "bearer", "key", "apikey", "signature", + "sig", + "hmac", + "pat", + "otp", ]); /** @@ -34,7 +47,21 @@ function words(key: string): string[] { .filter((word) => word.length > 0); } -/** Keys stripped from every evaluator entry regardless of shape. */ +/** + * Whether one word of a field name names a credential. + * + * The trailing-`s` rule is what makes `tokens`, `apiKeys` and `passwords` + * behave like their singulars. It is deliberately only a *stem* test and not a + * substring one, because the property the strip list must keep is that + * `keystone` and `monkeys` survive: `monkeys` stems to `monkey`, which is not a + * credential word, while a substring match would strip both. + */ +function isCredentialWord(word: string): boolean { + if (CREDENTIAL_WORDS.has(word)) return true; + return word.endsWith("s") && CREDENTIAL_WORDS.has(word.slice(0, -1)); +} + +/** Keys stripped wherever they appear in the policy, regardless of shape. */ function isStrippedKey(key: string): boolean { // The provider *name* is not a credential, but it names the operator's // infrastructure to a third party the operator did not choose — and the model @@ -44,7 +71,49 @@ function isStrippedKey(key: string): boolean { // project that fits today cannot start failing its gate closed by opting into // BYOK. if (key === "provider") return true; - return words(key).some((word) => CREDENTIAL_WORDS.has(word)); + return words(key).some(isCredentialWord); +} + +/** + * Bounds on the walk, so a hostile policy file cannot make sanitization + * expensive (or produce an output larger than the input). + * + * The depth bound doubles as the cycle guard: a YAML alias cycle + * (`a: &x {b: *x}`) is a perfectly well-formed file and an infinite structure, + * and the walk below would otherwise not terminate on one. + */ +const MAX_DEPTH = 8; +const MAX_NODES = 2_000; +const MAX_ARRAY_ITEMS = 256; +const DEPTH_ELISION = "[omitted: policy nesting limit]"; +const SIZE_ELISION = "[omitted: policy size limit]"; + +interface Budget { + left: number; +} + +/** Recursively rebuild `value` without any credential-shaped key. */ +function stripValue(value: unknown, depth: number, budget: Budget): unknown { + if (budget.left <= 0) return SIZE_ELISION; + budget.left -= 1; + if (value === null || typeof value !== "object") return value; + if (depth >= MAX_DEPTH) return DEPTH_ELISION; + + if (Array.isArray(value)) { + const kept: unknown[] = value + .slice(0, MAX_ARRAY_ITEMS) + .map((item) => stripValue(item, depth + 1, budget)); + if (value.length > MAX_ARRAY_ITEMS) kept.push(SIZE_ELISION); + return kept; + } + + const out: Record = {}; + for (const [key, nested] of Object.entries(value)) { + if (isStrippedKey(key)) continue; + if (budget.left <= 0) break; + out[key] = stripValue(nested, depth + 1, budget); + } + return out; } /** @@ -52,14 +121,19 @@ function isStrippedKey(key: string): boolean { * an LLM prompt, or a webhook request body). * * This is the last thing that runs before the policy is handed to a third - * party, so it strips by key shape rather than by known field. `webhook.secret` - * — the one field this used to handle — is used to sign outgoing requests and - * must never appear in a payload itself. The general rule exists because - * `sanitizeEvaluator` copies an entry of an unmodelled `type` through whole - * (a policy naming a future evaluator), so a credential-shaped field on such an - * entry would otherwise reach the model verbatim. `sanitizeLlmConfig`'s - * whitelist already keeps one off an `llm` entry; this is the backstop for - * every entry it does not model. + * party, so it strips by key shape rather than by known field, **recursively, + * over the whole policy object** — not just over the top level of each + * evaluator entry, which is all it used to do. Both of the holes that left are + * real: `parsePolicyContent` spreads unknown ROOT keys onto the policy, so an + * `apiKey:` written at the top of `.stratum/policy.yaml` reached the model + * prompt and the policy-supplied webhook URL (a host the policy file chose); + * and `sanitizeEvaluator` copies an entry of an unmodelled `type` through + * whole, so a credential nested one level inside such an entry — an object or + * an array, exactly the copy-through case the strip list exists for — was + * untouched. + * + * `webhook.secret` is the field this was originally written for: it signs + * outgoing requests and must never appear in a payload itself. * * Note what this is *not*: a defence against the model seeing secrets in the * diff. The diff is the thing being reviewed, and `redactSecrets` @@ -68,11 +142,16 @@ function isStrippedKey(key: string): boolean { * `EvalResult.reason` (`llm-evaluator.ts`). */ export function sanitizePolicy(policy: EvalPolicy): EvalPolicy { + const stripped = stripValue(policy, 0, { left: MAX_NODES }); + const base = + typeof stripped === "object" && stripped !== null && !Array.isArray(stripped) + ? (stripped as Record) + : {}; return { - ...policy, - evaluators: policy.evaluators.map((cfg: EvaluatorConfig) => { - const entries = Object.entries(cfg).filter(([key]) => !isStrippedKey(key)); - return Object.fromEntries(entries) as EvaluatorConfig; - }), - }; + ...base, + // Restated rather than trusted: `evaluators` is typed as an array and both + // callers serialize the result, so a policy large enough to exhaust the + // node budget must still hand back the shape the type promises. + evaluators: Array.isArray(base.evaluators) ? (base.evaluators as EvaluatorConfig[]) : [], + } as EvalPolicy; } diff --git a/src/evaluation/types.ts b/src/evaluation/types.ts index e03b3259..e313f786 100644 --- a/src/evaluation/types.ts +++ b/src/evaluation/types.ts @@ -46,6 +46,21 @@ export interface BillingContext { ownerType: "user" | "org"; /** The project whose policy is being enforced. */ projectId: string; + /** + * The user who ran this evaluation, when the call site knows one. + * + * Not the payer, and deliberately a separate field from `ownerId`: PRD §4a + * separates the subject a spend is RECORDED against from the subject a LIMIT + * is checked against, because an org-owned project bills the org and an org + * costs nothing to create. Recording still names the owner above; enforcement + * uses this. Absent where no acting user exists (a queue consumer), which is + * not the same as free — it means the check falls back to the recorded + * subject. + * + * Like every other field here, it must never be forwarded off-box: the + * webhook evaluator names the fields it sends for exactly this reason. + */ + actorUserId?: string; } /** diff --git a/src/merge/post-merge.ts b/src/merge/post-merge.ts index 82d5182f..7d60f3ac 100644 --- a/src/merge/post-merge.ts +++ b/src/merge/post-merge.ts @@ -94,6 +94,11 @@ export async function runPostMergeCheck( projectId: project.id, changeId: opts.changeId, ...(subject ?? {}), + // No `waitUntil`: this runs in the merge queue consumer, with no + // request to hang background work off, so delivery is best-effort + // inline (PRD §8). No actor either — a queue message names no user, + // so the check falls back to the recorded subject. + notify: { env }, }, [ { kind: "sandbox_ms", quantity: Date.now() - runStartedAt }, diff --git a/src/middleware/rate-limit.ts b/src/middleware/rate-limit.ts index 5ce8d700..0309f145 100644 --- a/src/middleware/rate-limit.ts +++ b/src/middleware/rate-limit.ts @@ -1,4 +1,6 @@ import type { MiddlewareHandler } from "hono"; +import { enforcementBinding } from "../billing/enforcement"; +import { RemoteEntitlements, UNLIMITED, entitlementsEnabled } from "../billing/entitlements"; import { isGitHttpPath } from "../routes/git-http"; import { isPostHogProxyPath, isPostHogSdkPath } from "../routes/posthog-proxy"; import type { Env } from "../types"; @@ -75,10 +77,35 @@ export function rateLimitMiddleware(opts?: RateLimitOptions): MiddlewareHandler< const isAuthenticated = Boolean(userId ?? agentId); const defaultLimit = isAuthenticated ? 1000 : 60; + // A plan's own request rate, when we hold one. Deliberately narrow: + // + // - **Users only.** Today's limiter keys on `userId ?? agentId ?? IP` while + // entitlements key on a billing subject, and resolving an agent to its + // owner needs a D1 hop inside this middleware. Folding a user's agents + // into one bucket is also a user-visible TIGHTENING (a user plus five + // agents get 6x1000 rpm today), so agents keep their own bucket at + // today's numbers and that change gets its own CHANGELOG entry. + // - **Cached read only.** `forOwner` never fetches; the auth middleware's + // warm is what makes it hit. A miss falls back to today's numbers rather + // than blocking the request on a billing round trip. + // - **`-1` is not a number to enforce.** Unlimited is the FAIL-OPEN default + // of the entitlements layer, so honouring it literally would let one + // uncached miss (or one billing outage) uncap a user's request rate. The + // default cap stands instead. + const planLimit = await entitledRequestsPerMinute(c.env, userId); + // Observe-only never tightens: with `ENTITLEMENTS_ENFORCE` off a plan may + // raise this ceiling but not lower it, so the month of measurement PRD §8 + // asks for cannot start 429-ing anybody. + const subjectLimit = + planLimit === null + ? defaultLimit + : enforcementBinding(c.env) + ? planLimit + : Math.max(planLimit, defaultLimit); // Far above what a real session produces and far below anything worth // relaying through someone else's Worker. const analyticsLimit = 600; - const limit = isAnalyticsIngest ? analyticsLimit : (opts?.requestsPerMinute ?? defaultLimit); + const limit = isAnalyticsIngest ? analyticsLimit : (opts?.requestsPerMinute ?? subjectLimit); const identifier = userId ?? agentId ?? c.req.header("CF-Connecting-IP") ?? "anonymous"; const minuteBucket = Math.floor(Date.now() / 60000); @@ -144,6 +171,30 @@ export function rateLimitMiddleware(opts?: RateLimitOptions): MiddlewareHandler< }; } +/** + * The user's own `requests_per_minute`, or `null` when there is none to use. + * + * `null` covers every case the caller must fall back to today's numbers for: the + * billing service is unconfigured (every self-hoster), the caller is not a user, + * the cache is cold, or the plan says "unlimited" — which here means "we do not + * know", because unlimited is what this layer fails open to. + */ +async function entitledRequestsPerMinute( + env: Env, + userId: string | undefined, +): Promise { + if (!userId) return null; + if (!entitlementsEnabled(env)) return null; + const resolved = await new RemoteEntitlements(env, logger).forOwner(userId, "user"); + if (!resolved.success) return null; + // "default" is a cache miss, not a plan. Only a value somebody actually sent + // us gets to move a limit. + if (resolved.data.source === "default") return null; + const limit = resolved.data.entitlements.rates.requests_per_minute; + if (!Number.isInteger(limit) || limit < 0 || limit === UNLIMITED) return null; + return limit; +} + /** * Rate limiting middleware specifically for import endpoints. * Provides both per-user and per-project rate limiting to prevent resource exhaustion. diff --git a/src/queue/usage-meter.ts b/src/queue/usage-meter.ts index 69c07bde..c5b07aaf 100644 --- a/src/queue/usage-meter.ts +++ b/src/queue/usage-meter.ts @@ -80,6 +80,20 @@ const RATE_WINDOWS: Record = { evaluations_per_hour: { windowSeconds: 60 * 60, bucketSeconds: 5 * 60 }, }; +/** + * How far back a rate meter's window looks, in seconds. + * + * Exported so the refusal copy can tell a caller when its burst allowance frees + * up without restating the window — a message promising a reset the object does + * not honour is worse than no message (PRD §4c). + * + * @param meter - The rate meter to describe + * @returns The window length in seconds + */ +export function rateWindowSeconds(meter: RateMeterKey): number { + return RATE_WINDOWS[meter].windowSeconds; +} + /** Bucket id (`floor(epochSeconds / bucketSeconds)`) -> quantity in that bucket. */ type RateBuckets = Record; diff --git a/src/routes/changes.ts b/src/routes/changes.ts index a29fdd39..c60c727c 100644 --- a/src/routes/changes.ts +++ b/src/routes/changes.ts @@ -1415,17 +1415,26 @@ app.post("/changes/:id/evaluate", async (c) => { baseOid, } = diffResult.data; - const evaluators = await buildEvaluators(c.env, policy, project, logger, { - remote: workspace.remote, - token: workspaceReadToken.data, - ref: evaluatedSha, - }); + const evaluators = await buildEvaluators( + c.env, + policy, + project, + logger, + { + remote: workspace.remote, + token: workspaceReadToken.data, + ref: evaluatedSha, + }, + getWaitUntil(c), + ); // The base this re-evaluation's diff was built against — not `change.baseSha`, // which records the base at creation and is exactly the value that has gone // stale by the time a change is re-evaluated (#274). const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha: baseOid, - billing: billingContextFor(project), + // This route is user-credentialed (agent tokens are rejected above), so the + // acting user is the caller — the subject PRD §4a checks a limit against. + billing: billingContextFor(project, userId), }); const recordResult = await recordEvalRuns(c.env.DB, logger, id, evalRuns); @@ -1448,6 +1457,7 @@ app.post("/changes/:id/evaluate", async (c) => { changeId: id, workspace: change.workspace, ...(evaluateSubject ?? {}), + notify: { env: c.env, actorUserId: userId, waitUntil: getWaitUntil(c) }, }, evaluateCostSamples, ); diff --git a/src/routes/deployments.ts b/src/routes/deployments.ts index 7c4f3a4b..c6234189 100644 --- a/src/routes/deployments.ts +++ b/src/routes/deployments.ts @@ -12,9 +12,13 @@ * path. `canReadProject` returns true unconditionally for a public project, * and a provider error payload can echo request context back at us. * - * No route on either router returns a secret value; there is no read path for - * one anywhere in the codebase (`loadSecretValues` is the deploy runner's, and - * is not reachable from here). + * No route on either router returns a secret value. `loadSecretValues` is the + * only function that yields one, and it is not reachable from this router at + * all — but it is no longer route-unreachable in general: BYOK resolves a + * project's provider credential on the change-creation path + * (`evaluation/llm-byok.ts`). That value never comes back out either; it goes + * straight into the provider request header. The invariant these routes keep is + * the narrower, checkable one: nothing here reads a value. * * The browser talks to these same routes. An HTML form can only issue GET or * POST and cannot read a JSON body, so the secret writes have form-friendly diff --git a/src/routes/projects.ts b/src/routes/projects.ts index d6fd49b9..ef48aa7a 100644 --- a/src/routes/projects.ts +++ b/src/routes/projects.ts @@ -1,5 +1,6 @@ import { Hono } from "hono"; import type { Context } from "hono"; +import { checkGauge, resolveEnforcementSubject } from "../billing/enforcement"; import { importRateLimitMiddleware, rateLimitMiddleware, @@ -65,6 +66,7 @@ import { filterReadableProjects, isDirectOwner, } from "../utils/authz"; +import { getWaitUntil } from "../utils/execution-ctx"; import { createLogger } from "../utils/logger"; import type { Logger } from "../utils/logger"; import { readJsonWithLimit } from "../utils/request-body"; @@ -181,6 +183,39 @@ app.post("/", async (c) => { visibility = body.visibility; } + // The private-project gauge (PRD §8): a thing that exists rather than a thing + // that is spent, so it is counted on demand and there is nothing to reserve. + // + // Counted over the namespace the project is being created INTO, not over + // everything the actor owns. For a personal namespace those are the same set. + // For an org namespace it is the org's own private projects checked against + // the actor's limit — which is not the §4a subject's own gauge, and is the + // reading that bounds BOTH cases: KV has no index by owner, so an actor-wide + // count is a full scan of every project on the instance, per creation. + if (visibility === "private") { + const gaugeSubject = await resolveEnforcementSubject(c.env, logger, { + actorUserId: userId, + owner: { ownerId: owner.id, ownerType: owner.type }, + waitUntil: getWaitUntil(c), + }); + if (gaugeSubject) { + const existing = await listProjectsByNamespace(c.env.STATE, namespace, logger); + // A failed listing does not block a creation: the gauge is a limit, not a + // correctness rule, and this layer fails open by design. + if (existing.success) { + const decision = await checkGauge(c.env, logger, { + subject: gaugeSubject, + count: "private_projects", + current: existing.data.filter((entry) => entry.visibility === "private").length, + what: "Creating another private project", + }); + if (!decision.admitted) { + return forbidden(decision.reason ?? "Private project limit reached"); + } + } + } + } + const seed = body.seed === "true" || body.seed === true || c.req.query("seed") === "true"; const files = body.files !== undefined diff --git a/src/routes/sync-management.ts b/src/routes/sync-management.ts index 82c41d9d..e2fc401d 100644 --- a/src/routes/sync-management.ts +++ b/src/routes/sync-management.ts @@ -26,6 +26,7 @@ import { import type { Env } from "../types"; import { projectDefaultBranch, projectDisplayName } from "../types"; import { canWriteProject } from "../utils/authz"; +import { getWaitUntil } from "../utils/execution-ctx"; import { createLogger } from "../utils/logger"; import { readJsonWithLimit } from "../utils/request-body"; import { notFound, ok } from "../utils/response"; @@ -669,14 +670,23 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { // pass BEFORE resolveConflict creates one — so there is no ref a sandbox // could check out. A policy naming `sandbox` fails closed via // UnavailableEvaluator, same as any other missing prerequisite. - const evaluators = await buildEvaluators(c.env, policy, project, logger); + const evaluators = await buildEvaluators( + c.env, + policy, + project, + logger, + undefined, + getWaitUntil(c), + ); // `buildManualResolutionDiff` resolved this base from the clone it built the // diff on, so it names the tree the resolution actually applies to (#274). // This path runs the LLM evaluator too, so it needs a payer as much as // change creation does. const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha, - billing: billingContextFor(project), + // The resolver is the actor: this route is user-credentialed, and PRD §4a + // checks the allowance against the person who ran the suite. + billing: billingContextFor(project, userId), }); // Recorded here, BEFORE the verdict is acted on, because the spend is @@ -696,6 +706,7 @@ app.post("/projects/conflicts/:id/resolve", async (c) => { ...(conflictCtx.changeId ? { changeId: conflictCtx.changeId } : {}), workspace: conflictCtx.workspaceName, ...(resolveSubject ?? {}), + notify: { env: c.env, actorUserId: userId, waitUntil: getWaitUntil(c) }, }, [{ kind: "git_ops", quantity: 1 }, ...evalRuns.flatMap(({ result }) => result.costs ?? [])], ); diff --git a/src/services/change-flow.ts b/src/services/change-flow.ts index 1b6c6427..fc795112 100644 --- a/src/services/change-flow.ts +++ b/src/services/change-flow.ts @@ -127,13 +127,25 @@ export class UnavailableEvaluator implements Evaluator { * aggregates under an empty key that no account will ever reconcile, where an * absent subject is at least visibly absent. */ -export function billingContextFor(project: ProjectEntry): BillingContext | undefined { +export function billingContextFor( + project: ProjectEntry, + /** + * The user who ran the evaluation, when the caller knows one. Carried + * alongside the payer rather than instead of it: PRD §4a checks a limit + * against the actor while the ledger keeps recording the owner, so a spend on + * an org-owned project can be billed to the org and counted against the + * person who caused it. For an agent-authored change this is the agent's + * OWNER — an agent is not a subject any more than it is a payer. + */ + actorUserId?: string, +): BillingContext | undefined { if (project.ownerType !== "user" && project.ownerType !== "org") return undefined; if (!project.ownerId || !project.id) return undefined; return { ownerId: project.ownerId, ownerType: project.ownerType, projectId: project.id, + ...(actorUserId ? { actorUserId } : {}), }; } @@ -161,6 +173,17 @@ export async function buildEvaluators( project: ProjectEntry, logger: Logger, workspaceRepo?: SandboxRepoAccess, + /** + * Cloudflare Workers `ExecutionContext.waitUntil`, when the caller has one. + * + * Used by the metered evaluator for one thing: warming an ORG's entitlements + * off the response path (PRD §4a). Auth middleware warms the caller, but it + * knows the caller and not the project, so nothing before this point can warm + * the org a project belongs to. Absent means no warm — the check then resolves + * to the actor, which is the safe direction and the documented first-contact + * behaviour, never a blocking fetch on the merge path. + */ + waitUntil?: (promise: Promise) => void, ): Promise> { // Async only because of this: a policy that selects a BYOK provider needs the // project's credential resolved before the evaluator exists. It is a no-op — @@ -193,16 +216,22 @@ export async function buildEvaluators( { type: "llm", evaluator: new UnavailableEvaluator("llm", llmProvider.reason) }, ]; } + const enforcement = { env, ...(waitUntil !== undefined ? { waitUntil } : {}) }; if (llmProvider.status === "byok") { return [ { type: "llm", - evaluator: new LLMEvaluator(llmProvider.provider, "byok"), + evaluator: new LLMEvaluator(llmProvider.provider, "byok", enforcement), }, ]; } if (env.AI) - return [{ type: "llm", evaluator: new LLMEvaluator(new WorkersAiProvider(env.AI)) }]; + return [ + { + type: "llm", + evaluator: new LLMEvaluator(new WorkersAiProvider(env.AI), "platform", enforcement), + }, + ]; return [ { type: "llm", @@ -480,17 +509,26 @@ export async function createChangeWithEvaluation( baseOid, } = diffResult.data; - const evaluators = await buildEvaluators(env, policy, project, logger, { - remote: workspaceRemote, - token: workspaceReadToken.data, - ref: evaluatedSha, - }); + const evaluators = await buildEvaluators( + env, + policy, + project, + logger, + { + remote: workspaceRemote, + token: workspaceReadToken.data, + ref: evaluatedSha, + }, + waitUntil, + ); // `baseOid`, not the `baseSha` resolved further up: that one is read before // the diff clone, so the default branch can advance in between. The evaluated // base must be the one the diff was actually built against (#274). const { evalRuns, evalResult } = await runEvaluation(evaluators, diff, policy, logger, { baseSha: baseOid, - billing: billingContextFor(project), + // The actor is the human behind the change — the agent's owner for an + // agent-authored one — which is who an allowance follows under PRD §4a. + billing: billingContextFor(project, createdByUserId), }); const newStatus: Change["status"] = evalResult.passed ? "accepted" : "needs_changes"; @@ -521,6 +559,11 @@ export async function createChangeWithEvaluation( changeId: change.id, workspace: workspaceName, ...(createSubject ?? {}), + notify: { + env, + ...(createdByUserId !== undefined ? { actorUserId: createdByUserId } : {}), + ...(waitUntil !== undefined ? { waitUntil } : {}), + }, }, createCostSamples, ); diff --git a/src/storage/costs.ts b/src/storage/costs.ts index f2f8f410..76f5673c 100644 --- a/src/storage/costs.ts +++ b/src/storage/costs.ts @@ -1,9 +1,17 @@ +import { noticeUsageThresholds } from "../billing/usage-notifications"; +import type { Env } from "../types"; import { AppError } from "../utils/errors"; import { newId } from "../utils/ids"; import type { Logger } from "../utils/logger"; import { type Result, err, ok } from "../utils/result"; import { getAgent } from "./agents"; -import { type UsageDelta, meterForCostKind, upsertUsage, usagePeriod } from "./usage"; +import { + type UsageDelta, + type UsageWriteTotal, + meterForCostKind, + upsertUsage, + usagePeriod, +} from "./usage"; export type CostKind = "llm_tokens" | "sandbox_ms" | "git_ops"; @@ -126,8 +134,8 @@ async function accumulateUsage( opts: { ownerId?: string; ownerType?: BillingSubject["ownerType"] }, samples: CostSample[], createdAt: string, -): Promise { - if (!opts.ownerId || !opts.ownerType) return; +): Promise<{ subject: BillingSubject; period: string; totals: UsageWriteTotal[] } | null> { + if (!opts.ownerId || !opts.ownerType) return null; const deltas: UsageDelta[] = []; for (const sample of samples) { // `git_ops` maps to no meter; see `meterForCostKind`. @@ -136,7 +144,7 @@ async function accumulateUsage( deltas.push({ meter, quantity: sample.quantity, source: sample.source ?? "platform" }); } } - if (deltas.length === 0) return; + if (deltas.length === 0) return null; // The Result is deliberately not folded into `recordCosts`'s own: the ledger // rows above ARE written, and reporting an error would tell the caller the // spend went unrecorded when it did not. `upsertUsage` logs its own failure. @@ -147,13 +155,12 @@ async function accumulateUsage( // migration 049), so a failed upsert is a permanent under-count of that // owner's month. Accepted because the alternative is failing a merge over an // accounting write, which is the worse trade in both directions. - await upsertUsage( - db, - logger, - { ownerId: opts.ownerId, ownerType: opts.ownerType }, - usagePeriod(new Date(createdAt)), - deltas, - ); + const subject: BillingSubject = { ownerId: opts.ownerId, ownerType: opts.ownerType }; + const period = usagePeriod(new Date(createdAt)); + const written = await upsertUsage(db, logger, subject, period, deltas); + // The post-write totals travel back up so the caller can notice an 80% + // crossing without a second read; a failed write reports nothing to notice. + return written.success ? { subject, period, totals: written.data } : null; } /** @@ -174,6 +181,20 @@ export async function recordCosts( workspace?: string; ownerId?: string; ownerType?: BillingSubject["ownerType"]; + /** + * Everything the threshold notification needs, when the call site can + * supply it. Absent means the write is recorded and nobody is told they are + * approaching a limit — which is the correct behaviour for a site with no + * `Env` in hand, and the reason this is optional rather than a second + * required argument. Inert unless the cloud billing service is configured. + */ + notify?: { + env: Env; + /** The acting user, for the §4a enforcement subject and the addressee. */ + actorUserId?: string; + /** Present on a request path; the queue consumers pass nothing. */ + waitUntil?: (promise: Promise) => void; + }; }, samples: CostSample[], ): Promise> { @@ -236,7 +257,16 @@ export async function recordCosts( }); // Only after the ledger write succeeded: the aggregate derives from those // rows, so it must never count spend the ledger does not record. - await accumulateUsage(db, logger, opts, usable, createdAt); + const accumulated = await accumulateUsage(db, logger, opts, usable, createdAt); + if (accumulated && opts.notify) { + noticeUsageThresholds(opts.notify.env, logger, { + recorded: accumulated.subject, + ...(opts.notify.actorUserId !== undefined ? { actorUserId: opts.notify.actorUserId } : {}), + period: accumulated.period, + totals: accumulated.totals, + ...(opts.notify.waitUntil !== undefined ? { waitUntil: opts.notify.waitUntil } : {}), + }); + } return ok(undefined); } catch (error) { const appError = diff --git a/src/storage/project-secrets.ts b/src/storage/project-secrets.ts index 21fcd038..616c9333 100644 --- a/src/storage/project-secrets.ts +++ b/src/storage/project-secrets.ts @@ -1,9 +1,16 @@ /** - * Per-project encrypted secret store, backing deploy provider credentials. + * Per-project encrypted secret store, backing deploy provider credentials and + * the per-project LLM provider keys BYOK runs on. * - * Write-only by construction: the only function that yields a plaintext value is - * {@link loadSecretValues}, which exists for the deploy runner and is deliberately - * not reachable from any route. Everything else returns names and metadata. + * Write-only by construction *to a caller*: the only function that yields a + * plaintext value is {@link loadSecretValues}, and no route returns, renders or + * logs what it resolves. It is no longer unreachable from a route, though — + * BYOK resolves a provider credential on the change-creation path, so the + * routes that create or re-evaluate a change (`routes/changes.ts`, + * `routes/git-http.ts`, `routes/sync-management.ts`) reach it through + * `evaluation/llm-byok.ts`. What keeps that safe is the rule below, not the + * call graph: the value it returns goes into one outbound request header and + * nowhere else. Everything else here returns names and metadata. */ import type { Env } from "../types"; import { type SecretScope, decryptSecret, deriveSecretKey, encryptSecret } from "../utils/crypto"; @@ -101,7 +108,43 @@ function validateName(name: string): ValidationError | null { ); } +/** + * Characters no credential contains and that a header cannot carry: C0, DEL, + * and NUL among them. + * + * Rejected at the WRITE, which is the only place the value is ever attacker- + * chosen. A stored value ends up in an outbound header, and `new Headers()` + * throws a `TypeError` whose message QUOTES the offending value — so a CR, LF + * or NUL in a stored secret turns into an error message carrying the secret, + * which then becomes an evaluation reason persisted on the change and rendered + * on a page that is world-readable for a public project. The provider closes + * the other end (`evaluation/llm-provider.ts` builds its headers outside the + * try and maps a failure to a constant), so neither end alone is load-bearing. + */ +/** + * Does `value` contain a C0 control character or DEL? + * + * A char-code scan rather than a regex: matching control characters is the + * whole point here, which Biome's `noControlCharactersInRegex` flags, and a + * suppression for a rule this deliberately trips reads worse than the loop. + */ +function hasControlCharacter(value: string): boolean { + for (let i = 0; i < value.length; i++) { + const code = value.charCodeAt(i); + if (code <= 0x1f || code === 0x7f) return true; + } + return false; +} + function validateValue(value: string): ValidationError | null { + if (hasControlCharacter(value)) { + // Deliberately says nothing about the value — not its length, not the + // offending character, and above all not the value itself: this message is + // rendered back to the caller. + return new ValidationError( + "Secret value must not contain control characters (including CR, LF, tab and NUL)", + ); + } const bytes = new TextEncoder().encode(value).length; if (bytes <= MAX_SECRET_VALUE_BYTES) return null; return new ValidationError(`Secret value exceeds ${MAX_SECRET_VALUE_BYTES} bytes`, { @@ -295,8 +338,16 @@ export async function deleteSecret( } /** - * Resolves plaintext secret values for the deploy runner. **The only read path - * for a value anywhere in the codebase** — do not call it from a route. + * Resolves plaintext secret values. **The only read path for a value anywhere + * in the codebase**, and the rule for a caller is that the plaintext goes + * straight into the one outbound request that needs it and is never returned, + * rendered, logged, or interpolated into an error. + * + * Two callers today: the deploy runner (`deploy/runner.ts`), and BYOK provider + * resolution (`evaluation/llm-byok.ts`), which the change-creation and + * re-evaluation routes reach synchronously. The comment that used to stand here + * said no route may call it; that stopped being true when BYOK landed, and a + * rule nobody can follow is worse than the honest one above. * * Derives the AES key once for the whole batch: PBKDF2 at 100k iterations is CPU * work and the deploy consumer runs under a CPU limit, so per-secret derivation diff --git a/src/storage/usage.ts b/src/storage/usage.ts index e0dc1c59..4599b909 100644 --- a/src/storage/usage.ts +++ b/src/storage/usage.ts @@ -28,6 +28,26 @@ export interface UsageDelta { source?: CostSource; } +/** + * One meter's total after an {@link upsertUsage} write, as the UPSERT itself + * reported it. + * + * Returned rather than re-read because the 80% threshold notification is + * edge-triggered (PRD §8): it fires on the CROSSING, which needs the totals + * either side of this write. A second `SELECT` per evaluated change would cost + * a D1 round trip on the change-creation path and still race a concurrent + * recorder — `RETURNING` reports the value the statement that did the adding + * actually wrote. + */ +export interface UsageWriteTotal { + meter: MeterKey; + source: CostSource; + /** The meter's total for this owner and period AFTER this write. */ + quantity: number; + /** What this write added, so the caller can reconstruct the total before it. */ + added: number; +} + export interface UsageSummaryEntry { meter: MeterKey; source: CostSource; @@ -112,7 +132,8 @@ export function usagePeriod(at: Date = new Date()): string { * @param subject - The account the usage is billed to * @param period - 'YYYY-MM' UTC, from `usagePeriod` * @param deltas - Increments to apply; same-meter entries are summed first - * @returns Success, or the database error that stopped the accumulation + * @returns The post-write total for each (meter, source) touched, or the + * database error that stopped the accumulation */ export async function upsertUsage( db: D1Database, @@ -120,7 +141,7 @@ export async function upsertUsage( subject: BillingSubject, period: string, deltas: readonly UsageDelta[], -): Promise> { +): Promise> { // Keyed on both dimensions the row is keyed on, so platform and BYOK spend // for one meter accumulate as two totals rather than one. const merged = new Map(); @@ -150,27 +171,38 @@ export async function upsertUsage( quantity: (existing?.quantity ?? 0) + delta.quantity, }); } - if (merged.size === 0) return ok(undefined); + if (merged.size === 0) return ok([]); const updatedAt = new Date().toISOString(); try { + const totals: UsageWriteTotal[] = []; for (const { meter, source, quantity } of merged.values()) { - await db + // `RETURNING quantity` on the same statement that does the adding: the + // post-write total is what makes the threshold crossing detectable + // without a second read, and reading it back separately would race + // another recorder and report a total this write did not produce. + const row = await db .prepare( "INSERT INTO usage_periods (owner_id, owner_type, period, meter, source, quantity, updated_at) " + "VALUES (?, ?, ?, ?, ?, ?, ?) " + "ON CONFLICT(owner_id, period, meter, source) DO UPDATE SET " + - "quantity = quantity + excluded.quantity, updated_at = excluded.updated_at", + "quantity = quantity + excluded.quantity, updated_at = excluded.updated_at " + + "RETURNING quantity", ) .bind(subject.ownerId, subject.ownerType, period, meter, source, quantity, updatedAt) - .run(); + .first<{ quantity: number }>(); + // A driver that does not hand RETURNING back is not an error: the row is + // written either way, and the only thing lost is the notification edge. + if (row && Number.isFinite(row.quantity)) { + totals.push({ meter, source, quantity: row.quantity, added: quantity }); + } } logger.debug("Usage aggregate updated", { ownerId: subject.ownerId, period, meters: merged.size, }); - return ok(undefined); + return ok(totals); } catch (error) { const appError = error instanceof AppError @@ -238,3 +270,58 @@ export async function getOwnerUsageSummary( return err(appError); } } + +/** + * One owner's totals for a period, **filtered to one source**, as a lookup by + * meter. + * + * Separate from {@link getOwnerUsageSummary}, which reports every source so a + * usage page can show what BYOK covered, because an enforcement read must not. + * Migration 049 put `source` in the primary key precisely so hosted and + * bring-your-own spend accumulate apart; a limit check that summed both would + * charge a project's own provider bill against the operator's allowance and + * contradict Goal 5. Callers checking a limit pass `"platform"`, always. + * + * Meters with no rows are absent rather than zero, as in the summary: "never + * used" and "used nothing" are one fact, and the caller knows its meter set. + * + * @param ownerId - The subject whose rows to sum — the ENFORCEMENT subject for + * a limit check (PRD §4a), which is not always the recorded billing subject + * @param period - 'YYYY-MM' UTC, from `usagePeriod` + * @param source - Which provider account's spend to count + * @returns Totals by meter, or the database error + */ +export async function getOwnerMeterTotals( + db: D1Database, + logger: Logger, + ownerId: string, + period: string, + source: CostSource, +): Promise>, AppError>> { + try { + const result = await db + .prepare( + "SELECT meter, SUM(quantity) AS quantity FROM usage_periods " + + "WHERE owner_id = ? AND period = ? AND source = ? GROUP BY meter", + ) + .bind(ownerId, period, source) + .all<{ meter: string; quantity: number }>(); + const totals: Partial> = {}; + for (const row of result.results) { + if (Number.isFinite(row.quantity)) totals[row.meter as MeterKey] = row.quantity; + } + return ok(totals); + } catch (error) { + const appError = + error instanceof AppError + ? error + : new AppError( + error instanceof Error ? error.message : "Failed to total usage", + "DATABASE_ERROR", + 500, + { operation: "getOwnerMeterTotals", ownerId, period }, + ); + logger.error("Failed to total owner usage", appError, { ownerId, period, source }); + return err(appError); + } +} diff --git a/src/types.ts b/src/types.ts index 3d3ae430..9a1e811a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -183,13 +183,15 @@ export interface Env { * nothing — `entitlementsConfigError` flags exactly that. */ ENTITLEMENTS_ENFORCE?: string; /** JSON array of the operator's LLM providers, e.g. - * `[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]`. + * `[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com/v1"}]`. * The ONLY place a provider endpoint may be named: a project's policy file * selects one of these by name and can never supply a `baseUrl`. Unset (the * default) means Workers AI only. Credentials are NOT here — they live per * project in `project_secrets`. Parsed by `parseLlmProviders`, which rejects * a non-https, credential-bearing, loopback, link-local or private-range - * `baseUrl` at parse time. */ + * `baseUrl` at parse time and stores it normalized. Include the API version + * path: the provider appends `/messages` or `/chat/completions` to whatever + * this names, so a `baseUrl` missing the `/v1` 404s on every call. */ LLM_PROVIDERS?: string; ANALYTICS?: AnalyticsEngineDataset; SANDBOX?: SandboxBinding; diff --git a/src/utils/validation.ts b/src/utils/validation.ts index 35faed2c..bef37d48 100644 --- a/src/utils/validation.ts +++ b/src/utils/validation.ts @@ -227,59 +227,160 @@ export function isValidRepoUrl(value: unknown): value is string { } /** - * Whether a hostname is a private/disallowed IP literal. Handles the obfuscated - * encodings a naive string blocklist misses: integer (2130706433), hex - * (0x7f000001, 0x7f.0.0.1), octal (0177.0.0.1), and IPv6 unspecified/expanded - * loopback ([::], [0:0:0:0:0:0:0:1]). The rule is allowlist-shaped: an IP-ish - * host must be a canonical public dotted-decimal IPv4 (or a public IPv6) — any - * non-canonical numeric form fails closed. Returns false for DNS names, which - * the caller filters separately. + * Why this hostname must not be fetched from the Worker, or null when it is fine. + * + * **The one host filter in this codebase.** `validateWebhookUrl` below and the + * LLM provider allowlist (`blockedHostReason` in `evaluation/llm-providers.ts`) + * both go through here, because two filters that must agree eventually do not: + * the provider copy missed CGNAT, `.internal` (GCP's metadata endpoint is + * `metadata.google.internal`) and everything in `fe80::/10` that is not spelled + * `fe80`. + * + * It handles the obfuscated encodings a naive string blocklist misses: integer + * (2130706433), hex (0x7f000001, 0x7f.0.0.1), octal (0177.0.0.1), and every + * IPv6 spelling of the same address — `[::]`, `[0:0:0:0:0:0:0:1]`, + * `[::ffff:127.0.0.1]` and `[::127.0.0.1]`, which the URL parser rewrites to + * `[::7f00:1]` before anything here sees it. The rule is allowlist-shaped: an + * IP-ish host must be a canonical public dotted-decimal IPv4 (or a public + * IPv6) — any non-canonical numeric form fails closed. + * + * **Known limit:** this is an address check, not a resolution check. A public + * DNS name that resolves to 127.0.0.1 defeats it, and nothing before the fetch + * can see that. */ -function isPrivateIpLiteral(hostname: string): boolean { - // Bracketed IPv6 literal, e.g. "[::1]". - if (hostname.startsWith("[")) { - const inner = hostname.slice(1, -1).toLowerCase(); - const compact = inner.replace(/[0:]/g, ""); - // "" ← "::" / "0:0:…:0" (unspecified); "1" ← "::1" / "0:0:…:1" (loopback). - if (compact === "" || compact === "1") return true; - if (/^f[cd]/.test(inner)) return true; // ULA fc00::/7 - if (/^fe80/.test(inner)) return true; // link-local - if (inner.includes("::ffff:")) return true; // IPv4-mapped - return false; +export function privateHostReason(hostname: string): string | null { + // A trailing dot is the same name to a resolver; lowercase for the literals. + const host = hostname.toLowerCase().replace(/\.$/, ""); + + if (host === "localhost" || host.endsWith(".localhost")) { + return "'localhost' names the loopback interface"; + } + if (host.startsWith("[")) return privateIpv6Reason(host); + if (host.endsWith(".internal")) { + // metadata.google.internal is the cloud metadata endpoint — the single most + // valuable SSRF target there is, and it is a DNS name, not an address. + return `${host} is an internal-only name (.internal)`; } + if (host.endsWith(".local")) return `${host} is a link-local mDNS name (.local)`; + return privateIpv4Reason(host); +} +/** @see privateHostReason — kept as a boolean for the call sites that only branch. */ +function isPrivateIpLiteral(hostname: string): boolean { + return privateHostReason(hostname) !== null; +} + +function privateIpv4Reason(host: string): string | null { // Only inspect hosts that look like a numeric IP literal (all digits/dots, or // containing a hex "0x"). A real DNS name — even a hex-word one like // "beef.cafe" — has other letters and is left to the DNS-name checks. - const looksNumeric = /^[0-9.]+$/.test(hostname) || /0x/i.test(hostname); - if (!looksNumeric) return false; + const looksNumeric = /^[0-9.]+$/.test(host) || /0x/i.test(host); + if (!looksNumeric) return null; // Accept ONLY canonical dotted-decimal IPv4; integer/hex/octal/short forms are // obfuscated addresses → reject. - const parts = hostname.split("."); - if (parts.length !== 4) return true; + const notCanonical = `${host} is not a canonical dotted-quad IPv4 address`; + const parts = host.split("."); + if (parts.length !== 4) return notCanonical; const octets: number[] = []; for (const part of parts) { - if (!/^\d{1,3}$/.test(part)) return true; // hex/non-decimal octet - if (part.length > 1 && part[0] === "0") return true; // leading zero (octal ambiguity) + if (!/^\d{1,3}$/.test(part)) return notCanonical; // hex/non-decimal octet + if (part.length > 1 && part[0] === "0") return notCanonical; // octal ambiguity const n = Number(part); - if (n > 255) return true; + if (n > 255) return notCanonical; octets.push(n); } const [a, b] = octets as [number, number, number, number]; - if (a === 0 || a === 10 || a === 127) return true; - if (a === 172 && b >= 16 && b <= 31) return true; - if (a === 192 && b === 168) return true; - if (a === 169 && b === 254) return true; - if (a === 100 && b >= 64 && b <= 127) return true; // CGNAT 100.64/10 - return false; + if (a === 127) return `${host} is in the loopback range 127.0.0.0/8`; + if (a === 0) return `${host} is in the unspecified range 0.0.0.0/8`; + if (a === 169 && b === 254) return `${host} is in the link-local range 169.254.0.0/16`; + if (a === 10) return `${host} is in the private range 10.0.0.0/8`; + if (a === 172 && b >= 16 && b <= 31) return `${host} is in the private range 172.16.0.0/12`; + if (a === 192 && b === 168) return `${host} is in the private range 192.168.0.0/16`; + if (a === 100 && b >= 64 && b <= 127) return `${host} is in the shared range 100.64.0.0/10`; + return null; +} + +/** + * The eight 16-bit groups of a bracketed IPv6 literal, or null when the text is + * not one this parser recognises. + * + * Expanded rather than string-matched because every interesting IPv6 range has + * more than one spelling: `fe80::/10` is not only the literal prefix "fe80", + * and `::127.0.0.1` reaches the loopback through a dotted tail the URL parser + * has already rewritten to hex. + */ +function expandIpv6(inner: string): number[] | null { + if (!/^[0-9a-f:.]+$/.test(inner)) return null; + + let text = inner; + // A trailing dotted quad (::ffff:127.0.0.1) is two hextets written in decimal. + const dotted = /(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$/.exec(text); + if (dotted) { + const octets = dotted.slice(1, 5).map(Number); + if (octets.some((n) => n > 255)) return null; + const high = ((octets[0] as number) << 8) | (octets[1] as number); + const low = ((octets[2] as number) << 8) | (octets[3] as number); + text = `${text.slice(0, dotted.index)}${high.toString(16)}:${low.toString(16)}`; + } + if (text.includes(".")) return null; + + const halves = text.split("::"); + if (halves.length > 2) return null; + const head = halves[0] === "" ? [] : (halves[0] as string).split(":"); + const tail = halves.length === 2 && halves[1] !== "" ? (halves[1] as string).split(":") : []; + const groups = [...head, ...tail]; + if (groups.some((group) => !/^[0-9a-f]{1,4}$/.test(group))) return null; + if (halves.length === 1 ? groups.length !== 8 : groups.length > 7) return null; + + const zeros = new Array(8 - groups.length).fill(0) as number[]; + return [ + ...head.map((group) => Number.parseInt(group, 16)), + ...zeros, + ...tail.map((group) => Number.parseInt(group, 16)), + ]; +} + +function privateIpv6Reason(host: string): string | null { + const hextets = expandIpv6(host.slice(1, -1)); + // Fail closed: a bracketed literal this parser cannot read is not one anybody + // should be pointing a Worker at. + if (hextets === null) return `${host} is not a recognizable IPv6 address`; + const [h0, h1, h2, h3, h4, h5, h6, h7] = hextets as [ + number, + number, + number, + number, + number, + number, + number, + number, + ]; + + if (h0 === 0 && h1 === 0 && h2 === 0 && h3 === 0) { + if (h4 === 0 && h5 === 0 && h6 === 0 && h7 === 0) { + return `${host} is the IPv6 unspecified address`; + } + if (h4 === 0 && h5 === 0 && h6 === 0 && h7 === 1) { + return `${host} is the IPv6 loopback address`; + } + if (h4 === 0 && h5 === 0xffff) return `${host} is an IPv4-mapped IPv6 address`; + // ::/64 holds nothing routable, and it is where an IPv4-compatible address + // (`::127.0.0.1`, which normalizes to `::7f00:1`) lands. + return `${host} is in the reserved ::/64 range, which includes IPv4-compatible addresses`; + } + if ((h0 & 0xfe00) === 0xfc00) return `${host} is an IPv6 unique-local address`; + if ((h0 & 0xffc0) === 0xfe80) return `${host} is an IPv6 link-local address`; + return null; } /** * Validates an outbound webhook URL. Requires http(s) and rejects hostnames - * that resolve to private space (loopback, RFC 1918, link-local, ULA, CGNAT), - * including obfuscated IP encodings. DNS-level rebinding is out of scope here; - * Workers egress is not a guaranteed second layer, so keep this the primary gate. + * that resolve to private space (loopback, RFC 1918, link-local, ULA, CGNAT, + * `.internal`/`.local`), including obfuscated IP encodings — see + * {@link privateHostReason}, which is the shared filter. DNS-level rebinding is + * out of scope here; Workers egress is not a guaranteed second layer, so keep + * this the primary gate. */ export function validateWebhookUrl(value: unknown, logger?: Logger): ValidationResult { const log = logger ?? defaultLogger; @@ -301,13 +402,9 @@ export function validateWebhookUrl(value: unknown, logger?: Logger): ValidationR const hostname = parsed.hostname.toLowerCase(); if ( - hostname === "localhost" || isPrivateIpLiteral(hostname) || - hostname.endsWith(".internal") || - hostname.endsWith(".local") || - hostname.endsWith(".localhost") || // A bare single label (no dot) is either an intranet name or an obfuscated - // integer/hex IP — reject. Bracketed IPv6 handled above. + // integer/hex IP — reject. Bracketed IPv6 is handled by the filter above. (!hostname.includes(".") && !hostname.startsWith("[")) ) { log.debug("Validation failed - webhook URL targets a private host", { hostname }); diff --git a/tests/enforcement.test.ts b/tests/enforcement.test.ts new file mode 100644 index 00000000..c3959600 --- /dev/null +++ b/tests/enforcement.test.ts @@ -0,0 +1,683 @@ +/** + * Enforcement: the decision layer that turns an entitlement into an answer. + * + * Four claims carry this suite, and each of them is the kind that decays + * silently if nothing asserts it: + * + * - **Inert by default.** With `BILLING_SERVICE_URL` unset nothing is fetched, + * the `USAGE_METER` binding is never touched, and no refusal can be produced. + * - **A refusal is a FAILING gate, never an absent one.** An exhausted allowance + * that skipped the LLM evaluator would stop a policy requiring AI review from + * requiring it, exactly when someone ran out of credit. + * - **Observe-only admits while recording.** Everything ships evaluating and + * logging decisions that bind only under `ENTITLEMENTS_ENFORCE=1`. + * - **The subject is the actor** (PRD §4a), an org pools only when positively + * known to be paid, and every usage read for a check is `source = 'platform'`. + */ +import { Hono } from "hono"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + checkGauge, + checkMeter, + resolveEnforcementSubject, + settleMeter, +} from "../src/billing/enforcement"; +import { type Entitlements, UNLIMITED, UnlimitedEntitlements } from "../src/billing/entitlements"; +import { runDeployMessage } from "../src/deploy/runner"; +import { LLMEvaluator } from "../src/evaluation/llm-evaluator"; +import type { LlmProvider } from "../src/evaluation/llm-provider"; +import type { EvalPolicy, EvaluationContext } from "../src/evaluation/types"; +import { rateLimitMiddleware } from "../src/middleware/rate-limit"; +import { recordCosts } from "../src/storage/costs"; +import { insertDeployment } from "../src/storage/deployments"; +import { setProject } from "../src/storage/state"; +import { getOwnerMeterTotals, upsertUsage } from "../src/storage/usage"; +import type { Env, ProjectEntry } from "../src/types"; +import type { Logger } from "../src/utils/logger"; +import { ok } from "../src/utils/result"; +import { makeFakeKV } from "./helpers/fake-kv"; +import { makeSqliteD1 } from "./helpers/sqlite-d1"; +import { makeUsageMeters } from "./helpers/usage-meter"; + +const PERIOD = "2026-09"; +const NOW = Date.UTC(2026, 8, 15, 12, 0, 0); + +const logged: Array<{ message: string; meta?: Record }> = []; +const logger: Logger = { + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn((message: string, meta?: Record) => { + logged.push({ message, ...(meta ? { meta } : {}) }); + }), + warn: vi.fn(), + error: vi.fn(), + fatal: vi.fn(), + child: vi.fn(() => logger), +} as unknown as Logger; + +type Meters = ReturnType; +type FakeKV = ReturnType; + +let db: D1Database; +let raw: ReturnType["raw"]; +let kv: FakeKV; +let meters: Meters; +let scheduled: Array>; + +/** Unique per test, so the module-level "already reconciled" memo cannot leak. */ +let seq = 0; +function nextUserId(): string { + seq += 1; + return `usr_${seq}_${Math.random().toString(36).slice(2, 8)}`; +} + +beforeEach(() => { + const made = makeSqliteD1(); + db = made.db; + raw = made.raw; + kv = makeFakeKV(); + meters = makeUsageMeters(); + scheduled = []; + logged.length = 0; + vi.clearAllMocks(); +}); + +function makeEnv(overrides: Partial = {}): Env { + return { + DB: db, + STATE: kv, + USAGE_METER: meters.namespace, + BILLING_SERVICE_URL: "https://billing.test", + BILLING_SERVICE_SECRET: "shh", + ...overrides, + } as unknown as Env; +} + +/** An env with the billing service unconfigured — every self-hoster. */ +function selfHostedEnv(overrides: Partial = {}): Env { + return makeEnv({ + BILLING_SERVICE_URL: undefined, + BILLING_SERVICE_SECRET: undefined, + ...overrides, + }); +} + +/** A plan with one monthly token allowance and everything else unlimited. */ +function tokenPlan(limit: number): Entitlements { + return plan({ meters: { ...UnlimitedEntitlements.meters, llm_tokens_month: limit } }); +} + +/** The plan a subject that has spent its whole monthly token allowance is on. */ +function exhaustedTokens(): Entitlements { + return tokenPlan(0); +} + +function plan(overrides: Partial = {}): Entitlements { + return { + ...UnlimitedEntitlements, + plan: "free", + meters: { ...UnlimitedEntitlements.meters }, + counts: { ...UnlimitedEntitlements.counts }, + rates: { ...UnlimitedEntitlements.rates }, + ...overrides, + }; +} + +/** Seed the entitlements cache the way a warm would, so `forOwner` hits. */ +function cachePlan(ownerType: "user" | "org", ownerId: string, entitlements: Entitlements): void { + kv.store.set(`entitlements:v2:${ownerType}:${ownerId}`, JSON.stringify(entitlements)); +} + +function createUser(id: string): string { + raw + .prepare("INSERT INTO users (id, email, token_hash) VALUES (?, ?, ?)") + .run(id, `${id}@example.test`, `hash-${id}`); + return id; +} + +function llmProvider(text = JSON.stringify({ score: 1, passed: true, reason: "fine" })) { + return { + run: vi.fn(async () => ok({ text, usage: { inputTokens: 100, outputTokens: 20 } })), + } as unknown as LlmProvider & { run: ReturnType }; +} + +const POLICY: EvalPolicy = { evaluators: [{ type: "llm" }] }; + +function billingContext(ownerId: string, actorUserId?: string): EvaluationContext { + return { + billing: { + ownerId, + ownerType: "user", + projectId: "prj_1", + ...(actorUserId ? { actorUserId } : {}), + }, + }; +} + +describe("inert with BILLING_SERVICE_URL unset (PRD Goal 3)", () => { + it("never touches the meter binding from the LLM gate", async () => { + const userId = nextUserId(); + // A plan that would refuse everything, cached — proving the short-circuit is + // the missing service and not the absence of a limit. + cachePlan("user", userId, exhaustedTokens()); + const provider = llmProvider(); + + const result = await new LLMEvaluator(provider, "platform", { + env: selfHostedEnv(), + }).evaluate("diff", POLICY, logger, billingContext(userId, userId)); + + expect(result.success && result.data.passed).toBe(true); + expect(provider.run).toHaveBeenCalledTimes(1); + expect(meters.calls).toEqual([]); + }); + + it("never touches the meter binding from a direct check", async () => { + const decision = await checkMeter(selfHostedEnv(), logger, { + subject: { ownerId: nextUserId(), ownerType: "user", pooled: false }, + meter: "llm_tokens_month", + estimate: 10_000, + nowMs: NOW, + period: PERIOD, + what: "AI review", + }); + + expect(decision).toMatchObject({ admitted: true, refused: false, checked: false }); + expect(meters.calls).toEqual([]); + }); +}); + +describe("the enforcement subject (PRD §4a)", () => { + it("charges the acting user, not the org the project is recorded against", async () => { + const actor = nextUserId(); + const subject = await resolveEnforcementSubject(makeEnv(), logger, { + actorUserId: actor, + owner: { ownerId: "org_1", ownerType: "org" }, + }); + + expect(subject).toEqual({ ownerId: actor, ownerType: "user", pooled: false }); + }); + + it("treats an UNCACHED org plan as free and charges the actor", async () => { + const actor = nextUserId(); + // Nothing cached for the org: the first check for any org is always a miss, + // because nothing warms org entitlements before this point. Reading that as + // "the org is its own subject" is the hole this direction closes. + const subject = await resolveEnforcementSubject(makeEnv(), logger, { + actorUserId: actor, + owner: { ownerId: "org_cold", ownerType: "org" }, + waitUntil: (promise) => scheduled.push(promise), + }); + + expect(subject?.ownerId).toBe(actor); + // ...and the warm was scheduled, so a paid org can pool from next time. + expect(scheduled.length).toBeGreaterThan(0); + await Promise.allSettled(scheduled); + }); + + it("treats a cached org plan that does not pool as free", async () => { + const actor = nextUserId(); + cachePlan("org", "org_solo", plan({ pooled: false })); + + const subject = await resolveEnforcementSubject(makeEnv(), logger, { + actorUserId: actor, + owner: { ownerId: "org_solo", ownerType: "org" }, + }); + + expect(subject?.ownerId).toBe(actor); + }); + + it("pools onto an org that is positively known to pool", async () => { + cachePlan("org", "org_paid", plan({ plan: "team", pooled: true })); + + const subject = await resolveEnforcementSubject(makeEnv(), logger, { + actorUserId: nextUserId(), + owner: { ownerId: "org_paid", ownerType: "org" }, + }); + + expect(subject).toEqual({ ownerId: "org_paid", ownerType: "org", pooled: true }); + }); + + it("walks an agent to its owning user when no actor is named", async () => { + const owner = createUser(nextUserId()); + raw + .prepare("INSERT INTO agents (id, name, owner_id, token_hash) VALUES (?, ?, ?, ?)") + .run("agt_1", "bot", owner, "hash"); + + const subject = await resolveEnforcementSubject(makeEnv(), logger, { + owner: { ownerId: "agt_1", ownerType: "agent" }, + }); + + expect(subject).toEqual({ ownerId: owner, ownerType: "user", pooled: false }); + }); + + it("names nobody when there is neither an actor nor a resolvable owner", async () => { + expect(await resolveEnforcementSubject(makeEnv(), logger, {})).toBeNull(); + }); +}); + +describe("reserve and settle", () => { + it("reserves the whole bound and settles the difference back", async () => { + const userId = nextUserId(); + const subject = { ownerId: userId, ownerType: "user" as const, pooled: false }; + cachePlan("user", userId, tokenPlan(10_000)); + + const decision = await checkMeter(makeEnv(), logger, { + subject, + meter: "llm_tokens_month", + estimate: 7_000, + nowMs: NOW, + period: PERIOD, + what: "AI review", + }); + expect(decision).toMatchObject({ admitted: true, reserved: true, count: 7_000 }); + + await settleMeter(makeEnv(), logger, { + subject, + meter: "llm_tokens_month", + delta: 900 - 7_000, + nowMs: NOW, + period: PERIOD, + }); + + const stub = meters.namespace.get(meters.namespace.idFromName(`user:${userId}`)); + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(900); + }); + + it("settles in a finally, so a thrown provider call cannot leak the reservation", async () => { + const userId = nextUserId(); + cachePlan("user", userId, tokenPlan(100_000)); + const provider = { + run: vi.fn(async () => { + throw new Error("connection reset"); + }), + } as unknown as LlmProvider; + + const result = await new LLMEvaluator(provider, "platform", { env: makeEnv() }).evaluate( + "diff", + POLICY, + logger, + billingContext(userId, userId), + ); + + expect(result.success).toBe(false); + const stub = meters.namespace.get(meters.namespace.idFromName(`user:${userId}`)); + // The whole reservation came back: nothing was spent, so nothing is charged. + expect((await stub.read(PERIOD)).counts.llm_tokens_month).toBe(0); + }); +}); + +describe("the LLM gate under an exhausted allowance", () => { + it("returns a FAILING gate naming cause and remedies, and never calls the provider", async () => { + const userId = nextUserId(); + cachePlan("user", userId, exhaustedTokens()); + const provider = llmProvider(); + + const result = await new LLMEvaluator(provider, "platform", { + env: makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial), + }).evaluate("diff", POLICY, logger, billingContext(userId, userId)); + + expect(result.success).toBe(true); + if (!result.success) return; + // A gate that ran and failed — NOT a gate that vanished. The evaluator is + // still in the set and still reports a verdict. + expect(result.data.passed).toBe(false); + expect(result.data.score).toBe(0); + expect(result.data.reason).toContain("monthly LLM token allowance"); + // Both remedies, per PRD §4c: bring your own key, or raise the plan. + expect(result.data.reason).toContain("provider:"); + expect(result.data.reason).toContain("plan limits"); + // ...and when it resets. + expect(result.data.reason).toContain("2026-10-01"); + expect(provider.run).not.toHaveBeenCalled(); + }); + + it("admits while recording the decision when enforcement is off", async () => { + const userId = nextUserId(); + cachePlan("user", userId, exhaustedTokens()); + const provider = llmProvider(); + + const result = await new LLMEvaluator(provider, "platform", { env: makeEnv() }).evaluate( + "diff", + POLICY, + logger, + billingContext(userId, userId), + ); + + expect(result.success && result.data.passed).toBe(true); + expect(provider.run).toHaveBeenCalledTimes(1); + // Recorded even though it admitted: a month of observation that logs nothing + // measures nothing. + const decision = logged.find( + (entry) => entry.message === "Entitlement decision" && entry.meta?.refused === true, + ); + expect(decision?.meta).toMatchObject({ enforcing: false, admitted: true, limit: 0 }); + }); + + it("refuses on the hourly evaluation rate even under BYOK (PRD §4b)", async () => { + const userId = nextUserId(); + cachePlan( + "user", + userId, + plan({ rates: { requests_per_minute: UNLIMITED, evaluations_per_hour: 0 } }), + ); + const provider = llmProvider(); + + const result = await new LLMEvaluator(provider, "byok", { + env: makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial), + }).evaluate("diff", POLICY, logger, billingContext(userId, userId)); + + expect(result.success && result.data.passed).toBe(false); + expect(result.success && result.data.reason).toContain("hourly evaluation allowance"); + // The string says so out loud, because it is the remedy an agent reaches for. + expect(result.success && result.data.reason).toContain("does NOT lift this one"); + expect(provider.run).not.toHaveBeenCalled(); + }); + + it("does not charge a BYOK run against the hosted token allowance", async () => { + const userId = nextUserId(); + cachePlan("user", userId, exhaustedTokens()); + const provider = llmProvider(); + + const result = await new LLMEvaluator(provider, "byok", { + env: makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial), + }).evaluate("diff", POLICY, logger, billingContext(userId, userId)); + + // The token allowance is the one thing BYOK lifts: the project is paying. + expect(result.success && result.data.passed).toBe(true); + expect(provider.run).toHaveBeenCalledTimes(1); + }); +}); + +describe("usage reads for a limit check (PRD §4a, Goal 5)", () => { + it("sums platform rows only, so BYOK spend is not charged to the hosted allowance", async () => { + const userId = nextUserId(); + const subject = { ownerId: userId, ownerType: "user" as const, pooled: false }; + await upsertUsage(db, logger, subject, PERIOD, [ + { meter: "llm_tokens_month", quantity: 600, source: "platform" }, + { meter: "llm_tokens_month", quantity: 5_000, source: "byok" }, + ]); + + const totals = await getOwnerMeterTotals(db, logger, userId, PERIOD, "platform"); + expect(totals.success && totals.data.llm_tokens_month).toBe(600); + + // And the counter is seeded from that same filtered total, so a project on + // its own key cannot exhaust the hosted meter. + cachePlan("user", userId, tokenPlan(10_000)); + const decision = await checkMeter(makeEnv(), logger, { + subject, + meter: "llm_tokens_month", + estimate: 100, + nowMs: NOW, + period: PERIOD, + what: "AI review", + }); + expect(decision.count).toBe(700); + }); +}); + +describe("the private project gauge", () => { + it("refuses one over the limit and admits under it", async () => { + const userId = nextUserId(); + cachePlan("user", userId, plan({ counts: { private_projects: 2 } })); + const env = makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial); + const subject = { ownerId: userId, ownerType: "user" as const, pooled: false }; + + const under = await checkGauge(env, logger, { + subject, + count: "private_projects", + current: 1, + what: "Creating another private project", + }); + const over = await checkGauge(env, logger, { + subject, + count: "private_projects", + current: 2, + what: "Creating another private project", + }); + + expect(under.admitted).toBe(true); + expect(over.admitted).toBe(false); + expect(over.reason).toContain("plan allows 2"); + }); +}); + +describe("rateLimitMiddleware's default limit", () => { + function app(env: Env, userId?: string) { + const application = new Hono<{ Bindings: Env }>(); + application.use("*", async (c, next) => { + if (userId) c.set("userId", userId); + await next(); + }); + application.use("*", rateLimitMiddleware()); + application.get("/x", (c) => c.text("ok")); + return (path = "/x") => application.fetch(new Request(`https://t.test${path}`), env); + } + + it("falls back to today's 1000 when the entitlements cache misses", async () => { + const response = await app(makeEnv(), nextUserId())(); + expect(response.headers.get("X-RateLimit-Limit")).toBe("1000"); + }); + + it("keeps today's numbers entirely when the billing service is unconfigured", async () => { + const response = await app(selfHostedEnv(), nextUserId())(); + expect(response.headers.get("X-RateLimit-Limit")).toBe("1000"); + }); + + it("uses a cached plan's own rate when enforcement is on", async () => { + const userId = nextUserId(); + cachePlan( + "user", + userId, + plan({ rates: { requests_per_minute: 25, evaluations_per_hour: 10 } }), + ); + const response = await app(makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial), userId)(); + expect(response.headers.get("X-RateLimit-Limit")).toBe("25"); + }); + + it("never tightens while observing: a lower plan limit does not bind", async () => { + const userId = nextUserId(); + cachePlan( + "user", + userId, + plan({ rates: { requests_per_minute: 25, evaluations_per_hour: 10 } }), + ); + const response = await app(makeEnv(), userId)(); + expect(response.headers.get("X-RateLimit-Limit")).toBe("1000"); + }); + + it("leaves an agent on its own bucket at today's numbers", async () => { + const application = new Hono<{ Bindings: Env }>(); + application.use("*", async (c, next) => { + c.set("agentId", "agt_1"); + await next(); + }); + application.use("*", rateLimitMiddleware()); + application.get("/x", (c) => c.text("ok")); + const response = await application.fetch( + new Request("https://t.test/x"), + makeEnv({ ENTITLEMENTS_ENFORCE: "1" } as Partial), + ); + expect(response.headers.get("X-RateLimit-Limit")).toBe("1000"); + }); +}); + +describe("the deploy volume check", () => { + const PROJECT_ID = "prj_deploy_enforce"; + const REMOTE = "https://acct.artifacts.cloudflare.net/git/acct/api.git"; + const SHA = "a".repeat(40); + + async function deployProject(ownerId: string): Promise { + const project: ProjectEntry = { + id: PROJECT_ID, + name: "api", + slug: "api", + namespace: "@alice", + ownerId, + ownerType: "user", + remote: REMOTE, + createdAt: "2026-09-01T00:00:00.000Z", + }; + const stored = await setProject(kv, project, logger); + if (!stored.success) throw stored.error; + return project; + } + + it("writes a persisted failed row with a named reason, and acks", async () => { + const userId = createUser(nextUserId()); + await deployProject(userId); + cachePlan( + "user", + userId, + plan({ meters: { ...UnlimitedEntitlements.meters, deploys_month: 0 } }), + ); + + const inserted = await insertDeployment(db, logger, { + projectId: PROJECT_ID, + project: "api", + commitSha: SHA, + name: "production", + target: "vercel", + requestedByType: "user", + requestedById: userId, + now: "2026-09-15T00:00:00.000Z", + }); + if (!inserted.success || !inserted.data.inserted) throw new Error("fixture insert failed"); + + const result = await runDeployMessage( + makeEnv({ + ENTITLEMENTS_ENFORCE: "1", + ARTIFACTS: { + get: async () => ({ createToken: async () => ({ plaintext: "repo-token" }) }), + }, + } as unknown as Partial), + { kind: "deployment", projectId: PROJECT_ID, deploymentId: inserted.data.deployment.id }, + logger, + { + now: () => NOW, + readFiles: async () => + ok( + new Map([ + [ + ".stratum/policy.yaml", + new TextEncoder().encode( + "evaluators: []\ndeploys:\n - name: production\n target: vercel\n", + ), + ], + ]), + ), + // A provider call would be a bug: the refusal happens before the claim. + fetch: vi.fn(async () => new Response("{}", { status: 200 })), + }, + ); + + // Acked, not retried: redelivery cannot change a limit. + expect(result.success).toBe(true); + const row = raw + .prepare("SELECT status, reason FROM deployments WHERE id = ?") + .get(inserted.data.deployment.id) as { status: string; reason: string }; + expect(row.status).toBe("failed"); + expect(row.reason).toContain("monthly deploy allowance"); + }); +}); + +describe("the 80% threshold notification", () => { + function envWithEmail(): { env: Env; send: ReturnType } { + const send = vi.fn(async () => undefined); + return { + env: makeEnv({ + EMAIL: { send } as unknown as Env["EMAIL"], + EMAIL_FROM_ADDRESS: "noreply@stratum.test", + }), + send, + }; + } + + async function spend(env: Env, ownerId: string, quantity: number): Promise { + await recordCosts( + db, + logger, + { + project: "api", + projectId: "prj_1", + ownerId, + ownerType: "user", + notify: { + env, + actorUserId: ownerId, + waitUntil: (promise) => scheduled.push(promise), + }, + }, + [{ kind: "llm_tokens", quantity }], + ); + await Promise.allSettled(scheduled); + scheduled.length = 0; + } + + it("fires once on the crossing and not again inside the same period", async () => { + const userId = createUser(nextUserId()); + const { env, send } = envWithEmail(); + cachePlan("user", userId, tokenPlan(1_000)); + + // 700 of 1000 is under 80%: approaching is not crossing. + await spend(env, userId, 700); + expect(send).not.toHaveBeenCalled(); + + // 850 crosses. + await spend(env, userId, 150); + expect(send).toHaveBeenCalledTimes(1); + expect(send.mock.calls[0]?.[0]).toMatchObject({ to: `${userId}@example.test` }); + + // Still past the threshold, but the edge has already been reported. + await spend(env, userId, 50); + expect(send).toHaveBeenCalledTimes(1); + }); + + it("re-arms when the limit is raised mid-period: the limit is in the receipt key", async () => { + const userId = createUser(nextUserId()); + const { env, send } = envWithEmail(); + cachePlan("user", userId, tokenPlan(1_000)); + + await spend(env, userId, 850); + expect(send).toHaveBeenCalledTimes(1); + + // The plan is upgraded: 850 is now well under 80% of 2000, so the next + // crossing is a real one and must not be silenced by the first receipt. + cachePlan("user", userId, tokenPlan(2_000)); + await spend(env, userId, 800); + expect(send).toHaveBeenCalledTimes(2); + }); + + it("ignores BYOK spend, which is not under the allowance at all", async () => { + const userId = createUser(nextUserId()); + const { env, send } = envWithEmail(); + cachePlan("user", userId, tokenPlan(1_000)); + + await recordCosts( + db, + logger, + { + project: "api", + ownerId: userId, + ownerType: "user", + notify: { env, actorUserId: userId, waitUntil: (promise) => scheduled.push(promise) }, + }, + [{ kind: "llm_tokens", quantity: 950, source: "byok" }], + ); + await Promise.allSettled(scheduled); + + expect(send).not.toHaveBeenCalled(); + }); + + it("says nothing at all when the billing service is unconfigured", async () => { + const userId = createUser(nextUserId()); + const send = vi.fn(async () => undefined); + const env = selfHostedEnv({ + EMAIL: { send } as unknown as Env["EMAIL"], + EMAIL_FROM_ADDRESS: "noreply@stratum.test", + }); + cachePlan("user", userId, tokenPlan(1_000)); + + await spend(env, userId, 950); + expect(send).not.toHaveBeenCalled(); + }); +}); diff --git a/tests/entitlements.test.ts b/tests/entitlements.test.ts index 9dac847f..c324ae15 100644 --- a/tests/entitlements.test.ts +++ b/tests/entitlements.test.ts @@ -40,15 +40,16 @@ function unconfiguredEnv(): Env { return env({ BILLING_SERVICE_URL: undefined, BILLING_SERVICE_SECRET: undefined }).env; } -const CACHE_KEY = "entitlements:v1:user:usr_1"; -const NEGATIVE_KEY = "entitlements:v1:neg:user:usr_1"; -const LOCK_KEY = "entitlements:v1:lock:user:usr_1"; +const CACHE_KEY = "entitlements:v2:user:usr_1"; +const NEGATIVE_KEY = "entitlements:v2:neg:user:usr_1"; +const LOCK_KEY = "entitlements:v2:lock:user:usr_1"; const paidPlan: Entitlements = { plan: "pro", + pooled: true, meters: { llm_tokens_month: 1_000_000, sandbox_ms_month: 60_000, deploys_month: 100 }, counts: { private_projects: 25 }, - rates: { requests_per_minute: 2_000 }, + rates: { requests_per_minute: 2_000, evaluations_per_hour: 200 }, }; function jsonResponse(body: unknown, status = 200): Response { @@ -159,10 +160,12 @@ describe("RemoteEntitlements.refresh", () => { expect(result.data.source).toBe("remote"); expect(result.data.entitlements).toEqual({ plan: "starter", + // A payload that says nothing about pooling is not a payload that pools. + pooled: false, // 0 survives: it is a hard block, not "unset". meters: { llm_tokens_month: 0, sandbox_ms_month: -1, deploys_month: -1 }, counts: { private_projects: -1 }, - rates: { requests_per_minute: -1 }, + rates: { requests_per_minute: -1, evaluations_per_hour: -1 }, }); }); diff --git a/tests/evaluator-billing-context.test.ts b/tests/evaluator-billing-context.test.ts index 29b2c5ab..939d9167 100644 --- a/tests/evaluator-billing-context.test.ts +++ b/tests/evaluator-billing-context.test.ts @@ -450,16 +450,25 @@ describe("createChangeWithEvaluation passes the project through", () => { expect(outcome.success).toBe(true); expect(secretScanContext()).toEqual({ baseSha: "base_from_diff_clone", - billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + billing: { + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + actorUserId: "user_alice", + }, }); }); - it("bills an org-owned project's org", async () => { + it("bills an org-owned project's org, and still names the actor", async () => { await create(projectEntry({ ownerId: "org_acme", ownerType: "org" })); + // The recorded subject is the org and the ACTOR is carried beside it: PRD + // §4a checks the limit against the person, while the ledger keeps naming + // the org. Collapsing the two is the hole that separation closes. expect(secretScanContext()?.billing).toEqual({ ownerId: "org_acme", ownerType: "org", projectId: "proj_abc", + actorUserId: "user_alice", }); }); @@ -479,6 +488,7 @@ describe("createChangeWithEvaluation passes the project through", () => { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc", + actorUserId: "user_alice", }); }); }); @@ -514,7 +524,12 @@ describe("POST /api/changes/:id/evaluate passes the project through", () => { // and no project id. The evaluators now see the loaded entry. expect(secretScanContext()).toEqual({ baseSha: "base_from_diff_clone", - billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + billing: { + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + actorUserId: "user_alice", + }, }); }); }); @@ -569,7 +584,12 @@ describe("POST /api/projects/conflicts/:id/resolve passes the project through", expect(res.status).toBe(200); expect(secretScanContext()).toEqual({ baseSha: "conflict_base_sha", - billing: { ownerId: "user_alice", ownerType: "user", projectId: "proj_abc" }, + billing: { + ownerId: "user_alice", + ownerType: "user", + projectId: "proj_abc", + actorUserId: "user_alice", + }, }); }); }); diff --git a/tests/llm-byok.test.ts b/tests/llm-byok.test.ts index 68c23fa8..e53269dc 100644 --- a/tests/llm-byok.test.ts +++ b/tests/llm-byok.test.ts @@ -42,7 +42,9 @@ const logger: Logger = { const ANTHROPIC_ENTRY = { name: "anthropic", kind: "anthropic", - baseUrl: "https://api.anthropic.com", + // With the "/v1": the provider POSTs `${baseUrl}/messages`, so the version-less + // host every example used to show 404s on every BYOK call. + baseUrl: "https://api.anthropic.com/v1", }; const PROVIDERS_JSON = JSON.stringify([ANTHROPIC_ENTRY]); @@ -93,7 +95,7 @@ describe("parseLlmProviders — the allowlist is the only place an endpoint is n expect(parse.providers.get("anthropic")).toEqual({ name: "anthropic", kind: "anthropic", - baseUrl: "https://api.anthropic.com", + baseUrl: "https://api.anthropic.com/v1", }); }); @@ -125,8 +127,18 @@ describe("parseLlmProviders — the allowlist is the only place an endpoint is n ["https://172.31.255.254/v1", "172.16.0.0/12"], ["https://192.168.1.1/v1", "192.168.0.0/16"], ["https://[fe80::1]/v1", "link-local"], + // The rest of fe80::/10, which a literal "fe80" prefix test let straight through. + ["https://[feb0::1]/v1", "link-local"], ["https://[fd00::1]/v1", "unique-local"], ["https://[::ffff:127.0.0.1]/v1", "IPv4-mapped"], + // IPv4-compatible IPv6: the URL parser rewrites this to [::7f00:1], which is + // neither "::1" nor "::ffff:" and was accepted. + ["https://[::127.0.0.1]/v1", "reserved ::/64"], + // CGNAT — carrier space, and reachable from plenty of private networks. + ["https://100.64.0.1/v1", "100.64.0.0/10"], + // The GCP metadata endpoint is a NAME, not an address: no IP check sees it. + ["https://metadata.google.internal/computeMetadata/v1", ".internal"], + ["https://printer.local/v1", ".local"], ])("rejects %s at parse time", (baseUrl, expected) => { const parse = parseLlmProviders(JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl }])); expect(parse.status, `${baseUrl} was accepted`).toBe("invalid"); @@ -142,6 +154,23 @@ describe("parseLlmProviders — the allowlist is the only place an endpoint is n } }); + it("stores the baseUrl normalized, so a bare '#' or '?' cannot swallow the path", () => { + // Both pass the empty-search and empty-hash checks — the parser reports both + // as "" — while the raw bytes would put `/messages` inside a fragment or a + // query, sending the prompt to the wrong path with the credential attached. + for (const [written, stored] of [ + ["https://api.anthropic.com/v1#", "https://api.anthropic.com/v1"], + ["https://api.anthropic.com/v1?", "https://api.anthropic.com/v1"], + ["https://api.anthropic.com/v1/", "https://api.anthropic.com/v1"], + ["https://api.anthropic.com", "https://api.anthropic.com"], + ]) { + const parse = parseLlmProviders(JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl: written }])); + expect(parse.status, `${written} was rejected`).toBe("ok"); + if (parse.status !== "ok") return; + expect(parse.providers.get("anthropic")?.baseUrl).toBe(stored); + } + }); + it("rejects credentials embedded in the URL", () => { const parse = parseLlmProviders( JSON.stringify([{ ...ANTHROPIC_ENTRY, baseUrl: "https://user:pass@api.anthropic.com" }]), @@ -227,19 +256,82 @@ describe("sanitizeLlmConfig — the policy may select a provider, never describe } it("keeps a provider the operator configured", () => { - const parsed = parsePolicy({ evaluators: [{ type: "llm", provider: "anthropic" }] }); + const parsed = parsePolicy({ + evaluators: [{ type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }], + }); expect(parsed.status).toBe("ok"); if (parsed.status !== "ok") return; - expect(parsed.policy.evaluators[0]).toEqual({ type: "llm", provider: "anthropic" }); + expect(parsed.policy.evaluators[0]).toEqual({ + type: "llm", + provider: "anthropic", + model: "claude-sonnet-4-5", + }); expect(parsed.policy.configError).toBeUndefined(); }); + it("fails closed on a provider entry that names no model", () => { + // The default model id is a Workers AI one (`@cf/meta/llama-3.1-8b-instruct`). + // Posted to Anthropic it fails every call — a gate that never runs, found one + // merge at a time — so the entry is refused where it is written instead. + const parsed = parsePolicy({ evaluators: [{ type: "llm", provider: "anthropic" }] }); + expect(parsed.status).toBe("malformed"); + if (parsed.status !== "malformed") return; + expect(parsed.reason).toContain('no "model"'); + // A policy with no provider keeps the Workers AI default, where the id is right. + expect(parsePolicy({ evaluators: [{ type: "llm" }] }).status).toBe("ok"); + }); + + it("names the provider in the reason that blocks the merge", () => { + // A bare count ("1 unusable entry in 'evaluators'") is what the person + // fixing the file sees; it has to say which entry and why. + const parsed = parsePolicy({ + evaluators: [{ type: "llm", provider: "evil-corp", model: "m" }], + }); + expect(parsed.status).toBe("malformed"); + if (parsed.status !== "malformed") return; + expect(parsed.reason).toContain("evil-corp"); + expect(parsed.reason).toContain("LLM_PROVIDERS"); + }); + + it("rejects a policy declaring more than one llm entry", () => { + // The fail-open this closes: `LLMEvaluator` reads the FIRST entry while + // `buildEvaluators` builds one per entry, so this pair used to yield two + // evaluators on the operator's Workers AI bill from a policy that asked to + // run on the project's own key — no error, no log. + const parsed = parsePolicy({ + evaluators: [ + { type: "llm" }, + { type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }, + ], + }); + expect(parsed.status).toBe("malformed"); + if (parsed.status !== "malformed") return; + expect(parsed.reason).toContain("'llm' entries"); + }); + + it("caps how many evaluator entries one policy file may declare", () => { + // The `deploys:` sibling has had this cap from the start, and an evaluator + // entry is an outbound call: without it one merge amplifies into unbounded + // provider requests. + const parsed = parsePolicy({ + evaluators: Array.from({ length: 17 }, () => ({ type: "diff" })), + }); + expect(parsed.status).toBe("malformed"); + if (parsed.status !== "malformed") return; + expect(parsed.reason).toContain("at most 16"); + expect( + parsePolicy({ evaluators: Array.from({ length: 16 }, () => ({ type: "diff" })) }).status, + ).toBe("ok"); + }); + it("fails the whole policy CLOSED on a provider name the operator has not configured", () => { // The evaluator-side convention (`policy-loader.ts`'s dropped-entry check), // deliberately opposite to a rejected `deploys:` entry: an unresolvable // provider means the gate cannot run, and a gate that cannot run must not // pass. - const parsed = parsePolicy({ evaluators: [{ type: "llm", provider: "evil-corp" }] }); + const parsed = parsePolicy({ + evaluators: [{ type: "llm", provider: "evil-corp", model: "m" }], + }); expect(parsed.status).toBe("malformed"); if (parsed.status !== "malformed") return; expect(parsed.reason).toContain("unusable entr"); @@ -251,7 +343,7 @@ describe("sanitizeLlmConfig — the policy may select a provider, never describe // silently dropped the entry would let the change through on its remaining // evaluators. const parsed = parsePolicy({ - evaluators: [{ type: "llm", provider: "evil-corp" }, { type: "diff" }], + evaluators: [{ type: "llm", provider: "evil-corp", model: "m" }, { type: "diff" }], }); expect(parsed.status).toBe("malformed"); }); @@ -260,7 +352,9 @@ describe("sanitizeLlmConfig — the policy may select a provider, never describe // Not a special case, and it must not be: an operator who never set // LLM_PROVIDERS has an empty catalog, which resolves no name. const parsed = parsePolicyContent( - JSON.stringify({ evaluators: [{ type: "llm", provider: "anthropic" }] }), + JSON.stringify({ + evaluators: [{ type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }], + }), "json", logger, catalogFrom(undefined), @@ -275,6 +369,7 @@ describe("sanitizeLlmConfig — the policy may select a provider, never describe { type: "llm", provider: "anthropic", + model: "claude-sonnet-4-5", baseUrl: "https://evil.example.com", base_url: "https://evil.example.com", endpoint: "https://evil.example.com", @@ -282,16 +377,23 @@ describe("sanitizeLlmConfig — the policy may select a provider, never describe logger, providers(), ); - expect(config).toEqual({ type: "llm", provider: "anthropic" }); + expect(config).toEqual({ type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }); expect(JSON.stringify(config)).not.toContain("evil.example.com"); }); it("drops an inline credential a policy tried to smuggle in", () => { const config = sanitizeLlmConfig( - { type: "llm", provider: "anthropic", apiKey: "sk-ant-inline", secret: "s" }, + { + type: "llm", + provider: "anthropic", + model: "claude-sonnet-4-5", + apiKey: "sk-ant-inline", + secret: "s", + }, logger, providers(), ); + expect(config).not.toBeNull(); expect(JSON.stringify(config)).not.toContain("sk-ant-inline"); }); @@ -368,7 +470,9 @@ describe("resolveLlmProvider — the credential path", () => { return db; } - const byokPolicy: EvalPolicy = { evaluators: [{ type: "llm", provider: "anthropic" }] }; + const byokPolicy: EvalPolicy = { + evaluators: [{ type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }], + }; it("skips the secret load entirely when the policy names no provider", async () => { // The skip is the performance requirement, so it is asserted the only way @@ -496,7 +600,9 @@ describe("resolveLlmProvider — the credential path", () => { describe("buildEvaluators — BYOK never falls back to the operator's Workers AI", () => { const KEY = "test-deploy-secret-key-0123456789"; - const byokPolicy: EvalPolicy = { evaluators: [{ type: "llm", provider: "anthropic" }] }; + const byokPolicy: EvalPolicy = { + evaluators: [{ type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }], + }; function envWith(db: D1Database, ai: AiBinding, overrides: Partial = {}): Env { return { @@ -595,6 +701,76 @@ describe("buildEvaluators — BYOK never falls back to the operator's Workers AI expect(costs?.[0]).toMatchObject({ kind: "llm_tokens", estimated: true }); }); + it("posts to the URL the operator configured, path included", async () => { + // The existing BYOK fetch stubs ignore their arguments, which is how a + // version-less example `baseUrl` survived in wrangler.toml and in the Env + // docstring: `${baseUrl}/messages` against "https://api.anthropic.com" is a + // 404 on every call, and no test looked at the URL. + const { db } = makeSqliteD1(); + await putSecret( + db, + logger, + { DEPLOY_SECRET_KEY: KEY }, + { + projectId: project().id, + name: "ANTHROPIC_API_KEY", + value: "sk-ant-real", + actorId: "user_alice", + }, + ); + + const fetchMock = vi.fn(async () => ({ + ok: true, + status: 200, + text: async () => + JSON.stringify({ + content: [{ type: "text", text: '{"score":1,"passed":true,"reason":"ok"}' }], + usage: { input_tokens: 9, output_tokens: 1 }, + }), + })); + vi.stubGlobal("fetch", fetchMock); + + const ai: AiBinding = { run: vi.fn() } as unknown as AiBinding; + await evaluateWith(envWith(db, ai)); + + const [url, init] = fetchMock.mock.calls[0] as unknown as [ + string, + { headers: Headers; redirect?: string }, + ]; + expect(url).toBe("https://api.anthropic.com/v1/messages"); + expect(init.headers.get("x-api-key")).toBe("sk-ant-real"); + // The allowlist binds the host this was SENT to; a followed redirect would + // move it, with the body and that header attached. + expect(init.redirect).toBe("manual"); + }); + + it("fails the gate closed when the policy declares two llm entries", async () => { + // Reachable with a policy that never went through the parser — a KV-cached + // one written before the rule existed. Neither entry may quietly win. + const { db } = makeSqliteD1(); + const ai: AiBinding = { + run: vi.fn(async () => ({ response: '{"score":1,"passed":true,"reason":"ok"}' })), + } as unknown as AiBinding; + const twoEntries: EvalPolicy = { + evaluators: [ + { type: "llm" }, + { type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }, + ], + }; + + const built = await buildEvaluators(envWith(db, ai), twoEntries, project(), logger); + const { evalResult } = await runEvaluation( + built.filter((entry) => entry.type === "llm"), + "diff --git a/x b/x", + twoEntries, + logger, + ); + + expect(evalResult.passed).toBe(false); + expect(evalResult.reason).toContain("llm unavailable"); + expect(ai.run).not.toHaveBeenCalled(); + }); + it("keeps a provider error body — including an echoed API key — out of the result and the logs", async () => { // Some providers echo the credential that failed. `llm-provider.ts` returns // the status and never the body, and this is the assertion that keeps it @@ -685,7 +861,77 @@ describe("sanitizePolicy — the strip list, and MAX_POLICY_CONTEXT_CHARS", () = const policy = { evaluators: [{ type: "diff", monkeys: 3, keystone: "kept" }], } as unknown as EvalPolicy; - expect(JSON.stringify(sanitizePolicy(policy))).toContain("keystone"); + const serialized = JSON.stringify(sanitizePolicy(policy)); + expect(serialized).toContain("keystone"); + expect(serialized).toContain("monkeys"); + }); + + it("strips a credential written at the TOP level of the policy file", () => { + // `parsePolicyContent` spreads unknown root keys onto the policy, so this + // reached the model prompt AND the body POSTed to a policy-supplied webhook + // URL — a host the repository's own file chose. + const policy = { + evaluators: [{ type: "diff" }], + apiKey: "sk-root-leak", + registry: { url: "https://kept.example" }, + } as unknown as EvalPolicy; + + const serialized = JSON.stringify(sanitizePolicy(policy)); + expect(serialized).not.toContain("sk-root-leak"); + expect(serialized).not.toContain("apiKey"); + expect(serialized).toContain("https://kept.example"); + }); + + it("strips a credential NESTED inside an entry, not just a top-level one", () => { + // The copy-through case the strip list exists for: an unmodelled `type` is + // passed to the model whole, and one level of key filtering left everything + // inside an object or an array untouched. + const policy = { + evaluators: [ + { + type: "future_evaluator", + headers: { Authorization: "Bearer sk-nested-1", accept: "application/json" }, + targets: [{ name: "ci", token: "sk-nested-2" }], + deep: { a: { b: { apiKey: "sk-nested-3" } } }, + }, + ], + } as unknown as EvalPolicy; + + const serialized = JSON.stringify(sanitizePolicy(policy)); + for (let i = 1; i <= 3; i++) expect(serialized).not.toContain(`sk-nested-${i}`); + expect(serialized).toContain("application/json"); + expect(serialized).toContain('"name":"ci"'); + }); + + it.each([ + "auth", + "AUTH", + "tokens", + "apiKeys", + "access_tokens", + "passwords", + "pwd", + "pat", + "githubPat", + "hmac", + "sig", + ])("strips the credential-shaped key %s", (key) => { + const policy = { + evaluators: [{ type: "future_evaluator", [key]: "sk-word-leak", keystone: "kept" }], + } as unknown as EvalPolicy; + const serialized = JSON.stringify(sanitizePolicy(policy)); + expect(serialized).not.toContain("sk-word-leak"); + // The property the stem test exists to keep: innocuous words survive. + expect(serialized).toContain("keystone"); + }); + + it("terminates on a cyclic policy rather than recursing forever", () => { + // A YAML alias cycle (`a: &x {b: *x}`) is a well-formed file and an infinite + // structure. The depth bound is the cycle guard. + const entry: Record = { type: "future_evaluator" }; + entry.self = entry; + const policy = { evaluators: [entry] } as unknown as EvalPolicy; + expect(() => JSON.stringify(sanitizePolicy(policy))).not.toThrow(); }); it("makes opting into BYOK free against the 8000-char policy context", () => { @@ -694,12 +940,15 @@ describe("sanitizePolicy — the strip list, and MAX_POLICY_CONTEXT_CHARS", () = // oversize one fails the gate CLOSED before any model call. Because // `provider` is stripped, adding it costs zero characters. const large: EvalPolicy = { - evaluators: [{ type: "diff", forbiddenPatterns: [`${"x".repeat(7_500)}`] }, { type: "llm" }], + evaluators: [ + { type: "diff", forbiddenPatterns: [`${"x".repeat(7_500)}`] }, + { type: "llm", model: "claude-sonnet-4-5" }, + ], }; const withProvider: EvalPolicy = { evaluators: [ { type: "diff", forbiddenPatterns: [`${"x".repeat(7_500)}`] }, - { type: "llm", provider: "anthropic" }, + { type: "llm", provider: "anthropic", model: "claude-sonnet-4-5" }, ], }; diff --git a/tests/llm-provider.test.ts b/tests/llm-provider.test.ts index 34eb6ef8..678399e7 100644 --- a/tests/llm-provider.test.ts +++ b/tests/llm-provider.test.ts @@ -58,9 +58,18 @@ function lastFetchCall() { const fetchMock = globalThis.fetch as unknown as ReturnType; const [url, init] = fetchMock.mock.calls[fetchMock.mock.calls.length - 1] as [ string, - { method: string; headers: Record; body: string }, + { method: string; headers: Headers; body: string; redirect?: string }, ]; - return { url, init, body: JSON.parse(init.body) as Record }; + // The provider builds a `Headers` object rather than a record — it has to, so + // that a value no header can carry fails where nothing interpolates it into a + // message. Flattened here (iteration lowercases the names) so the assertions + // below still read like the request. + const headers = Object.fromEntries(init.headers) as Record; + return { + url, + init: { ...init, headers }, + body: JSON.parse(init.body) as Record, + }; } beforeEach(() => { @@ -824,3 +833,129 @@ describe("LLMEvaluator — provider failures fail the gate closed", () => { expect(result.error.message).toContain("isolate exceeded CPU"); }); }); + +describe("HttpLlmProvider — the request is bound to the host the allowlist checked", () => { + it("refuses a redirect instead of following it", async () => { + // The allowlist validates the host a request is SENT to. Following a 3xx + // would re-send the prompt — the diff and the policy — to a host nobody + // validated, and on 307/308 the body goes with it; the Fetch spec strips + // `Authorization` cross-origin but not `x-api-key`, so the Anthropic path + // would hand the project's credential to whatever the redirect named. + const fetchMock = vi.fn().mockResolvedValue({ + ok: false, + status: 307, + text: async () => "", + }); + vi.stubGlobal("fetch", fetchMock); + + const result = await new AnthropicProvider("https://api.anthropic.com/v1", "sk-ant-secret").run( + "m", + [], + ); + + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + // Constant: a redirect target is attacker-chosen, so nothing about it is + // interpolated into a message that becomes a user-visible reason. + expect(result.error.message).toBe("provider redirected the request, which is not followed"); + + const [, init] = fetchMock.mock.calls[0] as [string, { redirect?: string }]; + expect(init.redirect).toBe("manual"); + }); + + it.each([301, 302, 303, 307, 308])("refuses HTTP %i the same way", async (status) => { + vi.stubGlobal("fetch", vi.fn().mockResolvedValue({ ok: false, status, text: async () => "" })); + const result = await new OpenAiCompatibleProvider("https://api.openai.com/v1", "k").run( + "m", + [], + ); + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.message).toBe("provider redirected the request, which is not followed"); + }); + + it("never puts a header value into an error message when the headers cannot be built", async () => { + // `new Headers()` throws a TypeError that QUOTES the offending value. The + // store rejects control characters so this should be unreachable, but the + // message is the one that becomes `EvalResult.reason` — persisted on the + // change and rendered on a world-readable page for a public project — so it + // is a constant, and the request is never made. + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + const result = await new AnthropicProvider( + "https://api.anthropic.com/v1", + "sk-ant-real\r\nX-Injected: 1", + ).run("m", []); + + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error).toBeInstanceOf(LlmProviderResponseError); + expect(result.error.message).toBe("provider request headers could not be constructed"); + expect(result.error.message).not.toContain("sk-ant-real"); + expect(JSON.stringify(result.error)).not.toContain("sk-ant-real"); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); + +describe("LLMEvaluator — a truncated verdict was still generated, and billed", () => { + it("records the Anthropic-reported counts for a response cut off at the cap", async () => { + // `failClosed` records the cost of a response it could not read; this is the + // same case one layer up. The provider ran the generation to its cap and + // charged for it, so recording zero would under-count a real spend. + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + content: [{ type: "text", text: '{"score":1,"passed"' }], + stop_reason: "max_tokens", + usage: { input_tokens: 900, output_tokens: 1024 }, + }), + ), + ); + + const result = await new LLMEvaluator( + new AnthropicProvider("https://api.anthropic.com/v1", "k"), + "byok", + ).evaluate("diff content", POLICY, mockLogger); + + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.passed).toBe(false); + expect(result.data.reason).toContain("truncated at the 1024-token limit"); + expect(result.data.costs).toEqual([{ kind: "llm_tokens", quantity: 1924, source: "byok" }]); + }); + + it("records the OpenAI-shaped counts too", async () => { + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue( + jsonResponse(200, { + choices: [{ message: { content: "{" }, finish_reason: "length" }], + usage: { prompt_tokens: 500, completion_tokens: 1024 }, + }), + ), + ); + + const result = await new LLMEvaluator( + new OpenAiCompatibleProvider("https://api.openai.com/v1", "k"), + ).evaluate("diff content", POLICY, mockLogger); + + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs).toEqual([{ kind: "llm_tokens", quantity: 1524 }]); + }); + + it("bills nothing for a failure that generated nothing", async () => { + // The distinction the branch turns on: a 401 produced no tokens, so it must + // not produce a cost sample either. + vi.stubGlobal("fetch", vi.fn().mockResolvedValue(jsonResponse(401, { error: "nope" }))); + const result = await new LLMEvaluator( + new AnthropicProvider("https://api.anthropic.com/v1", "k"), + ).evaluate("diff content", POLICY, mockLogger); + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.costs).toBeUndefined(); + }); +}); diff --git a/tests/project-secrets.test.ts b/tests/project-secrets.test.ts index 001ff1a8..63a656e6 100644 --- a/tests/project-secrets.test.ts +++ b/tests/project-secrets.test.ts @@ -297,6 +297,40 @@ describe("putSecret", () => { expect(result.success).toBe(false); }); + it.each([ + ["a carriage return", "sk-ant-real\rX-Injected: 1"], + ["a line feed", "sk-ant-real\nX-Injected: 1"], + ["a NUL", "sk-ant-real\u0000"], + ["a tab", "sk-ant-real\tmore"], + ])("rejects a value containing %s, without echoing it", async (_case, value) => { + // Not tidiness. A stored value goes into an outbound provider header, and + // `new Headers()` throws a TypeError that QUOTES the offending value — which + // would travel as an ExternalServiceError message into `EvalResult.reason`, + // be persisted on the change, and be rendered on a page that is + // world-readable for a public project. The store is the first of the two + // ends that close that path. + const result = await putSecret(db, logger, ENV, { + projectId: PROJECT, + name: "LEAKY", + value, + actorId: "usr_a", + }); + + expect(result.success).toBe(false); + if (result.success) return; + expect(result.error.code).toBe("VALIDATION_ERROR"); + // The refusal must not become the leak it prevents. + expect(result.error.message).not.toContain("sk-ant-real"); + expect(JSON.stringify(result.error.context ?? {})).not.toContain("sk-ant-real"); + + // …and nothing was written, so the value cannot be read back either. + const loaded = await loadSecretValues(db, logger, ENV, { + projectId: PROJECT, + names: ["LEAKY"], + }); + expect(loaded.success && loaded.data.missing).toEqual(["LEAKY"]); + }); + it("accepts an empty value: there is deliberately no minimum length", async () => { const result = await putSecret(db, logger, ENV, { projectId: PROJECT, diff --git a/website/src/content/docs/guides/getting-started.md b/website/src/content/docs/guides/getting-started.md index d1cc4aa9..a35d00f3 100644 --- a/website/src/content/docs/guides/getting-started.md +++ b/website/src/content/docs/guides/getting-started.md @@ -198,7 +198,14 @@ downgrade your governance. minimum passing score (0.0–1.0); `maxDiffChars` bounds how much diff is sent (default 24,000, max 100,000). An unavailable model or unparseable verdict fails closed. Token usage is recorded on the change as a cost - record. + record. **A policy may declare at most one `llm` entry** — a second one is a + merge-blocking policy error, because the two entries cannot both be the + configuration in force. + +A policy may declare **at most 16 `evaluators:` entries**; more than that is a +policy error and blocks merges, the same as any other unusable entry. (The +`deploys:` list has the same cap, for the same reason: one merge must not turn +into an unbounded number of external calls.) Two top-level knobs sit alongside `evaluators`: `requireAll` (default `true`) makes the aggregate verdict demand every evaluator pass — set it to `false` to diff --git a/wrangler.toml b/wrangler.toml index e09c2825..8d4a74a0 100644 --- a/wrangler.toml +++ b/wrangler.toml @@ -242,7 +242,10 @@ ENTITLEMENTS_ENFORCE = "0" # Parsed once per isolate by `parseLlmProviders` (src/evaluation/llm-providers.ts) # and validated there, not per request: `kind` must be "anthropic" or # "openai-compatible", `baseUrl` must be https with no embedded credentials, and -# a loopback, link-local or RFC1918 host is rejected. One bad entry rejects the +# a loopback, link-local, CGNAT, .internal/.local or RFC1918 host is rejected. +# `baseUrl` must include the API version path — the provider appends "/messages" +# (anthropic) or "/chat/completions" (openai-compatible) to it, so +# "https://api.anthropic.com" without the "/v1" 404s on every call. One bad entry rejects the # whole value — every project selecting a provider then blocks, projects on # Workers AI are unaffected, and `llmProvidersConfigError` logs the reason on the # first request after the deploy. @@ -255,7 +258,7 @@ ENTITLEMENTS_ENFORCE = "0" # every deploy secret undecryptable, and now blocks the merge gate on every BYOK # project until each re-enters its provider key. See the rotation warning in # docs/user-guide/deployments.md. -# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com/v1"}]' # Instance-wide telemetry kill switch. Not a secret — a plain var, listed here # next to POSTHOG_API_KEY because they are read together. Named envs do NOT @@ -419,6 +422,12 @@ GIT_PUSH_GATED_ENABLED = "false" # in the GitHub and Google OAuth app settings. OAUTH_REDIRECT_URI = "https://app.usestratum.dev/auth/github/callback" GOOGLE_REDIRECT_URI = "https://app.usestratum.dev/auth/google/callback" +# The From address every outbound email needs. Declared here and not only at top +# level for the usual reason (named envs inherit no vars) — and it is REQUIRED, +# not optional: every send path returns early without it, so production would +# otherwise ship the usage-threshold notice and the magic-link mail as features +# that silently never send, in the one environment that has limits at all. +EMAIL_FROM_ADDRESS = "noreply@app.usestratum.dev" # Open signup — the closed-beta invite gate is OFF in production ("available # today" is the positioning; Origin launched to every paid Cursor plan). # To re-enable the invite wall set BETA_GATE = "1"; the gate fails closed and @@ -442,7 +451,7 @@ ENTITLEMENTS_ENFORCE = "0" # what production runs on today. A project opts in by naming one of these in its # policy and storing its own key in that project's secrets. See the top-level # [vars] block for the validation rules and the DEPLOY_SECRET_KEY coupling. -# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com/v1"}]' # ============================================================================= # [env.staging] @@ -603,4 +612,4 @@ ENTITLEMENTS_ENFORCE = "0" # Operator-configured LLM provider allowlist; unset = Workers AI only. Staging is # where a new provider entry gets exercised end-to-end before production, since a # bad one is only visible as blocked merges on the projects selecting it. -# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com"}]' +# LLM_PROVIDERS = '[{"name":"anthropic","kind":"anthropic","baseUrl":"https://api.anthropic.com/v1"}]' From f79256b76bc3e108b6f742accbde7f471b87aec4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 12:25:16 +0000 Subject: [PATCH 14/26] Show people their usage, and let an agent ask for it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 9. One read behind three surfaces: a settings page, a banner, and an MCP tool. The read is the same one enforcement performs before a limit check — enforcement subject, platform rows only — so the page cannot disagree with the thing that refuses you. BYOK is a separate optional figure and is never folded into consumption, because spend on a project's own key is not under the hosted allowance. An unlimited entitlement says so in words rather than rendering an empty bar; a self-hoster sees this page too and everything is unlimited for them, so the page has to read sensibly rather than look broken. The banner costs one KV get per render and no D1. Nothing is computed at render time: the notification work already detects the 80% crossing on the write that causes it, from the totals the upsert returns, so this stores what it found and the page reads it back. It short-circuits before touching KV for signed-out visitors and for self-hosters, which the tests assert with a spy rather than by inspection. It renders on the pages people actually visit rather than only on the usage page, which would warn only the people who already went looking. stratum_get_usage is read-only, and its description says so — an agent that can raise its own limit does not have one, the same line this codebase already draws by refusing agent tokens on review verdicts. It wraps a REST endpoint rather than reading storage directly, because every other tool does and a tool with its own storage path is the parallel implementation the dispatcher exists to prevent. An agent token there resolves to its owner: the subject its evaluations are actually charged to. Typecheck clean; lint clean; 4077 tests pass, up 17. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Got86FaNHpxB6EDtp42gpX --- docs/developer/architecture.md | 2 +- src/billing/usage-banner.ts | 146 ++++++++ src/billing/usage-notifications.ts | 15 + src/billing/usage-report.ts | 206 +++++++++++ src/mcp/client.ts | 11 + src/mcp/tools.ts | 18 +- src/routes/ui.tsx | 65 +++- src/routes/users.ts | 36 ++ src/ui/layout.tsx | 48 ++- src/ui/pages/change-detail.tsx | 5 + src/ui/pages/changes.tsx | 13 +- src/ui/pages/home.tsx | 7 +- src/ui/pages/settings.tsx | 11 +- src/ui/pages/usage.tsx | 165 +++++++++ src/ui/styles.ts | 25 ++ tests/mcp-endpoint.test.ts | 4 +- tests/mcp-schema.test.ts | 2 +- tests/usage-page.test.tsx | 72 ++++ tests/usage-visibility.test.ts | 341 ++++++++++++++++++ .../public/.well-known/mcp/server-card.json | 1 + 20 files changed, 1179 insertions(+), 14 deletions(-) create mode 100644 src/billing/usage-banner.ts create mode 100644 src/billing/usage-report.ts create mode 100644 src/ui/pages/usage.tsx create mode 100644 tests/usage-page.test.tsx create mode 100644 tests/usage-visibility.test.ts diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 45fa6708..4ac84c8c 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -617,7 +617,7 @@ src/ │ ├── dispatch.ts # Runs tool calls against the real routers, in-process │ ├── protocol.ts # JSON-RPC 2.0 / MCP message layer │ ├── schema.ts # Tool arg schemas: JSON Schema + validator -│ └── tools.ts # The eighteen tools +│ └── tools.ts # The nineteen tools ├── merge/ │ ├── post-merge.ts # Post-merge actions │ └── protection.ts # Merge protection rules diff --git a/src/billing/usage-banner.ts b/src/billing/usage-banner.ts new file mode 100644 index 00000000..ce38c8e1 --- /dev/null +++ b/src/billing/usage-banner.ts @@ -0,0 +1,146 @@ +/** + * The 80% banner: the same crossing the email announces, shown on the pages a + * signed-in user actually visits (PRD §8, Task 9.3). + * + * ## Why it is a stored receipt and not a computed check + * + * "Are you near a limit?" is a question about D1 (`usage_periods`) and the + * billing service, and asking it on every page render would put a database + * read and an entitlements lookup in front of every dashboard, project and + * change page — for a banner that is absent for almost everyone, almost + * always. So nothing is computed here. `usage-notifications.ts` already + * detects the crossing exactly once, on the write that causes it, from the + * totals `upsertUsage` returns; {@link recordUsageBanner} stores what it + * found, and rendering costs **one KV `get`** and no D1 at all. + * + * The cost is bounded further: {@link loadUsageBanner} returns immediately — + * with no KV read whatsoever — when there is no signed-in user or when + * `entitlementsEnabled` is false, so a self-hoster's pages are byte-for-byte + * as cheap as before. + * + * ## Users only + * + * There is no org UI in this application, so an org's threshold reaches its + * actor by email alone (PRD §8). A banner is therefore written only when the + * enforcement subject IS a user, and is keyed on that user. + * + * ## What it costs in freshness, stated rather than discovered + * + * The stored notice is a snapshot of the crossing. Usage cannot go down within + * a period, so it never becomes optimistic; it can become pessimistic if the + * account's limit is RAISED mid-period, in which case the banner overstates + * until the period rolls. That is the same staleness the email's receipt + * already accepts, and the alternative — re-reading entitlements per page — + * is the cost this design exists to avoid. + */ +import type { MeterKey } from "../storage/usage"; +import type { Env } from "../types"; +import type { Logger } from "../utils/logger"; +import { entitlementsEnabled } from "./entitlements"; + +/** What a rendered banner says. Stored verbatim; every field is display data. */ +export interface UsageBannerNotice { + meter: MeterKey; + /** Hosted consumption at the moment of the crossing. */ + used: number; + /** The limit in force then. See the freshness note in the module doc. */ + limit: number; + /** The threshold that was crossed, as a percentage (80). */ + percent: number; + /** 'YYYY-MM' UTC — the period the figures belong to. */ + period: string; +} + +/** Versioned like the entitlements cache: a shape change must not read back as the old shape. */ +const BANNER_PREFIX = "usage-banner:v1:"; + +/** + * A month plus the grace the meter itself allows. The period is IN the key, so + * a new month silently reads as "no banner" without anything sweeping the old + * one; the TTL is only what stops KV accumulating them forever. + */ +const BANNER_TTL_SECONDS = 40 * 24 * 60 * 60; + +function bannerKey(userId: string, period: string): string { + return `${BANNER_PREFIX}${userId}:${period}`; +} + +/** + * Store the crossing so the next page render can show it. + * + * Best-effort by contract, like everything on the notification path: a banner + * that cannot be written must never fail the merge whose usage produced it. + */ +export async function recordUsageBanner( + env: Env, + logger: Logger, + userId: string, + notice: UsageBannerNotice, +): Promise { + if (!env.STATE || !userId) return; + try { + await env.STATE.put(bannerKey(userId, notice.period), JSON.stringify(notice), { + expirationTtl: BANNER_TTL_SECONDS, + }); + } catch (error) { + logger.warn("Usage banner could not be stored", { + userId, + error: error instanceof Error ? error.message : String(error), + }); + } +} + +/** + * The banner to render for this user right now, or `null`. + * + * One KV read, and none at all when the seam is off or nobody is signed in. + * Every failure and every unrecognised value resolves to `null`: a page is not + * worth failing for a notice, and a half-parsed one would render nonsense. + * + * @param userId - The signed-in user; agents are not shown pages + * @param period - 'YYYY-MM' UTC; the caller passes the period it is rendering + */ +export async function loadUsageBanner( + env: Env, + logger: Logger, + userId: string | undefined, + period: string, +): Promise { + if (!userId) return null; + if (!entitlementsEnabled(env)) return null; + if (!env.STATE) return null; + try { + const raw = await env.STATE.get(bannerKey(userId, period)); + if (!raw) return null; + return parseNotice(JSON.parse(raw), period); + } catch (error) { + logger.debug("Usage banner unreadable; rendering without it", { + userId, + error: error instanceof Error ? error.message : String(error), + }); + return null; + } +} + +/** + * Validated on the way out, not trusted because we wrote it: the entry outlives + * deploys, and a notice from an older shape must vanish rather than render as + * "NaN% of undefined". + */ +function parseNotice(value: unknown, period: string): UsageBannerNotice | null { + if (typeof value !== "object" || value === null) return null; + const record = value as Record; + const { meter, used, limit, percent } = record; + if (typeof meter !== "string") return null; + if (!Number.isFinite(used) || !Number.isFinite(limit) || !Number.isFinite(percent)) return null; + if ((limit as number) <= 0) return null; + return { + meter: meter as MeterKey, + used: used as number, + limit: limit as number, + percent: percent as number, + // The key already pins the period; echoing the caller's keeps a mismatched + // stored value from being displayed against the month it was not about. + period, + }; +} diff --git a/src/billing/usage-notifications.ts b/src/billing/usage-notifications.ts index f4b8bff8..91e5798e 100644 --- a/src/billing/usage-notifications.ts +++ b/src/billing/usage-notifications.ts @@ -39,6 +39,7 @@ import type { Env } from "../types"; import type { Logger } from "../utils/logger"; import { periodResetsAt, resolveEnforcementSubject } from "./enforcement"; import { RemoteEntitlements, UnlimitedEntitlements, entitlementsEnabled } from "./entitlements"; +import { recordUsageBanner } from "./usage-banner"; /** Fraction of a limit that earns a warning. One threshold, so one receipt shape. */ export const USAGE_NOTICE_THRESHOLD = 0.8; @@ -140,6 +141,20 @@ async function deliver(env: Env, logger: Logger, input: UsageNoticeInput): Promi `${total.meter}:${THRESHOLD_PERCENT}:${limit}`; if (await alreadySent(env, logger, key)) continue; + // The in-app banner (Task 9.3) is written from the crossing this loop has + // already detected, so a page render costs one KV read and no D1. Users + // only: there is no org UI to render it in, so a pooled org's threshold + // reaches its actor by email alone (PRD §8). + if (subject.ownerType === "user") { + await recordUsageBanner(env, logger, subject.ownerId, { + meter: total.meter, + used: total.quantity, + limit, + percent: THRESHOLD_PERCENT, + period: input.period, + }); + } + await send(env, logger, { subject, actorUserId: input.actorUserId, diff --git a/src/billing/usage-report.ts b/src/billing/usage-report.ts new file mode 100644 index 00000000..14eb0953 --- /dev/null +++ b/src/billing/usage-report.ts @@ -0,0 +1,206 @@ +/** + * What one account has consumed against what its plan allows. + * + * The single read behind both visibility surfaces of PRD §8 and §4c — the + * server-rendered `/settings/usage` page and the `stratum_get_usage` MCP tool — + * so a human and an agent are never shown two different answers to the same + * question. + * + * Three properties are load bearing, and all three are easy to get wrong: + * + * - **The subject is the §4a ENFORCEMENT subject**, resolved through + * {@link resolveEnforcementSubject}, not the recorded billing subject. A + * limit is checked against the acting person, so a page that reported the + * recorded owner's totals would show a number no limit is ever compared to. + * - **Consumption is `source = 'platform'` only.** Migration 049 keys usage on + * `(owner_id, period, meter, source)` precisely so BYOK spend accumulates + * apart, and Goal 5 says it is not under the hosted allowance. Summing both + * would tell someone their own provider bill had eaten an allowance it never + * touched. BYOK is reported as its OWN figure, for information, never folded + * into `used`. + * - **`-1` is unlimited and says so.** A self-hoster's every limit is `-1` + * (`UnlimitedEntitlements`), so a renderer that turned a limit into a + * percentage would show them an empty or a divide-by-zero bar on every row + * and look broken. {@link UsageMeterReport.unlimited} is the flag a caller + * branches on, and `percentUsed` is `null` rather than a made-up number. + * + * Read-only throughout: nothing here reserves, settles, records or mutates. + */ +import { type MeterKey, getOwnerMeterTotals, usagePeriod } from "../storage/usage"; +import type { Env } from "../types"; +import { type AppError, ValidationError } from "../utils/errors"; +import type { Logger } from "../utils/logger"; +import { type Result, err, ok } from "../utils/result"; +import { periodResetsAt, resolveEnforcementSubject } from "./enforcement"; +import { + METER_KEYS, + RATE_KEYS, + type RateKey, + RemoteEntitlements, + UNLIMITED, + UnlimitedEntitlements, + entitlementsEnabled, +} from "./entitlements"; + +/** One meter's consumption this period, and the allowance it is measured against. */ +export interface UsageMeterReport { + meter: MeterKey; + /** Hosted consumption — `source = 'platform'`, the only kind a limit binds. */ + used: number; + /** + * Bring-your-own-key consumption, reported apart and never added to + * {@link UsageMeterReport.used}: it is spent on the project's own provider + * account and no allowance applies to it (Goal 5). + */ + byok: number; + /** `-1` unlimited, `0` a hard block, otherwise the ceiling. */ + limit: number; + /** True when `limit` is `-1`. Callers render words, not an empty bar. */ + unlimited: boolean; + /** True when `limit` is `0` — not "unset", a plan that forbids this outright. */ + blocked: boolean; + /** `used / limit` as a percentage, or `null` when there is no finite limit. */ + percentUsed: number | null; + /** What is left of a finite allowance, floored at zero; `null` when unlimited. */ + remaining: number | null; +} + +/** A rate ceiling. Reported as a limit only — its window lives on the meter object, not in D1. */ +export interface UsageRateReport { + rate: RateKey; + limit: number; + unlimited: boolean; +} + +export interface UsageReport { + /** Whose allowance this is — the §4a enforcement subject. */ + subject: { ownerId: string; ownerType: "user" | "org" }; + /** Opaque plan name from the billing service; `"unlimited"` when there is none. */ + plan: string; + /** + * Whether any allowance is in force at all. False for every self-hoster (no + * `BILLING_SERVICE_URL`), which is why the page must read as "unlimited" + * rather than as "we could not load your usage". + */ + metered: boolean; + /** 'YYYY-MM' UTC. */ + period: string; + /** ISO 8601 instant the monthly meters roll over. */ + resetsAt: string; + meters: UsageMeterReport[]; + rates: UsageRateReport[]; +} + +/** + * A short display name per meter, for a column heading or a banner sentence. + * + * Deliberately not `enforcement.ts`'s `meterLabel`, which builds the fragment + * "monthly LLM token allowance" for the middle of a refusal sentence. This one + * is a heading; sharing one string would make one of the two read badly. + */ +export function meterTitle(meter: MeterKey): string { + switch (meter) { + case "llm_tokens_month": + return "LLM tokens"; + case "sandbox_ms_month": + return "Sandbox time"; + case "deploys_month": + return "Deploys"; + default: + return meter; + } +} + +function meterEntry( + meter: MeterKey, + limit: number, + platform: Partial>, + byok: Partial>, +): UsageMeterReport { + const used = platform[meter] ?? 0; + // A limit the billing service never sent, or sent as garbage, has already + // been discarded by `parseEntitlements`; anything still non-integer here is + // treated as "no usable limit" for the same reason `checkGauge` does — a + // limit derived from garbage must not render as a bar someone acts on. + const usable = Number.isInteger(limit); + const unlimited = !usable || limit === UNLIMITED; + const finite = usable && limit > 0; + return { + meter, + used, + byok: byok[meter] ?? 0, + limit, + unlimited, + blocked: usable && limit === 0, + percentUsed: finite ? Math.round((used / limit) * 100) : null, + remaining: finite ? Math.max(0, limit - used) : null, + }; +} + +/** + * Build one account's usage report for a period. + * + * @param opts.actorUserId - The signed-in user, or an agent's owner (§4a) + * @param opts.at - The moment whose period to report; defaults to now + * @param opts.includeByok - Read the `byok` totals too. The page shows them as + * a separate labelled figure; the MCP tool does not ask for them, and the + * flag keeps that read off a path that would only discard it. + * @returns The report, or the error that stopped it. A caller mistake (no + * actor) is an error rather than an empty report — "you have used nothing" + * is a claim, and it must not be made about nobody. + */ +export async function buildUsageReport( + env: Env, + logger: Logger, + opts: { actorUserId: string; at?: Date; includeByok?: boolean }, +): Promise> { + if (!opts.actorUserId) { + return err(new ValidationError("Usage requested without an actor")); + } + const at = opts.at ?? new Date(); + const period = usagePeriod(at); + + // No project in hand, so this resolves to the actor itself and performs no + // I/O — which is exactly §4a's answer for "whose allowance is this page + // about": the person reading it, never an org they happen to belong to. + const subject = await resolveEnforcementSubject(env, logger, { actorUserId: opts.actorUserId }); + if (!subject) { + return err(new ValidationError("Usage subject could not be resolved")); + } + + const resolved = await new RemoteEntitlements(env, logger).forOwner( + subject.ownerId, + subject.ownerType, + ); + // Fails open exactly as every other reader does: an unreachable billing + // service shows unlimited rather than a page full of errors. + const entitlements = resolved.success ? resolved.data.entitlements : UnlimitedEntitlements; + + // The same read `reconcileFloor` performs before a limit check, with the same + // `source` filter. Anything else would be a second definition of "used". + const platform = await getOwnerMeterTotals(env.DB, logger, subject.ownerId, period, "platform"); + if (!platform.success) return err(platform.error); + + let byok: Partial> = {}; + if (opts.includeByok) { + const read = await getOwnerMeterTotals(env.DB, logger, subject.ownerId, period, "byok"); + // A failed BYOK read is not worth failing the page over: it is an + // informational figure beside the allowance, not the allowance itself. + if (read.success) byok = read.data; + } + + return ok({ + subject: { ownerId: subject.ownerId, ownerType: subject.ownerType }, + plan: entitlements.plan, + metered: entitlementsEnabled(env), + period, + resetsAt: periodResetsAt(period), + meters: METER_KEYS.map((meter) => + meterEntry(meter, entitlements.meters[meter], platform.data, byok), + ), + rates: RATE_KEYS.map((rate) => { + const limit = entitlements.rates[rate]; + return { rate, limit, unlimited: !Number.isInteger(limit) || limit === UNLIMITED }; + }), + }); +} diff --git a/src/mcp/client.ts b/src/mcp/client.ts index 305f9cd6..bfe70a65 100644 --- a/src/mcp/client.ts +++ b/src/mcp/client.ts @@ -260,4 +260,15 @@ export class StratumClient { async me() { return this.request<{ id: string; email: string }>("GET", "/api/users/me"); } + + /** + * The caller's metered usage against their plan limits. + * + * Read-only by design and by surface: there is no billing WRITE anywhere on + * this client, because an agent that can raise its own limit does not have + * one (PRD §4c). + */ + async getUsage() { + return this.request("GET", "/api/users/me/usage"); + } } diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index 6f70bc5d..4fa85f76 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -1,5 +1,5 @@ /** - * The eighteen tools the MCP server exposes. + * The nineteen tools the MCP server exposes. * * The full agent contribution loop: read a project, fork a workspace, commit, * open an evaluation-gated change, follow it through review to merge, and track @@ -12,6 +12,16 @@ * Tool DESCRIPTIONS are the only documentation a model gets, so they state the * refusals as well as the capabilities: a model that knows `stratum_review_change` * will reject its agent token asks a human instead of retrying nineteen times. + * + * The billing surface is READ-ONLY and stays that way. `stratum_get_usage` + * exists because an agent that can see 12% of its allowance left batches its + * changes, moves the project to its own provider key, or stops and asks its + * human — where an agent that cannot see it discovers the limit by hitting a + * wall. There is deliberately no upgrade, top-up or payment-method tool, and no + * `stratum_set_provider_key`: a limit an agent can raise is not a limit, and a + * credential must never pass through a model's context window (PRD §4c). Both + * refusals are stated in `stratum_get_usage`'s own description so a model reads + * them where it would otherwise go looking. */ import type { StratumClient } from "./client"; import { InvalidArgumentError, parseProjectRef } from "./client"; @@ -115,6 +125,12 @@ export function buildTools(client: StratumClient): ToolDef[] { {}, () => client.me(), ), + tool( + "stratum_get_usage", + "Read this account's metered usage against its plan limits for the current billing period: per-meter consumption (LLM tokens, sandbox time, deploys), the limit on each, the rate ceilings, the period, and when it resets. A limit of -1 means unlimited, 0 means the plan forbids that meter outright; when a limit is unlimited there is no percentage to compute. Figures cover platform-billed usage only — spend on a project's own provider key (BYOK) is billed by that provider and is never counted against these allowances. An agent token reports its OWNER's allowance, because that is the account its evaluations are charged to. Read-only, and the only billing tool there is: nothing here can raise a limit, buy capacity, add a payment method or set a provider key, for the same reason agent tokens cannot submit review verdicts — a spend limit an agent can lift is not a limit. Do not look for one; hand the account owner https://docs.usestratum.dev and let a human do it in a browser.", + {}, + () => client.getUsage(), + ), // ── Projects ────────────────────────────────────────────────────────── tool( diff --git a/src/routes/ui.tsx b/src/routes/ui.tsx index 4ee339b6..c3a982b1 100644 --- a/src/routes/ui.tsx +++ b/src/routes/ui.tsx @@ -1,6 +1,8 @@ import { Hono } from "hono"; import type { Context } from "hono"; import { fetchInviteCodes, referralServiceConfigured } from "../beta/gate"; +import { type UsageBannerNotice, loadUsageBanner } from "../billing/usage-banner"; +import { buildUsageReport } from "../billing/usage-report"; import { cannotMintLegacyCredential } from "../middleware/auth"; import { NO_ANALYTICS_HEADER } from "../middleware/web-analytics"; import { createAgent, deleteAgent, getAgent, listAgents } from "../storage/agents"; @@ -48,6 +50,7 @@ import { listWorkspaces, } from "../storage/state"; import { getProjectSourceUrl, getSyncStatus } from "../storage/sync"; +import { usagePeriod } from "../storage/usage"; import { disableLegacyToken, getUser, @@ -81,6 +84,7 @@ import { } from "../ui/pages/settings"; import { SyncPage } from "../ui/pages/sync"; import { TagsPage } from "../ui/pages/tags"; +import { UsagePage } from "../ui/pages/usage"; import { WebhooksPage } from "../ui/pages/webhooks"; import { WorkspacesPage } from "../ui/pages/workspaces"; import { canReadProject, canWriteProject, filterMemberProjects } from "../utils/authz"; @@ -145,6 +149,23 @@ async function getCurrentUser( }; } +/** + * The 80% usage warning for the signed-in browser user, or `null`. + * + * Threaded into the pages a user actually visits — dashboard, settings, the + * changes list and a change's detail — rather than only onto `/settings/usage`, + * because a page nobody opens warns nobody (PRD §8). It costs **one KV read** + * and no D1: the crossing was already detected and stored by + * `usage-notifications.ts` at the moment the usage was written. Signed-out + * visitors and self-hosted instances take no read at all. + */ +async function usageBannerFor( + c: Context<{ Bindings: Env }>, + logger: ReturnType, +): Promise { + return loadUsageBanner(c.env, logger, c.get("userId"), usagePeriod()); +} + /** * Resolve issue author ids to display names ("@username" for users, the agent's * name for agents). Unresolvable authors fall back to their author type. @@ -215,7 +236,9 @@ app.get("/", async (c) => { })); logger.debug("Rendering home page", { projectCount: view.length }); - return c.html(); + return c.html( + , + ); }); // GET /new — New project form @@ -409,7 +432,7 @@ async function renderSettings( logger: ReturnType, extras: { freshToken?: FreshCredential; notice?: SettingsNotice } = {}, ) { - const [agents, tokens, grants, usernameChange, invites] = await Promise.all([ + const [agents, tokens, grants, usernameChange, invites, usageNotice] = await Promise.all([ loadAgentSummaries(c.env.DB, user.id, logger), loadApiTokens(c.env.DB, user.id, logger), loadOAuthGrants(c.env.DB, user.id, logger), @@ -420,6 +443,7 @@ async function renderSettings( referralServiceConfigured(c.env) ? fetchInviteCodes(c.env, user.id, logger) : Promise.resolve(undefined), + usageBannerFor(c, logger), ]); // A notice about the action just taken outranks a listing's own failure — // the caller needs to know what their POST did first. @@ -439,6 +463,7 @@ async function renderSettings( // is what keeps them from rendering identically — see `loadOAuthGrants`. oauthGrantsUnavailable={grants.notice !== undefined} telemetryOptOut={telemetryOptOut} + usageNotice={usageNotice} // Threaded here rather than at each call site: this helper is the single // render path for the settings page, so the CSP nonce main added cannot // be forgotten by a future caller. @@ -497,6 +522,39 @@ app.get("/settings", async (c) => { return c.html(page); }); +/** + * GET /settings/usage — this account's consumption against its allowances. + * + * A page of its own, not a section of `/settings`: that route renders one + * hash-anchored document (which `/profile` redirects into), and usage is a + * different question with a different read behind it. + * + * Users only, and a browser session only, like every other account page. There + * is deliberately no organization view — this application has no org UI, and + * under PRD §4a an allowance is checked against the person anyway, so what is + * shown here is exactly what every limit is compared to. + */ +app.get("/settings/usage", async (c) => { + const logger = createLogger({ path: c.req.path, userId: c.get("userId") }); + const access = await requireAccountSession(c, logger); + if ("response" in access) return access.response; + + const [report, usageNotice] = await Promise.all([ + buildUsageReport(c.env, logger, { actorUserId: access.user.id, includeByok: true }), + usageBannerFor(c, logger), + ]); + if (!report.success) { + // The allowances are still worth showing when only the totals failed to + // read, but this build cannot tell those apart, so the page says plainly + // that the figures are incomplete rather than rendering zeroes as fact. + logger.error("Failed to build usage report", report.error); + const message = "Your usage could not be loaded. Please try again."; + return c.html(errorPage(500, message, access.user), 500); + } + + return c.html(); +}); + // POST /settings/account — Set or clear the display name app.post("/settings/account", async (c) => { const logger = createLogger({ path: c.req.path, userId: c.get("userId") }); @@ -1218,6 +1276,7 @@ app.get("/p/:name/changes", async (c) => { changes={view} canWrite={canWrite} user={userResult} + usageNotice={await usageBannerFor(c, logger)} />, ); }); @@ -1344,6 +1403,7 @@ app.get("/changes/:id", async (c) => { canReview={canReview} projectRef={projectRef(projectResult.data)} user={userResult} + usageNotice={await usageBannerFor(c, logger)} />, ); }); @@ -1449,6 +1509,7 @@ app.get("/:namespace/:slug/changes", async (c) => { changes={changes} canWrite={canWrite} user={userResult} + usageNotice={await usageBannerFor(c, logger)} />, ); }); diff --git a/src/routes/users.ts b/src/routes/users.ts index d5ab3e8e..a903f0c7 100644 --- a/src/routes/users.ts +++ b/src/routes/users.ts @@ -1,5 +1,6 @@ import { Hono } from "hono"; import type { Context } from "hono"; +import { buildUsageReport } from "../billing/usage-report"; import { cannotMintLegacyCredential } from "../middleware/auth"; import { MAX_TOKEN_EXPIRY_DAYS, @@ -62,6 +63,41 @@ app.get("/me", async (c) => { return ok({ id: user.id, email: user.email, createdAt: user.createdAt }); }); +/** + * GET /api/users/me/usage — the caller's metered usage against their plan. + * + * Read-only, and deliberately the ONLY billing surface reachable from a token: + * there is no endpoint here to raise a limit, buy capacity or attach a payment + * method (PRD §4c). An agent that can lift its own ceiling does not have one, + * which is the same line this codebase already draws at review verdicts. + * + * An AGENT token resolves to its owner (`agentOwnerId`), because that is the + * §4a enforcement subject its evaluations are actually charged to — an agent + * shown its own empty allowance would be reading a number no check uses. + */ +app.get("/me/usage", async (c) => { + const logger = createLogger({ + requestId: crypto.randomUUID(), + path: c.req.path, + method: c.req.method, + userId: c.get("userId"), + }); + + // The acting user, or the person an agent acts for. Never a project's org: + // limits follow the person (PRD §4a). + const actorUserId = c.get("userId") ?? c.get("agentOwnerId"); + if (!actorUserId) { + return c.json({ error: "Unauthorized" }, 401); + } + + const report = await buildUsageReport(c.env, logger, { actorUserId }); + if (!report.success) { + logger.error("Failed to build usage report", report.error); + return internalError("Usage could not be read"); + } + return ok(report.data); +}); + // POST /api/users/me/rotate-token - Replace the caller's API key app.post("/me/rotate-token", async (c) => { const logger = createLogger({ diff --git a/src/ui/layout.tsx b/src/ui/layout.tsx index feb261b9..716cc471 100644 --- a/src/ui/layout.tsx +++ b/src/ui/layout.tsx @@ -1,4 +1,6 @@ import type { FC } from "hono/jsx"; +import type { UsageBannerNotice } from "../billing/usage-banner"; +import { meterTitle } from "../billing/usage-report"; import { SourceFooter } from "./components/source-footer"; /** @@ -21,11 +23,50 @@ interface LayoutProps { * see the URL, and a page under /new/import is still "new" to the reader. */ active?: NavItem; + /** + * The 80% usage warning, when this account has crossed one (PRD §8). + * + * A prop rather than something the layout fetches: the layout is a pure + * component with no request context, and the route that renders a page is + * the only place that can decide whether one KV read is worth taking. Pages + * that do not thread it simply do not warn. + */ + usageNotice?: UsageBannerNotice | null | undefined; children?: unknown; } +/** + * "You are near a limit", shown where a user is already looking. + * + * A page nobody visits warns nobody, which is why this lives in the shared + * chrome and not on `/settings/usage` alone. It has no dismiss control on + * purpose: dismissal is state, state without script means a cookie or a POST, + * and the notice is true until the period rolls over anyway. + */ +const UsageBanner: FC<{ notice: UsageBannerNotice }> = ({ notice }) => ( + // `` rather than a div with role="status": it is the semantic element + // for exactly this, which is what Biome's useSemanticElements asks for, and it + // carries the same implicit ARIA role to a screen reader. + + + You have used {notice.percent}% of this account's monthly{" "} + {meterTitle(notice.meter).toLowerCase()} allowance for {notice.period} ( + {notice.used.toLocaleString("en-US")} of {notice.limit.toLocaleString("en-US")}). Work that + needs it is refused once it runs out. + + View usage + +); + /** The page chrome shared by every server-rendered page: header nav, main column, footer, and the CSP-nonced scripts. */ -export const Layout: FC = ({ title, user, refreshSeconds, active, children }) => { +export const Layout: FC = ({ + title, + user, + refreshSeconds, + active, + usageNotice, + children, +}) => { const current = (item: NavItem) => (active === item ? "page" : undefined); return ( @@ -89,7 +130,10 @@ export const Layout: FC = ({ title, user, refreshSeconds, active, c )} -
{children}
+
+ {usageNotice && } + {children} +
diff --git a/src/ui/pages/change-detail.tsx b/src/ui/pages/change-detail.tsx index 886d4df3..b5466151 100644 --- a/src/ui/pages/change-detail.tsx +++ b/src/ui/pages/change-detail.tsx @@ -1,4 +1,5 @@ import type { FC } from "hono/jsx"; +import type { UsageBannerNotice } from "../../billing/usage-banner"; import type { ChangeComment, ChangeReview } from "../../storage/change-reviews"; import type { CostSummaryEntry } from "../../storage/costs"; import { type DiffFile, DiffView, LineCommentThreads } from "../components/diff-view"; @@ -44,6 +45,8 @@ interface ChangeDetailProps { /** Namespaced project ref for the header; legacy changes may lack one. */ projectRef?: { name: string; namespace: string; slug: string; visibility?: string } | null; user?: { id: string; email: string; username: string } | null; + /** The 80% warning, when this account has one. See `Layout`. */ + usageNotice?: UsageBannerNotice | null; } function statusBadgeClass(status: string): string { @@ -101,6 +104,7 @@ export const ChangeDetailPage: FC = ({ canReview = false, projectRef = null, user, + usageNotice, }) => { // Line-anchored comments (and their replies) render as threads beneath the // diff; anchorless comments stay in the change-level discussion below. @@ -116,6 +120,7 @@ export const ChangeDetailPage: FC = ({ {projectRef && } diff --git a/src/ui/pages/changes.tsx b/src/ui/pages/changes.tsx index 20f1ee1d..7cf39db1 100644 --- a/src/ui/pages/changes.tsx +++ b/src/ui/pages/changes.tsx @@ -1,4 +1,5 @@ import type { FC } from "hono/jsx"; +import type { UsageBannerNotice } from "../../billing/usage-banner"; import { ProjectHeader } from "../components/project-header"; import { Layout } from "../layout"; @@ -19,6 +20,8 @@ interface ChangesProps { createdAt: string; }>; user?: { id: string; email: string; username: string } | null; + /** The 80% warning, when this account has one. See `Layout`. */ + usageNotice?: UsageBannerNotice | null; } function statusBadgeClass(status: string): string { @@ -40,9 +43,15 @@ function statusBadgeClass(status: string): string { } } -export const ChangesPage: FC = ({ project, changes, canWrite, user }) => { +export const ChangesPage: FC = ({ + project, + changes, + canWrite, + user, + usageNotice, +}) => { return ( - + {changes.length === 0 ? ( diff --git a/src/ui/pages/home.tsx b/src/ui/pages/home.tsx index b7f72737..97c6964f 100644 --- a/src/ui/pages/home.tsx +++ b/src/ui/pages/home.tsx @@ -1,4 +1,5 @@ import type { FC } from "hono/jsx"; +import type { UsageBannerNotice } from "../../billing/usage-banner"; import { Layout } from "../layout"; interface HomeProps { @@ -11,11 +12,13 @@ interface HomeProps { visibility?: string; }>; user?: { id: string; email: string; username: string } | null; + /** The 80% warning, when this account has one. See `Layout`. */ + usageNotice?: UsageBannerNotice | null; } -export const HomePage: FC = ({ projects, user }) => { +export const HomePage: FC = ({ projects, user, usageNotice }) => { return ( - +

Dashboard

{user && ( diff --git a/src/ui/pages/settings.tsx b/src/ui/pages/settings.tsx index 9867ac4e..db6f5baa 100644 --- a/src/ui/pages/settings.tsx +++ b/src/ui/pages/settings.tsx @@ -1,5 +1,6 @@ import type { FC } from "hono/jsx"; import type { InviteCodesResult } from "../../beta/gate"; +import type { UsageBannerNotice } from "../../billing/usage-banner"; import { type ApiTokenSummary, MAX_ACTIVE_TOKENS_PER_USER, @@ -80,6 +81,8 @@ interface SettingsPageProps { nonce?: string; /** True when the user has opted out of product analytics (#257). */ telemetryOptOut: boolean; + /** The 80% warning, when this account has one. See `Layout`. */ + usageNotice?: UsageBannerNotice | null; } const SCOPE_LABEL: Record = { @@ -305,6 +308,7 @@ export const SettingsPage: FC = ({ notice, nonce, telemetryOptOut, + usageNotice, }) => { const now = Date.now(); // Mirrors the server-side cap in `createApiToken`: an expired row occupies no @@ -314,9 +318,14 @@ export const SettingsPage: FC = ({ (token) => token.revokedAt === undefined && !isExpired(token.expiresAt ?? null), ).length; return ( - +

Settings

+
+ + Usage + +