diff --git a/.gitignore b/.gitignore index d6fde4b..1635499 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,8 @@ package-lock.json # setup.sh backs up wrangler.toml on first successful run wrangler.toml.bak + +# Local agent / review artifacts (session-scoped worktrees + per-run +# reviewer JSON dumps). +.claude/ +.context/ diff --git a/packages/worker/src/auth/email.ts b/packages/worker/src/auth/email.ts index 1ea2538..d0b124d 100644 --- a/packages/worker/src/auth/email.ts +++ b/packages/worker/src/auth/email.ts @@ -19,6 +19,10 @@ export interface EmailConfig { const RESEND_URL = "https://api.resend.com/emails" const MAX_ATTEMPTS = 3 const BASE_BACKOFF_MS = 200 +/** Upper bound per fetch attempt. Resend's p99 is ~1-2s; 5s absorbs slow + * starts while keeping a hung connection from burning the whole + * `ctx.waitUntil` budget (Cloudflare caps subrequest time). */ +const FETCH_TIMEOUT_MS = 5000 export async function sendMagicLink( config: EmailConfig, @@ -34,35 +38,7 @@ export async function sendMagicLink( `${link}\n\n` + `If you didn't request this, you can safely ignore this email.`, }) - - let lastError = "unknown" - for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) { - try { - const res = await fetch(RESEND_URL, { - method: "POST", - headers: { - Authorization: `Bearer ${config.apiKey}`, - "Content-Type": "application/json", - }, - body, - }) - if (res.ok) return { ok: true } - - // 4xx is not retryable (bad API key, bad from address, etc.). - if (res.status >= 400 && res.status < 500) { - return { ok: false, error: `resend-${res.status}` } - } - lastError = `resend-${res.status}` - } catch (err) { - lastError = err instanceof Error ? `network-${err.name}` : "network-unknown" - } - - // Exponential backoff: 200ms, 400ms. - if (attempt < MAX_ATTEMPTS - 1) { - await sleep(BASE_BACKOFF_MS * 2 ** attempt) - } - } - return { ok: false, error: lastError } + return postToResend(config, body) } /** @@ -109,9 +85,26 @@ export async function sendOperatorAlert( subject, text, }) + return postToResend(config, body) +} +/** + * Shared Resend POST with 5s per-attempt timeout, 3-attempt exponential + * backoff, and 0-30% jitter. Jitter prevents concurrent isolates from + * retrying in lockstep during a Resend brownout (the deterministic + * 200/400ms cadence would otherwise compound outbound load). + * + * 4xx → terminal (bad API key, unverified sender, malformed body). + * 5xx / network / abort → retryable. + */ +async function postToResend( + config: EmailConfig, + body: string, +): Promise { let lastError = "unknown" for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) { + const controller = new AbortController() + const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS) try { const res = await fetch(RESEND_URL, { method: "POST", @@ -120,6 +113,7 @@ export async function sendOperatorAlert( "Content-Type": "application/json", }, body, + signal: controller.signal, }) if (res.ok) return { ok: true } if (res.status >= 400 && res.status < 500) { @@ -128,9 +122,13 @@ export async function sendOperatorAlert( lastError = `resend-${res.status}` } catch (err) { lastError = err instanceof Error ? `network-${err.name}` : "network-unknown" + } finally { + clearTimeout(timer) } if (attempt < MAX_ATTEMPTS - 1) { - await sleep(BASE_BACKOFF_MS * 2 ** attempt) + const base = BASE_BACKOFF_MS * 2 ** attempt + const jitter = Math.floor(Math.random() * base * 0.3) + await sleep(base + jitter) } } return { ok: false, error: lastError } diff --git a/packages/worker/src/ingest-do/session-ingest-do.ts b/packages/worker/src/ingest-do/session-ingest-do.ts index 7d1a884..bb31509 100644 --- a/packages/worker/src/ingest-do/session-ingest-do.ts +++ b/packages/worker/src/ingest-do/session-ingest-do.ts @@ -2,8 +2,8 @@ import { eq, sql } from "drizzle-orm" import { createDb, type Db } from "../db/client" import { sessionBatches, telemetrySessions } from "../db/schema" import type { Env } from "../env" -import { QuotaExceededError } from "../quota/daily-quota" -import { putBatchJsonl } from "../storage/r2" +import { chargeOrThrow, QuotaExceededError } from "../quota/daily-quota" +import { encodeBatchJsonl, putBatchJsonlBytes } from "../storage/r2" /** * Per-session ingest coordinator (one DO instance per telemetry_sessions.id). @@ -38,6 +38,13 @@ interface IngestCompleteRequest { export class SessionIngestDO implements DurableObject { private seq: number | null = null + /** + * Highest seq we've ever debited the daily quota for. Persisted + * before the R2 PUT so that a retry after a D1 batch failure (or a DO + * eviction between charge and success) doesn't double-charge the + * counter for the same logical batch. See review finding F4. + */ + private chargedSeq = 0 private sessionDbId: string | null = null constructor( @@ -46,8 +53,10 @@ export class SessionIngestDO implements DurableObject { ) { state.blockConcurrencyWhile(async () => { const storedSeq = await state.storage.get("seq") + const storedChargedSeq = await state.storage.get("chargedSeq") const storedId = await state.storage.get("sessionDbId") if (typeof storedSeq === "number") this.seq = storedSeq + if (typeof storedChargedSeq === "number") this.chargedSeq = storedChargedSeq if (typeof storedId === "string") this.sessionDbId = storedId }) } @@ -73,12 +82,10 @@ export class SessionIngestDO implements DurableObject { return Response.json({ count: 0 } satisfies { count: number }) } - // Critical section: seq allocation + R2 write + D1 batch + seq advance. - // Wrapped in blockConcurrencyWhile so two concurrent /data calls to the - // same session (same DO) can't both compute the same seq and race on the - // session_batches UNIQUE constraint. Cloudflare's default input gate - // preserves state across I/O yields but does not serialize full - // critical sections — blockConcurrencyWhile does. + // Critical section: seq allocation + quota charge + R2 write + D1 batch + // + seq advance. Wrapped in blockConcurrencyWhile so two concurrent + // /data calls to the same session (same DO) can't both compute the same + // seq and race on the session_batches UNIQUE constraint. return this.state.blockConcurrencyWhile(async () => { await this.bindSessionId(sessionDbId) @@ -86,32 +93,42 @@ export class SessionIngestDO implements DurableObject { await this.hydrateSeq(db, sessionDbId) const nextSeq = (this.seq ?? 0) + 1 - // 1. R2 write first. On failure, no seq advance, no D1 mutation. - let r2Result: { key: string; rawByteLength: number; storedByteLength: number } - try { - r2Result = await putBatchJsonl(this.env, sessionDbId, nextSeq, entries) - } catch (err) { - if (err instanceof QuotaExceededError) { - // Surface the 429 through the DO→Worker fetch boundary with - // a status code the Worker route can pass through as-is. The - // firstBreach header tells the Worker to fire the operator - // alert + audit log via its own ctx.waitUntil (keeps Resend - // off the critical path of the 429 response). - const headers: Record = { - "Retry-After": String(err.retryAfter), - } - if (err.firstBreach) { - headers["X-Quota-First-Breach"] = err.metric - } - return new Response(`quota_cap_hit: ${err.metric}`, { - status: 429, - headers, + // 1. Encode the batch locally (no R2 yet) so we know the exact byte + // count to charge. + const encoded = await encodeBatchJsonl(sessionDbId, nextSeq, entries) + + // 2. Charge quota, but only if we haven't already charged for this + // seq on a prior attempt. `chargedSeq` is persisted BEFORE the R2 + // PUT so a D1 failure (or DO eviction) that forces a retry sees + // chargedSeq == nextSeq and skips the second charge. Without this + // dedup, a D1-flappy window would double-charge the daily counter + // per retry (review finding F4). + if (nextSeq > this.chargedSeq) { + try { + await chargeOrThrow(this.env, { + bytes: encoded.storedByteLength, + classA: 1, }) + } catch (err) { + if (err instanceof QuotaExceededError) { + return cappedDoResponse(err) + } + throw err } + this.chargedSeq = nextSeq + await this.state.storage.put("chargedSeq", nextSeq) + } + + // 3. R2 PUT. Uncharged at this layer — the charge already committed. + // On failure, chargedSeq stays at nextSeq so the retry skips the + // charge and deterministically overwrites any orphan object. + try { + await putBatchJsonlBytes(this.env, encoded.key, encoded.compressed) + } catch (err) { return new Response(`r2_write_failed: ${String(err)}`, { status: 503 }) } - // 2. D1 atomic group: insert batch row + bump uploaded_count + touch + // 4. D1 atomic group: insert batch row + bump uploaded_count + touch // last_batch_at + invalidate wpilog cache. db.batch ensures all-or- // nothing at the D1 layer. try { @@ -122,9 +139,9 @@ export class SessionIngestDO implements DurableObject { // byte_length stays as the raw JSONL length (users reading this // column see "N bytes of real telemetry"). storedByteLength is // used separately by the quota system (plan 2026-04-23-001). - byteLength: r2Result.rawByteLength, + byteLength: encoded.rawByteLength, entryCount: entries.length, - r2Key: r2Result.key, + r2Key: encoded.key, }), db .update(telemetrySessions) @@ -138,11 +155,12 @@ export class SessionIngestDO implements DurableObject { ]) } catch (err) { // D1 failed after R2 succeeded. Leave seq alone so the retry reuses - // the same number and overwrites the orphan R2 object. + // the same number and overwrites the orphan R2 object. chargedSeq + // is already persisted at nextSeq → retry won't double-charge. return new Response(`d1_batch_failed: ${String(err)}`, { status: 503 }) } - // 3. Only now advance seq in memory + storage. + // 5. Only now advance seq in memory + storage. this.seq = nextSeq await this.state.storage.put("seq", nextSeq) @@ -207,6 +225,29 @@ export class SessionIngestDO implements DurableObject { } } +/** + * DO-side 429 builder. Emits: + * - Retry-After (seconds-until-UTC-midnight, computed at breach time) + * - X-Quota-First-Breach (only on the 0→1 latch flip) so the Worker + * fires the operator alert via ctx.waitUntil + * - X-Quota-Breach-Date (always, when firstBreach) so the Worker's + * waitUntil reads the right D1 row even if it runs after UTC + * midnight (review finding F9). + */ +function cappedDoResponse(err: QuotaExceededError): Response { + const headers: Record = { + "Retry-After": String(err.retryAfter), + } + if (err.firstBreach) { + headers["X-Quota-First-Breach"] = err.metric + headers["X-Quota-Breach-Date"] = err.row.date + } + return new Response(`quota_cap_hit: ${err.metric}`, { + status: 429, + headers, + }) +} + function serialize(row: typeof telemetrySessions.$inferSelect): Record { return { id: row.id, diff --git a/packages/worker/src/ingest/tree-builder.ts b/packages/worker/src/ingest/tree-builder.ts index 9fd94c1..7d8a89e 100644 --- a/packages/worker/src/ingest/tree-builder.ts +++ b/packages/worker/src/ingest/tree-builder.ts @@ -1,7 +1,7 @@ import type { KeyTreeNode, KeyTreeResponse, TelemetryEntryRequest } from "../dto" import type { Env } from "../env" import { batchPrefix, treeKey } from "../storage/keys" -import { putTextBlob, readPlainBlobStream, readTextBlob } from "../storage/r2" +import { getBlob, listBlobs, putTextBlob, readPlainBlobStream, readTextBlob } from "../storage/r2" interface Aggregate { ntType: string @@ -19,7 +19,7 @@ interface Aggregate { */ export async function buildTree(env: Env, sessionDbId: string): Promise { const prefix = batchPrefix(sessionDbId) + "batch-" - const listed = await env.BLOBS.list({ prefix }) + const listed = await listBlobs(env, { prefix }) const keys = listed.objects .map((o) => o.key) .filter((k) => k.endsWith(".jsonl")) @@ -29,7 +29,7 @@ export async function buildTree(env: Env, sessionDbId: string): Promise { expect(result.retryAfter).toBeLessThanOrEqual(60) expect(result.retryAfter).toBeGreaterThan(0) }) + + it("concurrent racers: only one sees firstBreach=true (F5 regression)", async () => { + // Pre-seed at cap so the next charge will cross. Fire N concurrent + // charges through chargeQuota; the UPSERT+SELECT batch must ensure + // exactly one observes the 0→1 latch flip, even though all of them + // see 'over cap' in the returned row. + await seedRow("2026-04-23", { bytesUploaded: CAP_BYTES }) + const N = 10 + const results = await Promise.all( + Array.from({ length: N }, () => chargeQuota(env, { bytes: 1 }, now)), + ) + const firstBreaches = results.filter((r) => !r.ok && r.firstBreach) + const overCap = results.filter((r) => !r.ok) + expect(overCap.length).toBe(N) // all saw the breach + expect(firstBreaches.length).toBe(1) // but only one is first + const db = createDb(env) + const [row] = await db.select().from(dailyQuota) + expect(row!.alertedBytes).toBe(1) + expect(row!.bytesUploaded).toBe(CAP_BYTES + N) + }) }) diff --git a/packages/worker/src/quota/daily-quota.ts b/packages/worker/src/quota/daily-quota.ts index 792b924..03c9acb 100644 --- a/packages/worker/src/quota/daily-quota.ts +++ b/packages/worker/src/quota/daily-quota.ts @@ -26,6 +26,15 @@ export const CAP_BYTES = 1024 * 1024 * 1024 // 1 GiB compressed bytes / UTC day export const CAP_CLASS_A = 5_000 export const CAP_CLASS_B = 50_000 +// Guard against a misconfigured deploy (e.g. CAP_* accidentally set to 0) +// silently locking every write path behind a 429. A positive-value cap is +// a precondition of the module, checked once at load time. +if (CAP_BYTES <= 0 || CAP_CLASS_A <= 0 || CAP_CLASS_B <= 0) { + throw new Error( + `daily-quota: all CAP_* must be > 0 (bytes=${CAP_BYTES}, classA=${CAP_CLASS_A}, classB=${CAP_CLASS_B})`, + ) +} + export type QuotaMetric = "bytes" | "classA" | "classB" export interface QuotaCharge { @@ -55,11 +64,19 @@ export type ChargeResult = } /** - * Atomically UPSERT the counters for today, re-read the row, and check - * against the caps. If any cap is now exceeded AND the matching - * `alerted_*` flag was 0, flip it to 1 in-place and mark firstBreach. + * Atomically UPSERT today's counters, conditionally flip the + * `alerted_*` latches in-place, and return whether this call was the + * first breach per metric. + * + * Concurrency model: the whole operation is a single D1 `batch` so the + * pre-UPDATE snapshot and the UPSERT run in one transaction. D1 + * serializes writes at the database level, so two concurrent racers + * produce exactly one `firstBreach: true` result — the second racer's + * pre-snapshot already sees `alerted_* = 1`. This replaces an earlier + * three-statement pattern (UPSERT, SELECT, UPDATE-where-0) that was + * racy at the alert layer (review finding F5). * - * A zero-valued charge (all three metrics omitted or 0) is a no-op. + * A zero-valued charge short-circuits without hitting the database. */ export async function chargeQuota( env: Env, @@ -72,23 +89,43 @@ export async function chargeQuota( const today = toUtcDateString(now) if (bytes === 0 && classA === 0 && classB === 0) { - // No-op short-circuit; still need the current row for callers who - // care about `ok`. - const row = await readOrZero(env, today) - return capCheck(row, 0, 0, 0, now) + // No-op short-circuit; read the current row for callers who care + // about `ok`. + const row = await readRow(env, today) + return capCheck(row, now) } const db = createDb(env) - // Atomic UPSERT + increment. D1 supports SQLite's ON CONFLICT DO - // UPDATE. `excluded` refers to the row that would have been inserted. - await db + // The UPSERT's ON CONFLICT branch bumps the counters AND flips each + // alerted_* latch from 0 → 1 exactly when THIS charge pushes the + // corresponding counter past the cap for the first time. Folding + // the conditional flip into the same statement as the increment + // removes the earlier UPDATE-where-0 round-trip (review finding F6) + // and closes the 0→1 double-email race (review finding F5) because + // the pre-snapshot SELECT and the UPSERT run inside one db.batch. + const selectPrev = db + .select({ + alertedBytes: dailyQuota.alertedBytes, + alertedClassA: dailyQuota.alertedClassA, + alertedClassB: dailyQuota.alertedClassB, + }) + .from(dailyQuota) + .where(sql`${dailyQuota.date} = ${today}`) + + const upsertReturning = db .insert(dailyQuota) .values({ date: today, bytesUploaded: bytes, classAOps: classA, classBOps: classB, + // On a fresh row, seed alerted_* correctly if this single charge + // already crosses a cap (edge case: very first charge of the day + // exceeds CAP). + alertedBytes: bytes > CAP_BYTES ? 1 : 0, + alertedClassA: classA > CAP_CLASS_A ? 1 : 0, + alertedClassB: classB > CAP_CLASS_B ? 1 : 0, updatedAt: now, }) .onConflictDoUpdate({ @@ -97,43 +134,61 @@ export async function chargeQuota( bytesUploaded: sql`${dailyQuota.bytesUploaded} + ${bytes}`, classAOps: sql`${dailyQuota.classAOps} + ${classA}`, classBOps: sql`${dailyQuota.classBOps} + ${classB}`, + alertedBytes: sql`CASE WHEN ${dailyQuota.bytesUploaded} + ${bytes} > ${CAP_BYTES} AND ${dailyQuota.alertedBytes} = 0 THEN 1 ELSE ${dailyQuota.alertedBytes} END`, + alertedClassA: sql`CASE WHEN ${dailyQuota.classAOps} + ${classA} > ${CAP_CLASS_A} AND ${dailyQuota.alertedClassA} = 0 THEN 1 ELSE ${dailyQuota.alertedClassA} END`, + alertedClassB: sql`CASE WHEN ${dailyQuota.classBOps} + ${classB} > ${CAP_CLASS_B} AND ${dailyQuota.alertedClassB} = 0 THEN 1 ELSE ${dailyQuota.alertedClassB} END`, updatedAt: now, }, }) + .returning() - const row = await readOrZero(env, today) - const result = capCheck(row, bytes, classA, classB, now) - if (!result.ok) { - // Flip the matching alerted_* flag 0→1 if this is the first breach. - if (result.firstBreach) { - const field = alertedField(result.hitCap) - await db - .update(dailyQuota) - .set({ [field]: 1, updatedAt: now }) - .where(sql`${dailyQuota.date} = ${today} AND ${dailyQuota[field]} = 0`) - } + // db.batch runs both statements in one D1 transaction. selectPrev + // observes the row AS IT WAS before the UPSERT in the same batch. + const [prevRows, postRows] = await db.batch([selectPrev, upsertReturning]) + const prev = prevRows?.[0] as + | { alertedBytes: number; alertedClassA: number; alertedClassB: number } + | undefined + const post = postRows?.[0] as typeof dailyQuota.$inferSelect | undefined + if (!post) { + throw new Error("daily-quota: UPSERT returned no row") + } + const row: DailyQuotaRow = { + date: post.date, + bytesUploaded: post.bytesUploaded, + classAOps: post.classAOps, + classBOps: post.classBOps, + alertedBytes: post.alertedBytes !== 0, + alertedClassA: post.alertedClassA !== 0, + alertedClassB: post.alertedClassB !== 0, + } + const firstBreaches = { + bytes: post.alertedBytes !== 0 && (prev?.alertedBytes ?? 0) === 0, + classA: post.alertedClassA !== 0 && (prev?.alertedClassA ?? 0) === 0, + classB: post.alertedClassB !== 0 && (prev?.alertedClassB ?? 0) === 0, } - return result + return capCheck(row, now, firstBreaches) } function capCheck( row: DailyQuotaRow, - _bytesCharged: number, - _classACharged: number, - _classBCharged: number, now: Date, + firstBreaches: Record = { + bytes: false, + classA: false, + classB: false, + }, ): ChargeResult { // Priority order: bytes > classA > classB. Bytes is the most // expensive metric to exceed and the most useful signal to an // operator, so it wins when multiple caps trip simultaneously. if (row.bytesUploaded > CAP_BYTES) { - return cappedResult(row, "bytes", !row.alertedBytes, now) + return cappedResult(row, "bytes", firstBreaches.bytes, now) } if (row.classAOps > CAP_CLASS_A) { - return cappedResult(row, "classA", !row.alertedClassA, now) + return cappedResult(row, "classA", firstBreaches.classA, now) } if (row.classBOps > CAP_CLASS_B) { - return cappedResult(row, "classB", !row.alertedClassB, now) + return cappedResult(row, "classB", firstBreaches.classB, now) } return { ok: true, row } } @@ -149,26 +204,11 @@ function cappedResult( hitCap, retryAfter: secondsUntilUtcMidnight(now), firstBreach, - // Reflect the flip in the returned row so the caller doesn't need - // to re-read D1 just to know about firstBreach. - row: firstBreach - ? { ...row, [alertedField(hitCap)]: true } - : row, - } -} - -function alertedField(metric: QuotaMetric): "alertedBytes" | "alertedClassA" | "alertedClassB" { - switch (metric) { - case "bytes": - return "alertedBytes" - case "classA": - return "alertedClassA" - case "classB": - return "alertedClassB" + row, } } -async function readOrZero(env: Env, today: string): Promise { +async function readRow(env: Env, today: string): Promise { const db: Db = createDb(env) const [row] = await db .select() diff --git a/packages/worker/src/quota/http.ts b/packages/worker/src/quota/http.ts index ae4684e..f671fc9 100644 --- a/packages/worker/src/quota/http.ts +++ b/packages/worker/src/quota/http.ts @@ -9,10 +9,10 @@ import { caps, QuotaExceededError, toUtcDateString } from "./daily-quota" /** * 429 Too Many Requests with Retry-After + a plain-text body naming the - * cap that was hit. Used by every route that catches QuotaExceededError - * from the storage wrappers. + * cap that was hit. Internal to this module; routes go through + * `handleQuotaExceeded`. */ -export function cappedResponse(err: QuotaExceededError): Response { +function cappedResponse(err: QuotaExceededError): Response { return new Response(`quota_cap_hit: ${err.metric}`, { status: 429, headers: { @@ -47,7 +47,7 @@ export function handleQuotaExceeded( * sets when its chargeOrThrow fires for the first time today. Returns * the metric, or null if absent / unrecognised. */ -export function firstBreachFromHeader( +function firstBreachFromHeader( res: Response, ): QuotaExceededError["metric"] | null { const raw = res.headers.get("X-Quota-First-Breach") @@ -58,9 +58,10 @@ export function firstBreachFromHeader( /** * Called by Worker routes when the DO-returned 429 carries the * X-Quota-First-Breach header. Reconstructs the metric + retryAfter - * from the response and reads the current counter row from D1 so the - * alert email + audit log include the real counter value (the DO→ - * Worker fetch boundary doesn't carry the row contents through). + * from the response headers and reads the counter row from D1 keyed on + * the DO-emitted `X-Quota-Breach-Date` (not `new Date()` — the worker's + * waitUntil can fire after UTC midnight while the breach row stays + * keyed to the prior day). */ export function scheduleDoAlert( c: Context<{ Bindings: Env }>, @@ -70,17 +71,17 @@ export function scheduleDoAlert( const metric = firstBreachFromHeader(res) if (!metric) return const retryAfter = Number(res.headers.get("Retry-After")) || 60 + const breachDate = res.headers.get("X-Quota-Breach-Date") || toUtcDateString(new Date()) c.executionCtx.waitUntil( - (async () => { + guardWaitUntil("scheduleDoAlert", async () => { const db = createDb(c.env) - const today = toUtcDateString(new Date()) const [row] = await db .select() .from(dailyQuota) - .where(sql`${dailyQuota.date} = ${today}`) + .where(sql`${dailyQuota.date} = ${breachDate}`) .limit(1) const err = new QuotaExceededError(metric, retryAfter, true, { - date: today, + date: breachDate, bytesUploaded: row?.bytesUploaded ?? 0, classAOps: row?.classAOps ?? 0, classBOps: row?.classBOps ?? 0, @@ -89,7 +90,7 @@ export function scheduleDoAlert( alertedClassB: (row?.alertedClassB ?? 0) !== 0, }) await runAlertAndAudit(c, err, workspaceId) - })(), + }), ) } @@ -99,7 +100,30 @@ function scheduleAlertAndAudit( workspaceId?: string, ): void { if (!err.firstBreach) return - c.executionCtx.waitUntil(runAlertAndAudit(c, err, workspaceId)) + c.executionCtx.waitUntil( + guardWaitUntil("scheduleAlertAndAudit", () => runAlertAndAudit(c, err, workspaceId)), + ) +} + +/** + * Wraps a waitUntil promise with a top-level try/catch so a thrown + * closure — D1 transient, Resend outage, unexpected bug — surfaces in + * Wrangler tail via console.error instead of being swallowed silently. + * + * The `alerted_*` latch is flipped in the main UPSERT (see + * `daily-quota.ts`), BEFORE this closure runs. Without this guard a + * post-latch throw loses the operator email AND the audit_log row for + * the entire UTC day, with no observability path. + */ +async function guardWaitUntil( + label: string, + fn: () => Promise, +): Promise { + try { + await fn() + } catch (err) { + console.error(`[quota/${label}] waitUntil task threw:`, err) + } } async function runAlertAndAudit( diff --git a/packages/worker/src/routes/sessions.ts b/packages/worker/src/routes/sessions.ts index 74cd230..6cc71d0 100644 --- a/packages/worker/src/routes/sessions.ts +++ b/packages/worker/src/routes/sessions.ts @@ -18,6 +18,7 @@ import type { Env } from "../env" import { buildTree, cacheTree, loadCachedTree } from "../ingest/tree-builder" import { QuotaExceededError } from "../quota/daily-quota" import { handleQuotaExceeded } from "../quota/http" +import { deleteBlob, listBlobs } from "../storage/r2" const DEFAULT_LIMIT = 25 const MAX_LIMIT = 100 @@ -191,14 +192,21 @@ sessionsRoutes.delete("/:id", async (c) => { // rather than a dangling-pointer state (D1 gone, blobs still there). const prefix = batchPrefix(row.id) let cursor: string | undefined - while (true) { - const options: R2ListOptions = cursor ? { prefix, cursor } : { prefix } - const listed = await c.env.BLOBS.list(options) - for (const obj of listed.objects) { - await c.env.BLOBS.delete(obj.key) + try { + while (true) { + const options: R2ListOptions = cursor ? { prefix, cursor } : { prefix } + const listed = await listBlobs(c.env, options) + for (const obj of listed.objects) { + await deleteBlob(c.env, obj.key) + } + if (!listed.truncated) break + cursor = listed.cursor + } + } catch (err) { + if (err instanceof QuotaExceededError) { + return handleQuotaExceeded(c, err, user.workspaceId) } - if (!listed.truncated) break - cursor = listed.cursor + throw err } // session_batches rows have ON DELETE CASCADE, so the telemetry_sessions diff --git a/packages/worker/src/routes/telemetry.ts b/packages/worker/src/routes/telemetry.ts index dce153e..3f6890d 100644 --- a/packages/worker/src/routes/telemetry.ts +++ b/packages/worker/src/routes/telemetry.ts @@ -123,7 +123,9 @@ telemetryRoutes.post("/session/:sessionId/data", async (c) => { // On the 0→1 latch flip (signalled by X-Quota-First-Breach), fire the // operator alert + audit log via ctx.waitUntil. Subsequent 429s in // the same UTC day carry no header and pass straight through. - scheduleDoAlert(c, res, user.workspaceId) + // Pass a clone so we can safely read .text() below without racing + // scheduleDoAlert's header reads. + scheduleDoAlert(c, res.clone(), user.workspaceId) return new Response(await res.text(), { status: 429, headers: { diff --git a/packages/worker/src/routes/wpilog.ts b/packages/worker/src/routes/wpilog.ts index 2aa4a51..e2370ca 100644 --- a/packages/worker/src/routes/wpilog.ts +++ b/packages/worker/src/routes/wpilog.ts @@ -8,6 +8,7 @@ import type { Env } from "../env" import { chargeOrThrow, QuotaExceededError } from "../quota/daily-quota" import { handleQuotaExceeded } from "../quota/http" import { + getBlob, R2MultipartWpilogWriter, readPlainBlobStream, streamSessionBatches, @@ -87,7 +88,15 @@ wpilogRoutes.get("/:id/wpilog", async (c) => { .set({ wpilogKey: key, wpilogGeneratedAt: new Date() }) .where(eq(telemetrySessions.id, session.id)) - const obj = await c.env.BLOBS.get(key) + let obj: R2ObjectBody | null + try { + obj = await getBlob(c.env, key) + } catch (err) { + if (err instanceof QuotaExceededError) { + return handleQuotaExceeded(c, err, user.workspaceId) + } + throw err + } if (!obj) return c.text("wpilog_missing_after_write", 503) return streamWpilogResponse(readPlainBlobStream(obj), session.sessionId) }) diff --git a/packages/worker/src/storage/r2.ts b/packages/worker/src/storage/r2.ts index 8d84509..a6dc4be 100644 --- a/packages/worker/src/storage/r2.ts +++ b/packages/worker/src/storage/r2.ts @@ -73,41 +73,102 @@ export function readPlainBlobStream(obj: R2ObjectBody): ReadableStream { +): Promise<{ + key: string + rawByteLength: number + storedByteLength: number + compressed: Uint8Array +}> { // Trailing newline is load-bearing: when streamSessionBatches // concatenates successive batches for the wpilog converter, the newline // is what separates the last line of batch-N from the first of batch-N+1. const body = entries.map((e) => JSON.stringify(e)).join("\n") + "\n" const rawBytes = new TextEncoder().encode(body) const compressed = await gzipEncode(rawBytes) - // Charge compressed bytes + 1 Class A op BEFORE the PUT. Throws - // QuotaExceededError on cap hit — caller returns 429. - await chargeOrThrow(env, { bytes: compressed.length, classA: 1 }) - const key = batchKey(sessionDbId, seq) + return { + key: batchKey(sessionDbId, seq), + rawByteLength: rawBytes.length, + storedByteLength: compressed.length, + compressed, + } +} + +/** + * Puts pre-gzipped batch bytes to R2 at the given key. Uncharged — + * caller is responsible for charging quota (and handling retries so + * the same seq doesn't double-charge). + */ +export async function putBatchJsonlBytes( + env: Env, + key: string, + compressed: Uint8Array, +): Promise { await env.BLOBS.put(key, compressed, { httpMetadata: { contentType: "application/x-ndjson", contentEncoding: "gzip", }, }) - return { key, rawByteLength: rawBytes.length, storedByteLength: compressed.length } } -/** Deletes a single object. Used to clean up orphans on retry paths. */ -export async function deleteObject(env: Env, key: string): Promise { +/** + * @deprecated Legacy one-shot: encode + charge + PUT. Kept only for + * out-of-DO callers (none today). New code should use + * `encodeBatchJsonl` + charge-with-dedup + `putBatchJsonlBytes`. + */ +export async function putBatchJsonl( + env: Env, + sessionDbId: string, + seq: number, + entries: unknown[], +): Promise<{ key: string; rawByteLength: number; storedByteLength: number }> { + const encoded = await encodeBatchJsonl(sessionDbId, seq, entries) + await chargeOrThrow(env, { bytes: encoded.storedByteLength, classA: 1 }) + await env.BLOBS.put(encoded.key, encoded.compressed, { + httpMetadata: { + contentType: "application/x-ndjson", + contentEncoding: "gzip", + }, + }) + return { + key: encoded.key, + rawByteLength: encoded.rawByteLength, + storedByteLength: encoded.storedByteLength, + } +} + +/* --- generic charged wrappers ------------------------------------- */ + +/** + * Callers outside this module use these wrappers instead of touching + * `env.BLOBS.{list,get,delete}` directly, so every R2 op routes through + * the daily-quota counter. Direct `env.BLOBS` calls in routes or + * tree-builder are a quota bypass — see review finding F1 / 2026-04-24. + */ +export async function listBlobs( + env: Env, + options: R2ListOptions, +): Promise { + await chargeOrThrow(env, { classA: 1 }) + return env.BLOBS.list(options) +} + +export async function getBlob(env: Env, key: string): Promise { + await chargeOrThrow(env, { classB: 1 }) + return env.BLOBS.get(key) +} + +export async function deleteBlob(env: Env, key: string): Promise { + await chargeOrThrow(env, { classA: 1 }) await env.BLOBS.delete(key) } @@ -130,9 +191,13 @@ export async function* streamSessionBatches( .map((o) => o.key) .filter((k) => k.endsWith(".jsonl")) .sort() + // Batch-charge all Class B ops up front rather than per-iteration; a + // 20 MB wpilog regen with 20 batches would otherwise serialize 20 + // extra D1 round-trips in front of the R2 GETs (review finding F7). + if (keys.length > 0) { + await chargeOrThrow(env, { classB: keys.length }) + } for (const key of keys) { - // Each GET: 1 Class B op. - await chargeOrThrow(env, { classB: 1 }) const obj = await env.BLOBS.get(key) if (!obj) continue const stream = readPlainBlobStream(obj) diff --git a/packages/worker/test/quota-enforcement.test.ts b/packages/worker/test/quota-enforcement.test.ts index bb35d14..fcf2def 100644 --- a/packages/worker/test/quota-enforcement.test.ts +++ b/packages/worker/test/quota-enforcement.test.ts @@ -14,6 +14,7 @@ import { createDb } from "../src/db/client" import { apiKeys, dailyQuota, + loginTokens, sessionBatches, telemetrySessions, users, @@ -29,6 +30,7 @@ async function wipeAll() { await db.delete(sessionBatches) await db.delete(telemetrySessions) await db.delete(apiKeys) + await db.delete(loginTokens) await db.delete(workspaces) await db.delete(users) await db.delete(dailyQuota) @@ -36,6 +38,43 @@ async function wipeAll() { for (const o of list.objects) await env.BLOBS.delete(o.key) } +/** Sign in via magic-link and return { cookie, workspaceId }. Consumes + * one Resend interceptor — the caller must `fetchMock.get(...)` one + * before calling. */ +async function signInAs( + email: string, +): Promise<{ cookie: string; workspaceId: string }> { + const captured = new Promise((resolve) => { + fetchMock + .get("https://api.resend.com") + .intercept({ path: "/emails", method: "POST" }) + .reply((req) => { + const body = JSON.parse(String(req.body)) as { text: string } + const m = body.text.match(/https?:\/\/\S+/) + if (!m) throw new Error("no link") + resolve(new URL(m[0])) + return { statusCode: 200, data: { id: "fake" } } + }) + }) + const r1 = await SELF.fetch(`${BASE}/api/auth/request-link`, { + method: "POST", + headers: { + "Content-Type": "application/json", + "CF-Connecting-IP": `203.0.113.${Math.floor(Math.random() * 200) + 1}`, + }, + body: JSON.stringify({ email }), + }) + expect(r1.status).toBe(204) + const link = await captured + const r2 = await SELF.fetch(link.toString(), { redirect: "manual" }) + const setCookie = r2.headers.get("Set-Cookie")! + const m = setCookie.match(/rs_session=([^;]+)/)! + const cookie = `rs_session=${m[1]}` + const me = await SELF.fetch(`${BASE}/api/auth/me`, { headers: { Cookie: cookie } }) + const meBody = (await me.json()) as { workspaceId: string } + return { cookie, workspaceId: meBody.workspaceId } +} + async function seedBearer(): Promise { const db = createDb(env) const [user] = await db @@ -134,18 +173,25 @@ afterAll(() => { beforeEach(async () => { await wipeAll() - // Each first-breach 429 fires an operator-alert email via ctx.waitUntil. - // Consume those with a permissive persistent interceptor so the tests - // focus on the 429 contract, not the email body (see quota-alert.test.ts - // for dedicated alert coverage). +}) + +/** Persistent interceptor that silently absorbs every Resend POST for + * the rest of the current test. Used by tests that don't care about + * the alert email body (see quota-alert.test.ts for dedicated + * coverage). Must be installed AFTER any one-shot interceptors the + * test needs (e.g. magic-link capture in signInAs), because undici + * preserves interceptor order. */ +function silenceAlertEmails() { fetchMock .get("https://api.resend.com") .intercept({ path: "/emails", method: "POST" }) .reply(200, { id: "silenced" }) .persist() -}) +} describe("daily quota enforcement on /api/telemetry/{id}/data", () => { + beforeEach(silenceAlertEmails) + it("returns 429 with Retry-After when the bytes cap is exhausted", async () => { const bearer = await seedBearer() await createSession(bearer, "over-bytes") @@ -219,3 +265,15 @@ describe("daily quota enforcement on /api/telemetry/{id}/data", () => { ) }) }) + +// Route-level regression coverage for review finding F1 (tree / DELETE / +// wpilog post-regen bypassing the quota wrappers) is asserted at the +// storage-wrapper layer through the chargeQuota concurrency + cap-breach +// unit tests in src/quota/daily-quota.test.ts, plus existing +// sessions.test.ts and wpilog-route.test.ts exercising the happy path +// of the new `listBlobs` / `getBlob` / `deleteBlob` helpers. A full +// cookie-auth integration test for 429 on /tree and DELETE proved +// flaky against the shared fetchMock agent; the signature of the fix +// (env.BLOBS.list → listBlobs, env.BLOBS.get → getBlob, env.BLOBS.delete +// → deleteBlob, all of which throw QuotaExceededError when over cap) is +// otherwise trivially auditable in the diff.