From 3c2114890278003b2b5eb354e8e7ab69d836736f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 31 Jul 2026 22:46:43 +0000 Subject: [PATCH 01/39] perf(entitlements): cache enforcement plan lookups Co-authored-by: Kent C. Dodds --- .../contributing/architecture/entitlements.md | 15 +++++- .../entitlements/entitlements.node.test.ts | 52 +++++++++++++++++++ packages/worker/src/entitlements/service.ts | 9 ++-- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 916671ffbc..c0507cb62d 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -16,8 +16,9 @@ Module: `packages/worker/src/entitlements/` `resolveEffectivePlan(manual, stripe)`. - `errors.ts` — the one typed error (`EntitlementLimitError`) and the one user-facing message builder every enforcement point uses. -- `service.ts` — `getUserPlan`, `assertWithinEntitlement`, built-in D1 usage - counters, and the daily-counter helpers for rate-style limits. +- `service.ts` — `getUserPlan`, `getCachedUserPlan` (60s TTL enforcement cache), + `assertWithinEntitlement`, built-in D1 usage counters, and the daily-counter + helpers for rate-style limits. ## Plan model @@ -190,6 +191,16 @@ use stable-id reverse resolution. `findUserAccountByStableUserId` in verified-state (for example the outbound fetch gateway), mirroring `findUserAccount` in `email/platform-address.ts`. +**Enforcement plan cache:** `assertWithinEntitlement` resolves the plan limit +via `getCachedUserPlan`, which wraps `getUserPlan` behind a 60-second per-isolate +TTL cache keyed by D1 binding and `(stable_user_id, normalized email)`. Cache +entries share the lookup semantics above; failures are never cached. Built-in +usage counters and any `getCurrent` override are read fresh on every call, so +only the plan limit can be stale — a plan change may take up to ~60s to bind at +enforcement points while current usage stays accurate. Surfaces that display the +user's plan (billing UI, email usage) should keep calling `getUserPlan` +directly. + ## The error shape Every enforcement point throws `EntitlementLimitError` from diff --git a/packages/worker/src/entitlements/entitlements.node.test.ts b/packages/worker/src/entitlements/entitlements.node.test.ts index 3b89d4cb87..2b3585ce6b 100644 --- a/packages/worker/src/entitlements/entitlements.node.test.ts +++ b/packages/worker/src/entitlements/entitlements.node.test.ts @@ -596,6 +596,58 @@ test('assertWithinEntitlement passes under the limit, throws at it, and enforces expect(error.message).toBe(buildEntitlementLimitMessage(error.details)) }) +test('assertWithinEntitlement reuses cached plan within TTL while still enforcing usage', async () => { + const userId = await createStableUserIdFromEmail(plannedEmail) + const freeLimit = planLimits.free.maxScheduledJobs + const counts = { jobs: freeLimit - 1 } + const { db, queries } = createEntitlementsTestDb({ + users: [{ email: plannedEmail, plan: 'free', stable_user_id: userId }], + counts, + }) + const planQueries = () => + queries.filter((query) => + query.sql.includes('email = ? AND stable_user_id = ?'), + ) + const usageQueries = () => + queries.filter((query) => query.sql.includes('FROM jobs')) + + await assertWithinEntitlement({ + db, + userId, + email: plannedEmail, + resource: 'scheduled_jobs', + }) + await assertWithinEntitlement({ + db, + userId, + email: plannedEmail, + resource: 'scheduled_jobs', + }) + expect(planQueries()).toHaveLength(1) + expect(usageQueries()).toHaveLength(2) + + counts.jobs = freeLimit + const denied = await assertWithinEntitlement({ + db, + userId, + email: plannedEmail, + resource: 'scheduled_jobs', + }).then( + () => null, + (thrown: unknown) => thrown, + ) + if (!(denied instanceof EntitlementLimitError)) { + throw new Error('Expected an EntitlementLimitError.') + } + expect(denied.details).toMatchObject({ + plan: 'free', + limit: freeLimit, + current: freeLimit, + }) + expect(planQueries()).toHaveLength(1) + expect(usageQueries()).toHaveLength(3) +}) + test('assertWithinEntitlement enforces free concurrent workflow limit without email', async () => { const freeLimit = planLimits.free.maxConcurrentWorkflows const { db } = createEntitlementsTestDb({ diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index f528cbd6c4..c7a7f534a1 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -791,8 +791,9 @@ export type AssertWithinEntitlementInput = { /** * Account email of the acting user when available. When present, plan * lookup requires the email + stable-id pair. When absent (package-job / - * workflow contexts with `email: ''`), {@link getUserPlan} reverse-resolves - * by stable userId; unknown accounts still fail closed to `free`. + * workflow contexts with `email: ''`), {@link getCachedUserPlan} + * reverse-resolves by stable userId; unknown accounts still fail closed + * to `free`. */ email: string | null | undefined resource: EntitlementResource @@ -812,7 +813,9 @@ export type AssertWithinEntitlementInput = { export async function assertWithinEntitlement( input: AssertWithinEntitlementInput, ): Promise { - const plan = await getUserPlan(input.db, { + // Cached plan resolution: quota limit lookup tolerates short-TTL + // staleness; usage counters below stay fresh on every call. + const plan = await getCachedUserPlan(input.db, { userId: input.userId, email: input.email, }) From 7d5d77741135a5d9d071cdc68d51b9a4579a3a1f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 31 Jul 2026 22:54:38 +0000 Subject: [PATCH 02/39] style(docs): format entitlement cache notes Co-authored-by: Kent C. Dodds --- docs/contributing/architecture/entitlements.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index c0507cb62d..5b66773a99 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -192,14 +192,14 @@ verified-state (for example the outbound fetch gateway), mirroring `findUserAccount` in `email/platform-address.ts`. **Enforcement plan cache:** `assertWithinEntitlement` resolves the plan limit -via `getCachedUserPlan`, which wraps `getUserPlan` behind a 60-second per-isolate -TTL cache keyed by D1 binding and `(stable_user_id, normalized email)`. Cache -entries share the lookup semantics above; failures are never cached. Built-in -usage counters and any `getCurrent` override are read fresh on every call, so -only the plan limit can be stale — a plan change may take up to ~60s to bind at -enforcement points while current usage stays accurate. Surfaces that display the -user's plan (billing UI, email usage) should keep calling `getUserPlan` -directly. +via `getCachedUserPlan`, which wraps `getUserPlan` behind a 60-second +per-isolate TTL cache keyed by D1 binding and +`(stable_user_id, normalized email)`. Cache entries share the lookup semantics +above; failures are never cached. Built-in usage counters and any `getCurrent` +override are read fresh on every call, so only the plan limit can be stale — a +plan change may take up to ~60s to bind at enforcement points while current +usage stays accurate. Surfaces that display the user's plan (billing UI, email +usage) should keep calling `getUserPlan` directly. ## The error shape From 59a82385100195c41ce4a57ab35ae4cec5531e3a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:11:34 +0000 Subject: [PATCH 03/39] feat(entitlements): add per-user quota meter Co-authored-by: Kent C. Dodds --- .../src/admin/user-usage-data.node.test.ts | 187 ++++- packages/worker/src/admin/user-usage-data.ts | 11 +- .../src/app/account-usage-data.node.test.ts | 252 +++++-- packages/worker/src/app/account-usage-data.ts | 11 +- packages/worker/src/email/inbound-delivery.ts | 181 +++-- .../inbound-entitlements.workers.test.ts | 588 ++++++++++++++- packages/worker/src/email/inbound.ts | 8 + .../worker/src/email/inbound.workers.test.ts | 68 +- packages/worker/src/email/outbound.ts | 7 +- .../worker/src/email/outbound.workers.test.ts | 92 ++- .../entitlements/entitlements.node.test.ts | 111 ++- packages/worker/src/entitlements/service.ts | 358 +++++++-- .../src/entitlements/user-meter-client.ts | 28 + .../worker/src/entitlements/user-meter-do.ts | 678 ++++++++++++++++++ .../entitlements/user-meter.workers.test.ts | 616 ++++++++++++++++ packages/worker/src/index.ts | 2 + .../mcp/capabilities/email/email-usage-get.ts | 7 +- .../email/email-usage-get.workers.test.ts | 235 +++--- .../worker/src/mcp/fetch-gateway.node.test.ts | 18 +- packages/worker/src/mcp/fetch-gateway.ts | 4 +- .../src/mcp/fetch-gateway.workers.test.ts | 63 +- packages/worker/src/mcp/tools/execute.ts | 2 + .../worker/src/test-support/user-meter.ts | 182 +++++ ...er-scoped-durable-object-name.node.test.ts | 9 + .../src/user-scoped-durable-object-name.ts | 5 + packages/worker/worker-configuration.d.ts | 9 +- packages/worker/wrangler.jsonc | 16 + 27 files changed, 3293 insertions(+), 455 deletions(-) create mode 100644 packages/worker/src/entitlements/user-meter-client.ts create mode 100644 packages/worker/src/entitlements/user-meter-do.ts create mode 100644 packages/worker/src/entitlements/user-meter.workers.test.ts create mode 100644 packages/worker/src/test-support/user-meter.ts diff --git a/packages/worker/src/admin/user-usage-data.node.test.ts b/packages/worker/src/admin/user-usage-data.node.test.ts index b7e91e8228..cb538919ec 100644 --- a/packages/worker/src/admin/user-usage-data.node.test.ts +++ b/packages/worker/src/admin/user-usage-data.node.test.ts @@ -1,8 +1,15 @@ import { expect, test } from 'vitest' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' import { type AdminUsageRollup } from '#app/loader-data.ts' import { loadAdminUserUsageData } from './user-usage-data.ts' +function withUserMeter(env: { APP_DB: D1Database } & Record) { + const meter = createInMemoryUserMeterEnv() + return { ...env, ...meter.env, meter } +} + type UserRow = { id: number username: string @@ -161,7 +168,7 @@ test('loadAdminUserUsageData returns null for unknown users and zeroed usage for const emptyDb = createAdminUserUsageTestDb({ users: [] }) expect( await loadAdminUserUsageData( - { APP_DB: emptyDb } as Env, + withUserMeter({ APP_DB: emptyDb }) as Env, 'missing-stable-user', new Date('2026-07-05T12:00:00.000Z'), ), @@ -185,7 +192,7 @@ test('loadAdminUserUsageData returns null for unknown users and zeroed usage for }) const data = await loadAdminUserUsageData( - { APP_DB: db } as Env, + withUserMeter({ APP_DB: db }) as Env, usageUserId, new Date('2026-07-05T12:00:00.000Z'), ) @@ -254,7 +261,7 @@ test('loadAdminUserUsageData warns above eighty percent of plan limits', async ( }) const data = await loadAdminUserUsageData( - { APP_DB: db } as Env, + withUserMeter({ APP_DB: db }) as Env, usageUserId, new Date('2026-07-05T12:00:00.000Z'), ) @@ -299,7 +306,7 @@ test('loadAdminUserUsageData rejects an invalid stored plan', async () => { await expect( loadAdminUserUsageData( - { APP_DB: db } as Env, + withUserMeter({ APP_DB: db }) as Env, usageUserId, new Date('2026-07-05T12:00:00.000Z'), ), @@ -361,7 +368,10 @@ test('loadAdminUserUsageData caches rollup reads in KV and serves repeat loads f }, }) const { kv, store } = createFakeKv() - const env = { APP_DB: countingDb, BUNDLE_ARTIFACTS_KV: kv } as Env + const env = withUserMeter({ + APP_DB: countingDb, + BUNDLE_ARTIFACTS_KV: kv, + }) as Env const now = new Date('2026-07-05T12:00:00.000Z') const first = await loadAdminUserUsageData(env, usageUserId, now) @@ -413,7 +423,7 @@ test('loadAdminUserUsageData keeps current-month and month-over-month rollups on }) const data = await loadAdminUserUsageData( - { APP_DB: db } as Env, + withUserMeter({ APP_DB: db }) as Env, usageUserId, new Date('2026-07-01T00:00:00.000Z'), ) @@ -427,6 +437,171 @@ test('loadAdminUserUsageData keeps current-month and month-over-month rollups on expect(getEventCount(data?.monthUsage[1]?.usage, 'execute')).toBe(12) }) +test('loadAdminUserUsageData reads daily counts from UserMeter (bootstrap then warm)', async () => { + const now = new Date('2026-07-05T12:00:00.000Z') + const day = utcDayKey(now) + + const bootstrapEmail = 'bootstrap-drilldown@example.com' + const bootstrapUserId = await createStableUserIdFromEmail(bootstrapEmail) + const bootstrapped = await loadAdminUserUsageData( + withUserMeter({ + APP_DB: createAdminUserUsageTestDb({ + users: [ + { + id: 3, + username: 'bootstrap', + email: bootstrapEmail, + plan: 'pro', + stable_user_id: bootstrapUserId, + }, + ], + dailyCounters: [ + { + user_id: bootstrapUserId, + resource: 'email_receives_per_day', + day, + count: 55, + }, + { + user_id: bootstrapUserId, + resource: 'outbound_fetches_per_day', + day, + count: 66, + }, + ], + resourceCounts: { + [bootstrapUserId]: { secrets: 3 }, + }, + }), + }) as Env, + bootstrapUserId, + now, + ) + expect( + bootstrapped?.entitlementConsumption.find( + (row) => row.resource === 'email_receives_per_day', + )?.current, + ).toBe(55) + expect( + bootstrapped?.entitlementConsumption.find( + (row) => row.resource === 'outbound_fetches_per_day', + )?.current, + ).toBe(66) + expect( + bootstrapped?.entitlementConsumption.find( + (row) => row.resource === 'secrets', + )?.current, + ).toBe(3) + + const meterEmail = 'meter-drilldown@example.com' + const meterUserId = await createStableUserIdFromEmail(meterEmail) + const meter = createInMemoryUserMeterEnv() + await meter.seed({ + userId: meterUserId, + resource: 'email_sends_per_day', + day, + count: 111, + }) + await meter.seed({ + userId: meterUserId, + resource: 'email_receives_per_day', + day, + count: 222, + }) + await meter.seed({ + userId: meterUserId, + resource: 'execute_calls_per_day', + day, + count: 333, + }) + await meter.seed({ + userId: meterUserId, + resource: 'outbound_fetches_per_day', + day, + count: 444, + }) + const warm = await loadAdminUserUsageData( + { + APP_DB: createAdminUserUsageTestDb({ + users: [ + { + id: 4, + username: 'meter', + email: meterEmail, + plan: 'pro', + stable_user_id: meterUserId, + }, + ], + dailyCounters: [ + { + user_id: meterUserId, + resource: 'email_sends_per_day', + day, + count: 11, + }, + { + user_id: meterUserId, + resource: 'email_receives_per_day', + day, + count: 22, + }, + { + user_id: meterUserId, + resource: 'execute_calls_per_day', + day, + count: 33, + }, + { + user_id: meterUserId, + resource: 'outbound_fetches_per_day', + day, + count: 44, + }, + ], + resourceCounts: { + [meterUserId]: { + saved_packages: 6, + stored_email_messages: 9, + }, + }, + }), + ...meter.env, + } as Env, + meterUserId, + now, + ) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'email_sends_per_day', + )?.current, + ).toBe(111) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'email_receives_per_day', + )?.current, + ).toBe(222) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'execute_calls_per_day', + )?.current, + ).toBe(333) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'outbound_fetches_per_day', + )?.current, + ).toBe(444) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'saved_packages', + )?.current, + ).toBe(6) + expect( + warm?.entitlementConsumption.find( + (row) => row.resource === 'stored_email_messages', + )?.current, + ).toBe(9) +}) + function getEventCount( usage: Array | undefined, metric: AdminUsageRollup['metric'], diff --git a/packages/worker/src/admin/user-usage-data.ts b/packages/worker/src/admin/user-usage-data.ts index 9a1c29deea..221fdbd277 100644 --- a/packages/worker/src/admin/user-usage-data.ts +++ b/packages/worker/src/admin/user-usage-data.ts @@ -8,7 +8,7 @@ import { type EntitlementResource, type PlanName, } from '#worker/entitlements/plans.ts' -import { readEntitlementResourceUsage } from '#worker/entitlements/service.ts' +import { readCurrentEntitlementResourceUsage } from '#worker/entitlements/service.ts' import { createKvCachifiedCache } from '#worker/kv-cachified.ts' import { resolveUserStableId } from '#worker/user-id.ts' import { @@ -113,7 +113,7 @@ export async function loadAdminUserUsageData( currentMonth, }), readEntitlementConsumption({ - db: env.APP_DB, + env, usageUserId, plan, now, @@ -140,15 +140,16 @@ export async function loadAdminUserUsageData( } async function readEntitlementConsumption(input: { - db: D1Database + env: Env usageUserId: string plan: PlanName now: Date }): Promise> { return await Promise.all( adminUserUsageEntitlementResources.map(async (resource) => { - const current = await readEntitlementResourceUsage({ - db: input.db, + const current = await readCurrentEntitlementResourceUsage({ + db: input.env.APP_DB, + env: input.env, userId: input.usageUserId, resource, now: input.now, diff --git a/packages/worker/src/app/account-usage-data.node.test.ts b/packages/worker/src/app/account-usage-data.node.test.ts index 0281b04c4e..fe49db6938 100644 --- a/packages/worker/src/app/account-usage-data.node.test.ts +++ b/packages/worker/src/app/account-usage-data.node.test.ts @@ -1,6 +1,15 @@ import { expect, test } from 'vitest' -import { loadAccountUsageData } from '#app/account-usage-data.ts' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' +import { loadAccountUsageData } from '#app/account-usage-data.ts' + +type DailyCounterRow = { + user_id: string + resource: string + day: string + count: number +} function createUsageTestDb(input: { userId: number @@ -8,73 +17,192 @@ function createUsageTestDb(input: { plan: string stripePlan?: string | null packageCount?: number + dailyCounters?: Array }) { const stableUserId = testStableUserIdFromEmail(input.email) + const dailyCounters = input.dailyCounters?.map((row) => ({ ...row })) ?? [] return { - prepare(query: string) { - const normalized = query.replace(/\s+/g, ' ').trim().toLowerCase() - return { - bind(...params: Array) { - return { - async first() { - if ( - normalized.includes('from users') && - normalized.includes('where id') - ) { - return { - id: input.userId, - plan: input.plan, - stripe_plan: input.stripePlan ?? null, - stable_user_id: stableUserId, - } as T - } - if (normalized.includes('from saved_packages')) { - return { count: input.packageCount ?? 0 } as T - } - if (normalized.includes('from entitlement_daily_counters')) { - return { count: 0 } as T - } - if ( - normalized.includes('count(*)') || - normalized.includes('sum(') - ) { - return { count: 0, total: 0 } as T - } - void params - return null - }, - async all() { - return { results: [] } - }, - } - }, - } - }, - } as unknown as D1Database + stableUserId, + db: { + prepare(query: string) { + const normalized = query.replace(/\s+/g, ' ').trim().toLowerCase() + return { + bind(...params: Array) { + return { + async first() { + if ( + normalized.includes('from users') && + normalized.includes('where id') + ) { + return { + id: input.userId, + plan: input.plan, + stripe_plan: input.stripePlan ?? null, + stable_user_id: stableUserId, + } as T + } + if (normalized.includes('from saved_packages')) { + return { count: input.packageCount ?? 0 } as T + } + if (normalized.includes('from entitlement_daily_counters')) { + const row = dailyCounters.find( + (counter) => + counter.user_id === params[0] && + counter.resource === params[1] && + counter.day === params[2], + ) + return (row ? { count: row.count } : { count: 0 }) as T + } + if ( + normalized.includes('count(*)') || + normalized.includes('sum(') || + normalized.includes('d1_storage_bytes') + ) { + return { count: 0, total: 0, bytes: 0 } as T + } + void params + return null + }, + async all() { + return { results: [] } + }, + } + }, + } + }, + } as unknown as D1Database, + } } -test('loadAccountUsageData returns plan and entitlement rows for the account', async () => { - const env = { - APP_DB: createUsageTestDb({ - userId: 7, - email: 'usage@example.com', - plan: 'free', - packageCount: 2, - }), - } as Env +function currentFor( + data: Awaited>, + resource: string, +) { + return data?.entitlementConsumption.find((row) => row.resource === resource) + ?.current +} - const data = await loadAccountUsageData({ - env, +test('loadAccountUsageData returns plan rows and authoritative UserMeter daily counts', async () => { + const now = new Date('2026-07-25T12:00:00.000Z') + const day = utcDayKey(now) + + const baselineMeter = createInMemoryUserMeterEnv() + const { db: emptyDailyDb } = createUsageTestDb({ + userId: 7, + email: 'usage@example.com', + plan: 'free', + packageCount: 2, + }) + const baseline = await loadAccountUsageData({ + env: { APP_DB: emptyDailyDb, ...baselineMeter.env } as Env, userId: 7, - now: new Date('2026-07-25T12:00:00.000Z'), + now, + }) + expect(baseline?.ok).toBe(true) + expect(baseline?.plan).toBe('free') + expect(baseline?.today).toBe('2026-07-25') + expect(currentFor(baseline, 'saved_packages')).toBe(2) + expect(baseline?.entitlementConsumption.length).toBeGreaterThan(5) + + const bootstrapEmail = 'usage-bootstrap@example.com' + const bootstrapUserId = testStableUserIdFromEmail(bootstrapEmail) + const bootstrapMeter = createInMemoryUserMeterEnv() + const { db: bootstrapDb } = createUsageTestDb({ + userId: 8, + email: bootstrapEmail, + plan: 'pro', + packageCount: 1, + dailyCounters: [ + { + user_id: bootstrapUserId, + resource: 'email_sends_per_day', + day, + count: 17, + }, + { + user_id: bootstrapUserId, + resource: 'execute_calls_per_day', + day, + count: 91, + }, + ], + }) + const bootstrapped = await loadAccountUsageData({ + env: { APP_DB: bootstrapDb, ...bootstrapMeter.env } as Env, + userId: 8, + now, + }) + expect(currentFor(bootstrapped, 'email_sends_per_day')).toBe(17) + expect(currentFor(bootstrapped, 'execute_calls_per_day')).toBe(91) + expect(currentFor(bootstrapped, 'saved_packages')).toBe(1) + + const meterEmail = 'usage-meter@example.com' + const meterUserId = testStableUserIdFromEmail(meterEmail) + const warmMeter = createInMemoryUserMeterEnv() + const { db: warmDb } = createUsageTestDb({ + userId: 9, + email: meterEmail, + plan: 'pro', + packageCount: 4, + dailyCounters: [ + { + user_id: meterUserId, + resource: 'email_sends_per_day', + day, + count: 11, + }, + { + user_id: meterUserId, + resource: 'email_receives_per_day', + day, + count: 22, + }, + { + user_id: meterUserId, + resource: 'execute_calls_per_day', + day, + count: 33, + }, + { + user_id: meterUserId, + resource: 'outbound_fetches_per_day', + day, + count: 44, + }, + ], + }) + await warmMeter.seed({ + userId: meterUserId, + resource: 'email_sends_per_day', + day, + count: 101, + }) + await warmMeter.seed({ + userId: meterUserId, + resource: 'email_receives_per_day', + day, + count: 202, + }) + await warmMeter.seed({ + userId: meterUserId, + resource: 'execute_calls_per_day', + day, + count: 303, + }) + await warmMeter.seed({ + userId: meterUserId, + resource: 'outbound_fetches_per_day', + day, + count: 404, + }) + const authoritative = await loadAccountUsageData({ + env: { APP_DB: warmDb, ...warmMeter.env } as Env, + userId: 9, + now, }) - expect(data?.ok).toBe(true) - expect(data?.plan).toBe('free') - expect(data?.today).toBe('2026-07-25') - const packages = data?.entitlementConsumption.find( - (row) => row.resource === 'saved_packages', - ) - expect(packages?.current).toBe(2) - expect(packages?.limit).toBeGreaterThan(0) - expect(data?.entitlementConsumption.length).toBeGreaterThan(5) + expect(currentFor(authoritative, 'email_sends_per_day')).toBe(101) + expect(currentFor(authoritative, 'email_receives_per_day')).toBe(202) + expect(currentFor(authoritative, 'execute_calls_per_day')).toBe(303) + expect(currentFor(authoritative, 'outbound_fetches_per_day')).toBe(404) + expect(currentFor(authoritative, 'saved_packages')).toBe(4) }) diff --git a/packages/worker/src/app/account-usage-data.ts b/packages/worker/src/app/account-usage-data.ts index 6a41dd74d6..fe416ee63d 100644 --- a/packages/worker/src/app/account-usage-data.ts +++ b/packages/worker/src/app/account-usage-data.ts @@ -7,7 +7,7 @@ import { type EntitlementResource, type PlanName, } from '#worker/entitlements/plans.ts' -import { readEntitlementResourceUsage } from '#worker/entitlements/service.ts' +import { readCurrentEntitlementResourceUsage } from '#worker/entitlements/service.ts' import { resolveUserStableId } from '#worker/user-id.ts' import { type AccountUsageEntitlementConsumption, @@ -62,7 +62,7 @@ export async function loadAccountUsageData(input: { ) const usageUserId = resolveUserStableId(row) const entitlementConsumption = await readEntitlementConsumption({ - db: input.env.APP_DB, + env: input.env, usageUserId, plan, now, @@ -78,15 +78,16 @@ export async function loadAccountUsageData(input: { } async function readEntitlementConsumption(input: { - db: D1Database + env: Env usageUserId: string plan: PlanName now: Date }): Promise> { return await Promise.all( accountUsageEntitlementResources.map(async (resource) => { - const current = await readEntitlementResourceUsage({ - db: input.db, + const current = await readCurrentEntitlementResourceUsage({ + db: input.env.APP_DB, + env: input.env, userId: input.usageUserId, resource, now: input.now, diff --git a/packages/worker/src/email/inbound-delivery.ts b/packages/worker/src/email/inbound-delivery.ts index 6102a1df1d..f5499bf1fb 100644 --- a/packages/worker/src/email/inbound-delivery.ts +++ b/packages/worker/src/email/inbound-delivery.ts @@ -5,6 +5,11 @@ import { EntitlementLimitError, } from '#worker/entitlements/errors.ts' import { type PlanName } from '#worker/entitlements/plans.ts' +import { scheduleAbsoluteDailyEntitlementMirror } from '#worker/entitlements/service.ts' +import { + userMeterRpc, + type UserMeterEnv, +} from '#worker/entitlements/user-meter-client.ts' import { normalizeEmailAddress } from './address.ts' import { emailRawMimeKey, @@ -447,17 +452,47 @@ async function readCounter(input: { return Number(row?.count ?? 0) } +/** Point-read today's user inbound receive count from UserMeter. */ export async function readUserInboundReceiveCount(input: { db: D1Database + env: UserMeterEnv userId: string day: string + now?: Date }) { - return await readCounter({ - db: input.db, - table: 'entitlement_daily_counters', - userId: input.userId, + const now = input.now ?? new Date() + const updatedAt = now.toISOString() + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + let result = await meter.read({ + resource: 'email_receives_per_day', day: input.day, + now: updatedAt, }) + if (result.outcome === 'needs_bootstrap') { + const baseline = await readCounter({ + db: input.db, + table: 'entitlement_daily_counters', + userId: input.userId, + day: input.day, + }) + await meter.initialize({ + resource: 'email_receives_per_day', + day: input.day, + count: baseline, + updatedAt, + }) + result = await meter.read({ + resource: 'email_receives_per_day', + day: input.day, + now: updatedAt, + }) + if (result.outcome === 'needs_bootstrap') { + throw new Error( + 'UserMeter inbound receive read still needs bootstrap after initialize.', + ) + } + } + return result.count } export async function readSystemInboundReceiveCount(input: { @@ -484,85 +519,101 @@ async function resolveConcurrentDelivery(input: { }) } +/** + * Charge one inbound receive via UserMeter, then persist the D1 delivery + * event. Delivery-id idempotency lives in the DO; D1 mirror is best-effort. + */ export async function chargeUserInboundDeliveryOnce(input: { db: D1Database + env: UserMeterEnv delivery: InboundDelivery plan: PlanName limit: number now: Date + waitUntil?: (promise: Promise) => void }): Promise { const existing = await resolveConcurrentDelivery(input) if (existing) return existing - const current = await readCounter({ - db: input.db, - table: 'entitlement_daily_counters', + + const updatedAt = input.now.toISOString() + const meter = userMeterRpc({ + env: input.env, userId: input.delivery.userId, + }) + let meterResult = await meter.consumeInboundDelivery({ + deliveryId: input.delivery.deliveryId, + resource: 'email_receives_per_day', day: input.delivery.quotaDay, + limit: input.limit, + updatedAt, }) - const throwLimitError = (): never => { + if (meterResult.outcome === 'needs_bootstrap') { + const baseline = await readCounter({ + db: input.db, + table: 'entitlement_daily_counters', + userId: input.delivery.userId, + day: input.delivery.quotaDay, + }) + await meter.initialize({ + resource: 'email_receives_per_day', + day: input.delivery.quotaDay, + count: baseline, + updatedAt, + }) + meterResult = await meter.consumeInboundDelivery({ + deliveryId: input.delivery.deliveryId, + resource: 'email_receives_per_day', + day: input.delivery.quotaDay, + limit: input.limit, + updatedAt, + }) + if (meterResult.outcome === 'needs_bootstrap') { + throw new Error( + 'UserMeter inbound delivery consume still needs bootstrap after initialize.', + ) + } + } + + const accepted = meterResult.consumed || meterResult.replayed + if (!accepted) { throw new EntitlementLimitError({ resource: 'email_receives_per_day', plan: input.plan, limit: input.limit, - current, + current: meterResult.count, upgradeHint: buildEntitlementUpgradeHint('email_receives_per_day'), }) } - if (input.limit < 1 || current >= input.limit) throwLimitError() - const operationTimestamp = randomOperationTimestamp(input.now) + + scheduleAbsoluteDailyEntitlementMirror({ + db: input.db, + userId: input.delivery.userId, + resource: meterResult.resource, + day: meterResult.day, + count: meterResult.count, + mirrorUpdatedAt: meterResult.mirrorUpdatedAt, + waitUntil: input.waitUntil, + }) + try { - const results = await input.db.batch([ - input.db - .prepare( - `INSERT INTO entitlement_daily_counters ( - user_id, resource, day, count, updated_at - ) VALUES (?, 'email_receives_per_day', ?, 1, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = entitlement_daily_counters.count + 1, - updated_at = excluded.updated_at - WHERE entitlement_daily_counters.count + 1 <= ? - AND NOT EXISTS ( - SELECT 1 FROM email_delivery_events - WHERE id = ? AND user_id = ? - )`, - ) - .bind( - input.delivery.userId, - input.delivery.quotaDay, - operationTimestamp, - input.limit, - input.delivery.deliveryId, - input.delivery.userId, - ), - input.db - .prepare( - `INSERT OR IGNORE INTO email_delivery_events ( - id, message_id, user_id, inbox_id, event_type, provider, - provider_event_id, detail_json, created_at - ) - SELECT ?, NULL, ?, ?, 'receive_started', ?, ?, ?, ? - WHERE EXISTS ( - SELECT 1 FROM entitlement_daily_counters - WHERE user_id = ? - AND resource = 'email_receives_per_day' - AND day = ? - AND updated_at = ? - )`, - ) - .bind( - input.delivery.deliveryId, - input.delivery.userId, - input.delivery.inboxId, - inboundProvider, - input.delivery.deliveryId, - JSON.stringify(input.delivery), - input.now.toISOString(), - input.delivery.userId, - input.delivery.quotaDay, - operationTimestamp, - ), - ]) - if (Number(results[1]?.meta.changes ?? 0) > 0) return input.delivery + const insert = await input.db + .prepare( + `INSERT OR IGNORE INTO email_delivery_events ( + id, message_id, user_id, inbox_id, event_type, provider, + provider_event_id, detail_json, created_at + ) VALUES (?, NULL, ?, ?, 'receive_started', ?, ?, ?, ?)`, + ) + .bind( + input.delivery.deliveryId, + input.delivery.userId, + input.delivery.inboxId, + inboundProvider, + input.delivery.deliveryId, + JSON.stringify(input.delivery), + input.now.toISOString(), + ) + .run() + if (Number(insert.meta.changes ?? 0) > 0) return input.delivery } catch (error) { const committed = await resolveConcurrentDelivery(input).catch(() => null) if (committed) return committed @@ -570,7 +621,9 @@ export async function chargeUserInboundDeliveryOnce(input: { } const raced = await resolveConcurrentDelivery(input) if (raced) return raced - return throwLimitError() + throw new Error( + 'Inbound delivery charge was accepted but the delivery event was not persisted.', + ) } export async function chargeSystemInboundDeliveryOnce(input: { diff --git a/packages/worker/src/email/inbound-entitlements.workers.test.ts b/packages/worker/src/email/inbound-entitlements.workers.test.ts index b32a87ac98..41757b39cd 100644 --- a/packages/worker/src/email/inbound-entitlements.workers.test.ts +++ b/packages/worker/src/email/inbound-entitlements.workers.test.ts @@ -1,9 +1,23 @@ +import { runInDurableObject } from 'cloudflare:test' import { env } from 'cloudflare:workers' import { expect, test } from 'vitest' -import { maxPlanEmailLimits, planLimits } from '#worker/entitlements/plans.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { isEntitlementLimitError } from '#worker/entitlements/errors.ts' +import { maxPlanEmailLimits, planLimits } from '#worker/entitlements/plans.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' +import { UserMeter } from '#worker/entitlements/user-meter-do.ts' +import { + createWaitUntilDrain, + withPatchedDbPrepare, +} from '#worker/test-support/user-meter.ts' import { ensureUsageRollupsTestSchema } from '#worker/usage/test-schema.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' +import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' +import { + buildInboundDelivery, + chargeUserInboundDeliveryOnce, + readUserInboundReceiveCount, +} from './inbound-delivery.ts' import { handleInboundEmail } from './inbound.ts' import { listEmailMessages } from './repo.ts' import { maxDetailedEmailRejectionEventsPerDay } from './service.ts' @@ -11,7 +25,6 @@ import { createForwardableEmailMessage } from './test-fixtures.ts' import { ensureEmailTestSchema } from './test-schema.ts' const platformBaseUrl = 'https://kody.example.com' -// User mail lives on the inbox. subdomain derived from APP_BASE_URL. const platformDomain = 'inbox.kody.example.com' function createInboundEnv() { @@ -40,7 +53,7 @@ async function seedAccountWithPlan(input: { stableUserId, ) .run() - return { username, address: `${username}@${platformDomain}` } + return { username, address: `${username}@${platformDomain}`, stableUserId } } function buildInboundMessage(address: string) { @@ -58,23 +71,45 @@ function buildInboundMessage(address: string) { }) } -async function setDailyReceiveCounter(userId: string, count: number) { +async function seedDailyReceiveCounter(userId: string, count: number) { + const day = utcDayKey() + const updatedAt = new Date().toISOString() + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(userId)), + ) + await runInDurableObject(stub, async (instance: UserMeter, state) => { + expect(instance).toBeInstanceOf(UserMeter) + await instance.read({ resource: 'email_receives_per_day', day }) + state.storage.sql.exec( + `INSERT INTO daily_counters (resource, day, count, revision, updated_at) + VALUES (?, ?, ?, 1, ?) + ON CONFLICT(resource, day) DO UPDATE SET + count = excluded.count, + revision = excluded.revision, + updated_at = excluded.updated_at`, + 'email_receives_per_day', + day, + count, + updatedAt, + ) + }) await env.APP_DB.prepare( `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, 'email_receives_per_day', ?, ?, ?)`, + VALUES (?, 'email_receives_per_day', ?, ?, ?) + ON CONFLICT(user_id, resource, day) DO UPDATE SET + count = excluded.count, + updated_at = excluded.updated_at`, ) - .bind(userId, utcDayKey(), count, new Date().toISOString()) + .bind(userId, day, count, updatedAt) .run() } async function readDailyReceiveCounter(userId: string) { - const row = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, - ) - .bind(userId, utcDayKey()) - .first<{ count: number }>() - return Number(row?.count ?? 0) + const result = await userMeterRpc({ env, userId }).read({ + resource: 'email_receives_per_day', + day: utcDayKey(), + }) + return result.outcome === 'ready' ? result.count : 0 } async function bulkInsertStoredMessages(userId: string, count: number) { @@ -141,6 +176,71 @@ async function readEmailReceivedRollup(userId: string) { .first<{ event_count: number; error_count: number; total_bytes: number }>() } +async function buildChargeableDelivery(input: { + userId: string + inboxId: string + address: string + raw?: string +}) { + const now = new Date() + const raw = + input.raw ?? + [ + 'From: Sender ', + `To: ${input.address}`, + 'Subject: Charge once', + `Message-ID: `, + '', + 'Body', + ].join('\r\n') + return await buildInboundDelivery({ + userId: input.userId, + inboxId: input.inboxId, + recipient: input.address, + envelopeFrom: 'sender@example.net', + rawMime: raw, + quotaDay: utcDayKey(now), + now, + }) +} + +function withFailingReceiveStartedInsert(input: { + deliveryId: string + message: string +}) { + let failNext = true + return withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + const statement = originalPrepare(query) + const isDeliveryInsert = + query.includes('INSERT OR IGNORE INTO email_delivery_events') && + query.includes("'receive_started'") + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: bound.first.bind(bound), + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: async () => { + if ( + isDeliveryInsert && + failNext && + params[0] === input.deliveryId + ) { + failNext = false + throw new Error(input.message) + } + return await bound.run() + }, + } + }, + } + }) as D1Database['prepare'] + }) +} + test('inbound email enforces free-plan receive, storage, and size limits then stores under quota', async () => { await ensureEmailTestSchema(env.APP_DB) await ensureUsageRollupsTestSchema(env.APP_DB) @@ -154,7 +254,7 @@ test('inbound email enforces free-plan receive, storage, and size limits then st email: receiveLimitEmail, plan: 'free', }) - await setDailyReceiveCounter(receiveLimitUserId, receiveLimit) + await seedDailyReceiveCounter(receiveLimitUserId, receiveLimit) const receiveLimitMessage = buildInboundMessage(receiveLimitAddress) await handleInboundEmail(receiveLimitMessage, createInboundEnv()) @@ -307,7 +407,7 @@ test('inbound email applies the max-plan email receive backstop', async () => { const userId = await createStableUserIdFromEmail(email) const { address } = await seedAccountWithPlan({ email, plan: 'max' }) const receiveLimit = maxPlanEmailLimits.email_receives_per_day - await setDailyReceiveCounter(userId, receiveLimit) + await seedDailyReceiveCounter(userId, receiveLimit) const message = buildInboundMessage(address) await handleInboundEmail(message, createInboundEnv()) @@ -328,7 +428,7 @@ test('inbound email stores at most five detailed rejection events per inbox per const limit = planLimits.free.maxEmailReceivesPerDay if (limit === null) throw new Error('Expected a numeric free limit.') const { address } = await seedAccountWithPlan({ email, plan: 'free' }) - await setDailyReceiveCounter(userId, limit) + await seedDailyReceiveCounter(userId, limit) const attempts = maxDetailedEmailRejectionEventsPerDay + 3 for (let index = 0; index < attempts; index += 1) { @@ -351,3 +451,459 @@ test('inbound email stores at most five detailed rejection events per inbox per error_count: attempts, }) }) + +test('inbound UserMeter delivery claims are idempotent, isolated, and over-limit exact', async () => { + await ensureEmailTestSchema(env.APP_DB) + const receiveLimit = planLimits.free.maxEmailReceivesPerDay + if (receiveLimit === null) throw new Error('Expected a numeric free limit.') + + const userAEmail = `meter-a-${crypto.randomUUID()}@example.com` + const userBEmail = `meter-b-${crypto.randomUUID()}@example.com` + const userAId = await createStableUserIdFromEmail(userAEmail) + const userBId = await createStableUserIdFromEmail(userBEmail) + const userA = await seedAccountWithPlan({ email: userAEmail, plan: 'free' }) + const userB = await seedAccountWithPlan({ email: userBEmail, plan: 'free' }) + const inboxA = `inbox-a-${crypto.randomUUID()}` + const inboxB = `inbox-b-${crypto.randomUUID()}` + const now = new Date() + const day = utcDayKey(now) + const drain = createWaitUntilDrain() + + const deliveryA = await buildChargeableDelivery({ + userId: userAId, + inboxId: inboxA, + address: userA.address, + }) + const deliveryARetry = { ...deliveryA } + const sharedDeliveryId = deliveryA.deliveryId + + const concurrent = await Promise.all([ + chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: deliveryA, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }), + chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: deliveryARetry, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }), + ]) + expect(concurrent[0]?.deliveryId).toBe(sharedDeliveryId) + expect(concurrent[1]?.deliveryId).toBe(sharedDeliveryId) + expect(await readDailyReceiveCounter(userAId)).toBe(1) + expect( + await env.APP_DB.prepare( + `SELECT COUNT(*) AS count FROM email_delivery_events + WHERE id = ? AND user_id = ?`, + ) + .bind(sharedDeliveryId, userAId) + .first<{ count: number }>(), + ).toEqual({ count: 1 }) + + const failingDelivery = await buildChargeableDelivery({ + userId: userAId, + inboxId: inboxA, + address: userA.address, + raw: [ + 'From: Sender ', + `To: ${userA.address}`, + 'Subject: Fail then retry', + `Message-ID: `, + '', + 'Retry body.', + ].join('\r\n'), + }) + using _patch = withFailingReceiveStartedInsert({ + deliveryId: failingDelivery.deliveryId, + message: 'simulated D1 delivery event insert failure', + }) + + await expect( + chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: failingDelivery, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }), + ).rejects.toThrow('simulated D1 delivery event insert failure') + expect(await readDailyReceiveCounter(userAId)).toBe(2) + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_delivery_events + WHERE id = ? AND user_id = ?`, + ) + .bind(failingDelivery.deliveryId, userAId) + .first(), + ).toBeNull() + + const recovered = await chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: failingDelivery, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }) + expect(recovered.deliveryId).toBe(failingDelivery.deliveryId) + expect(await readDailyReceiveCounter(userAId)).toBe(2) + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_delivery_events + WHERE id = ? AND user_id = ?`, + ) + .bind(failingDelivery.deliveryId, userAId) + .first(), + ).not.toBeNull() + + const deliveryB = await buildChargeableDelivery({ + userId: userBId, + inboxId: inboxB, + address: userB.address, + }) + await chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: deliveryB, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }) + expect(await readDailyReceiveCounter(userAId)).toBe(2) + expect(await readDailyReceiveCounter(userBId)).toBe(1) + expect( + await readUserInboundReceiveCount({ + db: env.APP_DB, + env, + userId: userBId, + day, + now, + }), + ).toBe(1) + + const sharedClaimId = `email-inbound-delivery:shared-${crypto.randomUUID()}` + const meterA = userMeterRpc({ env, userId: userAId }) + const meterB = userMeterRpc({ env, userId: userBId }) + const claimA = await meterA.consumeInboundDelivery({ + deliveryId: sharedClaimId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt: now.toISOString(), + }) + const claimB = await meterB.consumeInboundDelivery({ + deliveryId: sharedClaimId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt: now.toISOString(), + }) + expect(claimA).toMatchObject({ + outcome: 'ready', + consumed: true, + replayed: false, + count: 3, + day, + resource: 'email_receives_per_day', + }) + expect(claimB).toMatchObject({ + outcome: 'ready', + consumed: true, + replayed: false, + count: 2, + day, + resource: 'email_receives_per_day', + }) + const replayA = await meterA.consumeInboundDelivery({ + deliveryId: sharedClaimId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt: now.toISOString(), + }) + expect(replayA).toMatchObject({ + outcome: 'ready', + consumed: false, + replayed: true, + count: 3, + day, + resource: 'email_receives_per_day', + }) + expect(await readDailyReceiveCounter(userAId)).toBe(3) + expect(await readDailyReceiveCounter(userBId)).toBe(2) + + await seedDailyReceiveCounter(userAId, receiveLimit) + const overLimitDelivery = await buildChargeableDelivery({ + userId: userAId, + inboxId: inboxA, + address: userA.address, + raw: [ + 'From: Sender ', + `To: ${userA.address}`, + 'Subject: Over limit', + `Message-ID: `, + '', + 'Denied.', + ].join('\r\n'), + }) + const denied = await chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: overLimitDelivery, + plan: 'free', + limit: receiveLimit, + now, + waitUntil: drain.waitUntil, + }).then( + () => null, + (thrown: unknown) => thrown, + ) + expect(isEntitlementLimitError(denied)).toBe(true) + expect(denied).toMatchObject({ + details: { + code: 'entitlement_limit_exceeded', + resource: 'email_receives_per_day', + plan: 'free', + limit: receiveLimit, + current: receiveLimit, + }, + }) + expect(await readDailyReceiveCounter(userAId)).toBe(receiveLimit) + await drain.drain() +}, 30_000) + +test('inbound UserMeter claim+consume rolls back together when claim insert fails', async () => { + await ensureEmailTestSchema(env.APP_DB) + const email = `meter-atomic-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(email) + await seedAccountWithPlan({ email, plan: 'free' }) + const day = '2026-07-30' + const updatedAt = '2026-07-30T12:00:00.000Z' + const deliveryId = `email-inbound-delivery:atomic-${crypto.randomUUID()}` + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(userId)), + ) + + await runInDurableObject(stub, async (instance: UserMeter, state) => { + expect(instance).toBeInstanceOf(UserMeter) + await instance.initialize({ + resource: 'email_receives_per_day', + day, + count: 0, + updatedAt, + }) + + const sql = state.storage.sql + const originalExec = sql.exec.bind(sql) + sql.exec = ((query: string, ...bindings: Array) => { + if ( + typeof query === 'string' && + query.includes('INSERT INTO inbound_delivery_claims') + ) { + throw new Error('injected inbound claim insert failure') + } + return originalExec(query, ...bindings) + }) as typeof sql.exec + + const receiveLimit = planLimits.free.maxEmailReceivesPerDay + if (receiveLimit === null) throw new Error('Expected a numeric free limit.') + + await expect( + instance.consumeInboundDelivery({ + deliveryId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt, + }), + ).rejects.toThrow('injected inbound claim insert failure') + + sql.exec = originalExec + + expect( + await instance.read({ resource: 'email_receives_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 0 }) + expect( + state.storage.sql + .exec<{ n: number }>( + `SELECT COUNT(*) AS n FROM inbound_delivery_claims + WHERE delivery_id = ?`, + deliveryId, + ) + .toArray()[0]?.n, + ).toBe(0) + + const accepted = await instance.consumeInboundDelivery({ + deliveryId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt, + }) + expect(accepted).toMatchObject({ + outcome: 'ready', + consumed: true, + replayed: false, + count: 1, + day, + resource: 'email_receives_per_day', + }) + expect( + state.storage.sql + .exec<{ n: number }>( + `SELECT COUNT(*) AS n FROM inbound_delivery_claims + WHERE delivery_id = ?`, + deliveryId, + ) + .toArray()[0]?.n, + ).toBe(1) + + const foreignDeliveryId = `email-inbound-delivery:foreign-${crypto.randomUUID()}` + state.storage.sql.exec( + `INSERT INTO inbound_delivery_claims ( + delivery_id, resource, day, count_after, revision, claimed_at + ) VALUES (?, 'email_sends_per_day', ?, 1, 1, ?)`, + foreignDeliveryId, + day, + updatedAt, + ) + await expect( + instance.consumeInboundDelivery({ + deliveryId: foreignDeliveryId, + resource: 'email_receives_per_day', + day, + limit: receiveLimit, + updatedAt, + }), + ).rejects.toThrow(/was claimed for "email_sends_per_day"/) + }) +}) + +test('inbound UserMeter replay keeps original claim day across UTC midnight', async () => { + await ensureEmailTestSchema(env.APP_DB) + const receiveLimit = planLimits.free.maxEmailReceivesPerDay + if (receiveLimit === null) throw new Error('Expected a numeric free limit.') + + const email = `meter-cross-day-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(email) + const account = await seedAccountWithPlan({ email, plan: 'free' }) + const inboxId = `inbox-cross-${crypto.randomUUID()}` + const dayN = '2026-07-30' + const dayN1 = '2026-07-31' + const drain = createWaitUntilDrain() + const meter = userMeterRpc({ env, userId }) + + await meter.initialize({ + resource: 'email_receives_per_day', + day: dayN, + count: 0, + updatedAt: `${dayN}T23:50:00.000Z`, + }) + + const delivery = await buildInboundDelivery({ + userId, + inboxId, + recipient: account.address, + envelopeFrom: 'sender@example.net', + rawMime: [ + 'From: Sender ', + `To: ${account.address}`, + 'Subject: Cross-day retry', + `Message-ID: `, + '', + 'Body', + ].join('\r\n'), + quotaDay: dayN, + now: new Date(`${dayN}T23:59:00.000Z`), + }) + + using _patch = withFailingReceiveStartedInsert({ + deliveryId: delivery.deliveryId, + message: 'simulated cross-day D1 insert failure', + }) + + await expect( + chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery, + plan: 'free', + limit: receiveLimit, + now: new Date(`${dayN}T23:59:30.000Z`), + waitUntil: drain.waitUntil, + }), + ).rejects.toThrow('simulated cross-day D1 insert failure') + + expect( + await meter.read({ resource: 'email_receives_per_day', day: dayN }), + ).toMatchObject({ outcome: 'ready', count: 1 }) + expect( + await meter.read({ resource: 'email_receives_per_day', day: dayN1 }), + ).toMatchObject({ outcome: 'needs_bootstrap' }) + + const replay = await meter.consumeInboundDelivery({ + deliveryId: delivery.deliveryId, + resource: 'email_receives_per_day', + day: dayN1, + limit: receiveLimit, + updatedAt: `${dayN1}T00:01:00.000Z`, + }) + expect(replay).toMatchObject({ + outcome: 'ready', + consumed: false, + replayed: true, + count: 1, + day: dayN, + resource: 'email_receives_per_day', + }) + + const recovered = await chargeUserInboundDeliveryOnce({ + db: env.APP_DB, + env, + delivery: { + ...delivery, + quotaDay: dayN1, + }, + plan: 'free', + limit: receiveLimit, + now: new Date(`${dayN1}T00:01:00.000Z`), + waitUntil: drain.waitUntil, + }) + expect(recovered.deliveryId).toBe(delivery.deliveryId) + await drain.drain() + + expect( + await meter.read({ resource: 'email_receives_per_day', day: dayN }), + ).toMatchObject({ outcome: 'ready', count: 1 }) + expect( + await meter.read({ resource: 'email_receives_per_day', day: dayN1 }), + ).toMatchObject({ outcome: 'needs_bootstrap' }) + + const mirroredDayN = await env.APP_DB.prepare( + `SELECT count FROM entitlement_daily_counters + WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, + ) + .bind(userId, dayN) + .first<{ count: number }>() + const mirroredDayN1 = await env.APP_DB.prepare( + `SELECT count FROM entitlement_daily_counters + WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, + ) + .bind(userId, dayN1) + .first<{ count: number }>() + expect(Number(mirroredDayN?.count ?? 0)).toBe(1) + expect(mirroredDayN1).toBeNull() +}, 30_000) diff --git a/packages/worker/src/email/inbound.ts b/packages/worker/src/email/inbound.ts index 31e24e5fa5..4f5a72c58e 100644 --- a/packages/worker/src/email/inbound.ts +++ b/packages/worker/src/email/inbound.ts @@ -196,6 +196,7 @@ export async function handleInboundEmail( | 'APP_BASE_URL' | 'USER_EMAIL_DOMAIN' | 'USAGE_EVENTS' + | 'USER_METER' >, ctx?: ExecutionContext, ) { @@ -482,8 +483,10 @@ export async function handleInboundEmail( ) const receivesToday = await readUserInboundReceiveCount({ db: env.APP_DB, + env, userId, day: userInboundQuotaDay(quotaNow), + now: quotaNow, }) if (receivesToday >= receiveLimit) { message.setReject('Recipient mailbox is over quota.') @@ -530,12 +533,16 @@ export async function handleInboundEmail( }) } } + const waitUntil = ctx + ? (promise: Promise) => ctx.waitUntil(promise) + : undefined let claimedDelivery: InboundDelivery try { claimedDelivery = existingDelivery ?? (await chargeUserInboundDeliveryOnce({ db: env.APP_DB, + env, delivery: { ...delivery, quotaDay: userInboundQuotaDay(quotaNow), @@ -546,6 +553,7 @@ export async function handleInboundEmail( 'email_receives_per_day', ), now: quotaNow, + waitUntil, })) } catch (error) { if (!isEntitlementLimitError(error)) throw error diff --git a/packages/worker/src/email/inbound.workers.test.ts b/packages/worker/src/email/inbound.workers.test.ts index eeda94dd88..7d9ff38520 100644 --- a/packages/worker/src/email/inbound.workers.test.ts +++ b/packages/worker/src/email/inbound.workers.test.ts @@ -1,5 +1,6 @@ import { env } from 'cloudflare:workers' import { expect, test, vi } from 'vitest' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' import { handleInboundEmail } from './inbound.ts' import { processInboundDeliveryEffects } from './inbound-effects.ts' import { @@ -848,13 +849,18 @@ test('inbound email stores raw MIME in R2 and readers and deletes follow the blo }) async function readUserDailyReceiveCount(userId: string) { - const row = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, - ) - .bind(userId, new Date().toISOString().slice(0, 10)) - .first<{ count: number }>() - return Number(row?.count ?? 0) + const result = await userMeterRpc({ env, userId }).read({ + resource: 'email_receives_per_day', + day: new Date().toISOString().slice(0, 10), + }) + return result.outcome === 'ready' ? result.count : 0 +} + +async function readUserTotalReceiveCount(userId: string) { + const exported = await userMeterRpc({ env, userId }).exportCounters({}) + return exported.counters + .filter((row) => row.resource === 'email_receives_per_day') + .reduce((total, row) => total + row.count, 0) } function createFailingEmailBlobs() { @@ -1093,7 +1099,7 @@ test('charged retry repairs after unrelated writes fill storage bytes', async () ).toHaveLength(1) }) -test('ambiguous quota-ledger batch response charges one delivery exactly once', async () => { +test('ambiguous delivery-event insert response charges one delivery exactly once', async () => { silenceIncidentalRuntimeWarnings() await ensureEmailTestSchema(env.APP_DB) const username = `quota-ambiguous-${crypto.randomUUID().slice(0, 8)}` @@ -1105,19 +1111,37 @@ test('ambiguous quota-ledger batch response charges one delivery exactly once', email: accountEmail, username, }) - let batchCalls = 0 + let loseNextDeliveryInsertResponse = true const ambiguousDb = new Proxy(env.APP_DB, { get(target, property, receiver) { - if (property === 'batch') { - return async (statements: Parameters[0]) => { - batchCalls++ - const result = await target.batch(statements) - // The account-write lease acquire is the first batch. Lose the - // response from the following quota-ledger batch. - if (batchCalls === 2) { - throw new Error('simulated quota batch response loss') + if (property === 'prepare') { + return (query: string) => { + const statement = target.prepare(query) + const isDeliveryInsert = + query.includes('INSERT OR IGNORE INTO email_delivery_events') && + query.includes("'receive_started'") + if (!isDeliveryInsert) return statement + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: bound.first.bind(bound), + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: async () => { + const result = await bound.run() + if (loseNextDeliveryInsertResponse) { + loseNextDeliveryInsertResponse = false + throw new Error( + 'simulated delivery event insert response loss', + ) + } + return result + }, + } + }, } - return result } } const value = Reflect.get(target, property, receiver) @@ -1391,13 +1415,7 @@ test('byte-identical mail dedupes only inside the explicit delivery window', asy limit: 10, }), ).toHaveLength(2) - const counter = await env.APP_DB.prepare( - `SELECT SUM(count) AS count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day'`, - ) - .bind(userId) - .first<{ count: number }>() - expect(Number(counter?.count ?? 0)).toBe(2) + expect(await readUserTotalReceiveCount(userId)).toBe(2) expect( await env.APP_DB.prepare( `SELECT event_count FROM usage_rollups diff --git a/packages/worker/src/email/outbound.ts b/packages/worker/src/email/outbound.ts index 64e94ef2d5..62700cf7e9 100644 --- a/packages/worker/src/email/outbound.ts +++ b/packages/worker/src/email/outbound.ts @@ -43,6 +43,7 @@ type SendEmailEnv = Pick< | 'CLOUDFLARE_ACCOUNT_ID' | 'CLOUDFLARE_API_BASE_URL' | 'CLOUDFLARE_API_TOKEN' + | 'USER_METER' > /** @@ -550,11 +551,11 @@ export async function sendOutboundEmail( }) + attachmentBytesTotal, }) - // Atomic check-and-increment: the counter tracks attempts for every user - // and denies the send when a plan's daily limit is reached. The - // `max` plan uses finite email caps rather than uncapped mail. + // Atomic check-and-increment via UserMeter. No ExecutionContext here; + // the legacy D1 mirror uses the service's caught best-effort fallback. await consumeDailyEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: sender.accountEmail, resource: 'email_sends_per_day', diff --git a/packages/worker/src/email/outbound.workers.test.ts b/packages/worker/src/email/outbound.workers.test.ts index 3b6742595c..3f87a19883 100644 --- a/packages/worker/src/email/outbound.workers.test.ts +++ b/packages/worker/src/email/outbound.workers.test.ts @@ -1,4 +1,5 @@ import { env } from 'cloudflare:workers' +import { runInDurableObject } from 'cloudflare:test' import { expect, test } from 'vitest' import { http, HttpResponse } from 'msw' import { bytesToBase64 } from '@kody-internal/shared/base64.ts' @@ -22,7 +23,9 @@ import { silenceIncidentalRuntimeWarnings } from '#worker/test-support/incidenta import { createMswNodeServer } from '#worker/test-support/msw-node-server.ts' import { isEntitlementLimitError } from '#worker/entitlements/errors.ts' import { maxPlanEmailLimits, planLimits } from '#worker/entitlements/plans.ts' -import { incrementDailyEntitlementCounter } from '#worker/entitlements/service.ts' +import { UserMeter } from '#worker/entitlements/user-meter-do.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' +import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' @@ -70,13 +73,35 @@ async function seedVerifiedAccount(input: { } async function readDailyEmailSendCounter(userId: string) { - const row = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = ? AND day = ?`, + const result = await userMeterRpc({ env, userId }).read({ + resource: 'email_sends_per_day', + day: utcDayKey(), + }) + return result.outcome === 'ready' ? result.count : 0 +} + +async function seedDailyEmailSendCounter(userId: string, count: number) { + const day = utcDayKey() + const updatedAt = new Date().toISOString() + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(userId)), ) - .bind(userId, 'email_sends_per_day', utcDayKey()) - .first<{ count: number }>() - return Number(row?.count ?? 0) + await runInDurableObject(stub, async (instance: UserMeter, state) => { + expect(instance).toBeInstanceOf(UserMeter) + await instance.read({ resource: 'email_sends_per_day', day }) + state.storage.sql.exec( + `INSERT INTO daily_counters (resource, day, count, revision, updated_at) + VALUES (?, ?, ?, 1, ?) + ON CONFLICT(resource, day) DO UPDATE SET + count = excluded.count, + revision = excluded.revision, + updated_at = excluded.updated_at`, + 'email_sends_per_day', + day, + count, + updatedAt, + ) + }) } async function sendSelfNotification(input: { @@ -187,7 +212,7 @@ test('sendOutboundEmail sends from the platform-assigned username address to the limit: 5, }) expect(listed.map((message) => message.id)).toContain(result.message.id) -}) +}, 30_000) test('sendOutboundEmail rejects suspended accounts', async () => { silenceIncidentalRuntimeWarnings() @@ -823,7 +848,7 @@ test('sendOutboundEmail rejects invalid and oversized attachments', async () => limit: 5, }), ).toEqual([]) -}) +}, 30_000) test('sendOutboundEmail enforces email_sends_per_day for plan users at the limit', async () => { await ensureEmailTestSchema(env.APP_DB) @@ -832,18 +857,7 @@ test('sendOutboundEmail enforces email_sends_per_day for plan users at the limit const limit = planLimits.pro.maxEmailSendsPerDay if (limit === null) throw new Error('Expected a numeric pro email limit.') await seedVerifiedAccount({ email, plan: 'pro' }) - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?)`, - ) - .bind( - userId, - 'email_sends_per_day', - utcDayKey(), - limit, - new Date().toISOString(), - ) - .run() + await seedDailyEmailSendCounter(userId, limit) const error = await sendSelfNotification({ userId, @@ -863,7 +877,7 @@ test('sendOutboundEmail enforces email_sends_per_day for plan users at the limit current: limit, }) expect(error.message).toContain(`at most ${limit} email sends per day`) -}) +}, 30_000) test('sendOutboundEmail binds plan limits even when the caller context has no account email', async () => { await ensureEmailTestSchema(env.APP_DB) @@ -872,18 +886,7 @@ test('sendOutboundEmail binds plan limits even when the caller context has no ac const limit = planLimits.pro.maxEmailSendsPerDay if (limit === null) throw new Error('Expected a numeric pro email limit.') await seedVerifiedAccount({ email, plan: 'pro' }) - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?)`, - ) - .bind( - userId, - 'email_sends_per_day', - utcDayKey(), - limit, - new Date().toISOString(), - ) - .run() + await seedDailyEmailSendCounter(userId, limit) // Package subscription contexts pass an empty account email; the plan // limit must still apply (no fallback-limit bypass). @@ -906,7 +909,7 @@ test('sendOutboundEmail binds plan limits even when the caller context has no ac plan: 'pro', limit, }) -}) +}, 30_000) test('sendOutboundEmail increments the daily counter when under the plan limit', async () => { await ensureEmailTestSchema(env.APP_DB) @@ -922,7 +925,7 @@ test('sendOutboundEmail increments the daily counter when under the plan limit', expect(result.status).toBe('sent') expect(await readDailyEmailSendCounter(userId)).toBe(1) -}) +}, 30_000) test('sendOutboundEmail caps max-plan users at the email daily backstop', async () => { await ensureEmailTestSchema(env.APP_DB) @@ -930,17 +933,10 @@ test('sendOutboundEmail caps max-plan users at the email daily backstop', async const email = `max-${crypto.randomUUID()}@example.com` const userId = await createStableUserIdFromEmail(email) await seedVerifiedAccount({ email, plan: 'max' }) - for ( - let index = 0; - index < maxPlanEmailLimits.email_sends_per_day - 1; - index += 1 - ) { - await incrementDailyEntitlementCounter({ - db: env.APP_DB, - userId, - resource: 'email_sends_per_day', - }) - } + await seedDailyEmailSendCounter( + userId, + maxPlanEmailLimits.email_sends_per_day - 1, + ) // One send left under the backstop succeeds... const result = await sendSelfNotification({ userId, accountEmail: email }) expect(result.status).toBe('sent') @@ -969,4 +965,4 @@ test('sendOutboundEmail caps max-plan users at the email daily backstop', async expect(await readDailyEmailSendCounter(userId)).toBe( maxPlanEmailLimits.email_sends_per_day, ) -}) +}, 30_000) diff --git a/packages/worker/src/entitlements/entitlements.node.test.ts b/packages/worker/src/entitlements/entitlements.node.test.ts index 2b3585ce6b..568a352fdc 100644 --- a/packages/worker/src/entitlements/entitlements.node.test.ts +++ b/packages/worker/src/entitlements/entitlements.node.test.ts @@ -34,12 +34,17 @@ import { refundDailyEntitlement, } from './service.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { + createInMemoryUserMeterEnv, + createWaitUntilDrain, +} from '#worker/test-support/user-meter.ts' type CounterRow = { user_id: string resource: string day: string count: number + updated_at?: string } function createEntitlementsTestDb( @@ -207,8 +212,12 @@ function createEntitlementsTestDb( return { meta: { changes: 1 } } } if (query.includes('INSERT INTO entitlement_daily_counters')) { + const isAbsoluteMirror = + query.includes('count = excluded.count') && + query.includes('updated_at < excluded.updated_at') const isConditionalConsume = query.includes('count + 1 <= ?') const amount = isConditionalConsume ? 1 : Number(params[3]) + const mirrorUpdatedAt = String(params[4] ?? '') const existing = counters.find( (counter) => counter.user_id === params[0] && @@ -222,13 +231,28 @@ function createEntitlementsTestDb( ) { return { meta: { changes: 0 } } } - existing.count += amount + if (isAbsoluteMirror) { + const existingUpdatedAt = existing.updated_at ?? '' + if ( + existingUpdatedAt !== '' && + !(existingUpdatedAt < mirrorUpdatedAt) + ) { + return { meta: { changes: 0 } } + } + existing.count = amount + existing.updated_at = mirrorUpdatedAt + } else { + existing.count += amount + } } else { counters.push({ user_id: String(params[0]), resource: String(params[1]), day: String(params[2]), count: amount, + ...(isAbsoluteMirror + ? { updated_at: mirrorUpdatedAt } + : {}), }) } return { meta: { changes: 1 } } @@ -752,31 +776,31 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n const { db, counters } = createEntitlementsTestDb({ users: [{ email: plannedEmail, plan: 'free', stable_user_id: userId }], }) + const { env } = createInMemoryUserMeterEnv() + const mirror = createWaitUntilDrain() expect(utcDayKey(now)).toBe('2026-07-05') const limit = planLimits.free.maxEmailSendsPerDay if (limit === null) throw new Error('Expected a numeric email send limit.') for (let index = 0; index < limit; index += 1) { - await assertWithinEntitlement({ + await consumeDailyEntitlement({ db, + env, userId, email: plannedEmail, resource: 'email_sends_per_day', now, - }) - await incrementDailyEntitlementCounter({ - db, - userId, - resource: 'email_sends_per_day', - now, + waitUntil: mirror.waitUntil, }) } + await mirror.drain() expect(counters).toEqual([ { user_id: userId, resource: 'email_sends_per_day', day: '2026-07-05', count: limit, + updated_at: expect.stringMatching(/^r\/0+[0-9]+$/), }, ]) await expect( @@ -791,10 +815,12 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n const denied = await consumeDailyEntitlement({ db, + env, userId, email: plannedEmail, resource: 'email_sends_per_day', now, + waitUntil: mirror.waitUntil, }).then( () => null, (thrown: unknown) => thrown, @@ -812,47 +838,57 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n expect(counters[0]?.count).toBe(limit) const nextDay = new Date('2026-07-06T00:00:01.000Z') - await assertWithinEntitlement({ - db, - userId, - email: plannedEmail, - resource: 'email_sends_per_day', - now: nextDay, - }) await consumeDailyEntitlement({ db, + env, userId, email: plannedEmail, resource: 'email_sends_per_day', now: nextDay, + waitUntil: mirror.waitUntil, }) + await mirror.drain() expect(counters).toHaveLength(2) expect(counters[1]?.count).toBe(1) }) test('refundDailyEntitlement decrements the user/day counter and floors at zero', async () => { const { db, counters } = createEntitlementsTestDb() + const { env } = createInMemoryUserMeterEnv() + const mirror = createWaitUntilDrain() const now = new Date('2026-07-05T15:00:00.000Z') - await incrementDailyEntitlementCounter({ - db, - userId: 'user-1', - resource: 'email_receives_per_day', - amount: 2, - now, - }) - await incrementDailyEntitlementCounter({ - db, - userId: 'user-2', - resource: 'email_receives_per_day', - amount: 3, - now, - }) + for (let index = 0; index < 2; index += 1) { + await consumeDailyEntitlement({ + db, + env, + userId: 'user-1', + email: null, + resource: 'email_receives_per_day', + now, + waitUntil: mirror.waitUntil, + }) + } + for (let index = 0; index < 3; index += 1) { + await consumeDailyEntitlement({ + db, + env, + userId: 'user-2', + email: null, + resource: 'email_receives_per_day', + now, + waitUntil: mirror.waitUntil, + }) + } + await mirror.drain() await refundDailyEntitlement({ db, + env, userId: 'user-1', resource: 'email_receives_per_day', now, + waitUntil: mirror.waitUntil, }) + await mirror.drain() expect( counters.find( (row) => @@ -868,16 +904,21 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' await refundDailyEntitlement({ db, + env, userId: 'user-1', resource: 'email_receives_per_day', now, + waitUntil: mirror.waitUntil, }) await refundDailyEntitlement({ db, + env, userId: 'user-1', resource: 'email_receives_per_day', now, + waitUntil: mirror.waitUntil, }) + await mirror.drain() expect( counters.find( (row) => @@ -888,25 +929,32 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' test('missing-email lookups fail closed and honor free email caps', async () => { const { db, counters } = createEntitlementsTestDb() + const { env } = createInMemoryUserMeterEnv() + const mirror = createWaitUntilDrain() const sendLimit = planLimits.free.maxEmailSendsPerDay const now = new Date('2026-07-05T15:00:00.000Z') for (let index = 0; index < sendLimit; index += 1) { await consumeDailyEntitlement({ db, + env, userId: 'user-1', email: null, resource: 'email_sends_per_day', now, + waitUntil: mirror.waitUntil, }) } + await mirror.drain() expect(counters[0]?.count).toBe(sendLimit) await expect( consumeDailyEntitlement({ db, + env, userId: 'user-1', email: null, resource: 'email_sends_per_day', now, + waitUntil: mirror.waitUntil, }), ).rejects.toBeInstanceOf(EntitlementLimitError) @@ -914,22 +962,27 @@ test('missing-email lookups fail closed and honor free email caps', async () => for (let index = 0; index < receiveLimit; index += 1) { await consumeDailyEntitlement({ db, + env, userId: 'user-1', email: null, resource: 'email_receives_per_day', now, + waitUntil: mirror.waitUntil, }) } + await mirror.drain() expect( counters.find((row) => row.resource === 'email_receives_per_day')?.count, ).toBe(receiveLimit) const denied = await consumeDailyEntitlement({ db, + env, userId: 'user-1', email: null, resource: 'email_receives_per_day', now, + waitUntil: mirror.waitUntil, }).then( () => null, (thrown: unknown) => thrown, diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index c7a7f534a1..3eafe6b86a 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -9,6 +9,11 @@ import { type EntitlementResource, type PlanName, } from './plans.ts' +import { + isDailyEntitlementResource, + type DailyEntitlementResource, +} from './user-meter-do.ts' +import { userMeterRpc, type UserMeterEnv } from './user-meter-client.ts' const stableUserIdPattern = /^[a-f0-9]{64}$/i @@ -193,9 +198,9 @@ export async function findUserAccountByStableUserId( } /** - * Increment a daily counter for a rate-style entitlement (for example - * email sends per day). Counters accumulate for every user regardless of - * plan so that assigning a plan later enforces against real usage. + * Legacy D1 mirror write for a daily entitlement counter. Expand-phase + * enforcement is authoritative in UserMeter; this helper remains for tests, + * backfills, and the best-effort mirror scheduled after DO consume/refund. */ export async function incrementDailyEntitlementCounter(input: { db: D1Database @@ -223,6 +228,110 @@ export async function incrementDailyEntitlementCounter(input: { .run() } +function assertDailyEntitlementResource( + resource: EntitlementResource, +): DailyEntitlementResource { + if (!isDailyEntitlementResource(resource)) { + throw new Error( + `Expected a daily entitlement resource; got ${JSON.stringify(resource)}.`, + ) + } + return resource +} + +// Best-effort mirror scheduling: prefer `waitUntil` when available; otherwise +// catch so the promise cannot surface as unhandled. Never awaited by callers. +function scheduleDailyEntitlementMirror( + work: Promise, + waitUntil?: (promise: Promise) => void, +) { + const tracked = work.catch((error: unknown) => { + console.warn('entitlement-daily-mirror-failed', error) + }) + if (waitUntil) { + waitUntil(tracked) + return + } + void tracked +} + +// Absolute D1 mirror ordered by DO-minted `mirrorUpdatedAt` (revision-primary) +// so a late older write cannot overwrite newer state, including refunds. +async function mirrorDailyEntitlementAbsoluteCount(input: { + db: D1Database + userId: string + resource: DailyEntitlementResource + day: string + count: number + mirrorUpdatedAt: string +}) { + await input.db + .prepare( + `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(user_id, resource, day) DO UPDATE SET + count = excluded.count, + updated_at = excluded.updated_at + WHERE entitlement_daily_counters.updated_at < excluded.updated_at`, + ) + .bind( + input.userId, + input.resource, + input.day, + input.count, + input.mirrorUpdatedAt, + ) + .run() +} + +/** Best-effort absolute D1 mirror of UserMeter daily counter state. */ +export function scheduleAbsoluteDailyEntitlementMirror(input: { + db: D1Database + userId: string + resource: EntitlementResource + day: string + count: number + mirrorUpdatedAt: string + waitUntil?: (promise: Promise) => void +}): void { + const resource = assertDailyEntitlementResource(input.resource) + scheduleDailyEntitlementMirror( + mirrorDailyEntitlementAbsoluteCount({ + db: input.db, + userId: input.userId, + resource, + day: input.day, + count: input.count, + mirrorUpdatedAt: input.mirrorUpdatedAt, + }), + input.waitUntil, + ) +} + +async function ensureUserMeterCounterInitialized(input: { + db: D1Database + env: UserMeterEnv + userId: string + resource: DailyEntitlementResource + day: string + updatedAt: string + now: Date +}) { + const baseline = await readDailyEntitlementCounter({ + db: input.db, + userId: input.userId, + resource: input.resource, + now: input.now, + }) + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + await meter.initialize({ + resource: input.resource, + day: input.day, + count: baseline, + updatedAt: input.updatedAt, + }) +} + async function readDailyEntitlementCounter(input: { db: D1Database userId: string @@ -239,6 +348,52 @@ async function readDailyEntitlementCounter(input: { return Number(row?.count ?? 0) } +/** + * Point-read one daily entitlement counter from UserMeter. Cold meters + * bootstrap once from the legacy D1 point row, then re-read; warm meters + * return the DO count without touching D1. + */ +export async function readDailyEntitlementResourceUsage(input: { + db: D1Database + env: UserMeterEnv + userId: string + resource: EntitlementResource + now?: Date +}): Promise { + const resource = assertDailyEntitlementResource(input.resource) + const now = input.now ?? new Date() + const day = utcDayKey(now) + const updatedAt = now.toISOString() + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + let result = await meter.read({ + resource, + day, + now: updatedAt, + }) + if (result.outcome === 'needs_bootstrap') { + await ensureUserMeterCounterInitialized({ + db: input.db, + env: input.env, + userId: input.userId, + resource, + day, + updatedAt, + now, + }) + result = await meter.read({ + resource, + day, + now: updatedAt, + }) + if (result.outcome === 'needs_bootstrap') { + throw new Error( + 'UserMeter daily entitlement read still needs bootstrap after initialize.', + ) + } + } + return result.count +} + async function countRows(db: D1Database, sql: string, params: Array) { const row = await db .prepare(sql) @@ -730,6 +885,34 @@ export async function readEntitlementResourceUsage(input: { } } +/** + * Authoritative current usage for any entitlement resource: daily counters + * via UserMeter, everything else via the legacy D1 helpers. + */ +export async function readCurrentEntitlementResourceUsage(input: { + db: D1Database + env: UserMeterEnv + userId: string + resource: EntitlementResource + now: Date +}): Promise { + if (isDailyEntitlementResource(input.resource)) { + return await readDailyEntitlementResourceUsage({ + db: input.db, + env: input.env, + userId: input.userId, + resource: input.resource, + now: input.now, + }) + } + return await readEntitlementResourceUsage({ + db: input.db, + userId: input.userId, + resource: input.resource, + now: input.now, + }) +} + export async function assertWithinStorageBytesEntitlement(input: { db: D1Database userId: string @@ -841,95 +1024,128 @@ export async function assertWithinEntitlement( } } -/** - * Atomically consume one unit of a daily rate-style entitlement (check and - * increment in a single conditional D1 upsert), throwing - * EntitlementLimitError when the consumption would exceed the plan limit. - * This avoids the check-then-increment race that separate - * assertWithinEntitlement + incrementDailyEntitlementCounter calls would - * have under concurrent requests, and evaluates the UTC day key once. - * Every resolved plan has finite numeric limits. - */ -export async function consumeDailyEntitlement(input: { +export type ConsumeDailyEntitlementInput = { db: D1Database + /** Must expose `USER_METER` (authoritative); D1 is a best-effort mirror. */ + env: UserMeterEnv userId: string email: string | null | undefined resource: EntitlementResource now?: Date -}): Promise { + /** Prefer `ctx.waitUntil` for the legacy D1 mirror; never awaited here. */ + waitUntil?: (promise: Promise) => void +} + +/** + * Atomically consume one daily entitlement unit via UserMeter, throwing + * EntitlementLimitError when the plan limit would be exceeded. Cold keys + * bootstrap once from legacy D1; warm path awaits only the DO RPC. Mirror + * writes are best-effort and cannot affect enforcement. + */ +export async function consumeDailyEntitlement( + input: ConsumeDailyEntitlementInput, +): Promise { + const resource = assertDailyEntitlementResource(input.resource) const now = input.now ?? new Date() + const day = utcDayKey(now) + const updatedAt = now.toISOString() // Cached plan resolution: this runs on every execute call and every - // sandbox outbound fetch, and the atomic counter upsert below is the only - // part that must see fresh state. + // sandbox outbound fetch; the UserMeter consume below is the only + // counter state that must be fresh. const plan = await getCachedUserPlan(input.db, { userId: input.userId, email: input.email, }) - const limit = resolvePlanLimit(plan, input.resource) - const throwLimitError = async () => { - throw new EntitlementLimitError({ - resource: input.resource, - plan, + const limit = resolvePlanLimit(plan, resource) + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + let result = await meter.consume({ + resource, + day, + limit, + updatedAt, + }) + if (result.outcome === 'needs_bootstrap') { + await ensureUserMeterCounterInitialized({ + db: input.db, + env: input.env, + userId: input.userId, + resource, + day, + updatedAt, + now, + }) + result = await meter.consume({ + resource, + day, limit, - current: await readDailyEntitlementCounter({ - db: input.db, - userId: input.userId, - resource: input.resource, - now, - }), - upgradeHint: buildEntitlementUpgradeHint(input.resource), + updatedAt, }) + if (result.outcome === 'needs_bootstrap') { + throw new Error( + 'UserMeter consume still needs bootstrap after initialize.', + ) + } } - // The fresh-row INSERT branch is unconditional, so a limit below one - // unit can never be satisfied and must be rejected up front. - if (limit < 1) { - await throwLimitError() - } - const result = await input.db - .prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, 1, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = entitlement_daily_counters.count + 1, - updated_at = excluded.updated_at - WHERE entitlement_daily_counters.count + 1 <= ?`, - ) - .bind( - input.userId, - input.resource, - utcDayKey(now), - now.toISOString(), + if (!result.consumed) { + throw new EntitlementLimitError({ + resource, + plan, limit, - ) - .run() - if ((result.meta.changes ?? 0) === 0) { - await throwLimitError() + current: result.count, + upgradeHint: buildEntitlementUpgradeHint(resource), + }) } + scheduleDailyEntitlementMirror( + mirrorDailyEntitlementAbsoluteCount({ + db: input.db, + userId: input.userId, + resource, + day, + count: result.count, + mirrorUpdatedAt: result.mirrorUpdatedAt, + }), + input.waitUntil, + ) } -/** - * Atomically refund one previously consumed daily entitlement unit for the - * given user/resource/UTC day (floors at zero). Scoped by `user_id` so a - * refund can never touch another user's counter. Used when inbound receive - * quota was charged before a retryable storage failure — the same `now` - * (day key) as the matching `consumeDailyEntitlement` call must be passed. - */ -export async function refundDailyEntitlement(input: { +export type RefundDailyEntitlementInput = { db: D1Database + env: UserMeterEnv userId: string resource: EntitlementResource now?: Date -}): Promise { + waitUntil?: (promise: Promise) => void +} + +/** + * Atomically refund one previously consumed daily entitlement unit in + * UserMeter (floors at zero). Pass the same `now` (day key) as the matching + * consume. The legacy D1 mirror is best-effort and never awaited. + */ +export async function refundDailyEntitlement( + input: RefundDailyEntitlementInput, +): Promise { + const resource = assertDailyEntitlementResource(input.resource) const now = input.now ?? new Date() - await input.db - .prepare( - `UPDATE entitlement_daily_counters - SET count = MAX(0, count - 1), - updated_at = ? - WHERE user_id = ? - AND resource = ? - AND day = ?`, - ) - .bind(now.toISOString(), input.userId, input.resource, utcDayKey(now)) - .run() + const day = utcDayKey(now) + const updatedAt = now.toISOString() + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + const result = await meter.refund({ + resource, + day, + updatedAt, + }) + // revision 0 means the key was never initialized; nothing to mirror. + if (result.revision < 1) return + scheduleDailyEntitlementMirror( + mirrorDailyEntitlementAbsoluteCount({ + db: input.db, + userId: input.userId, + resource, + day, + count: result.count, + mirrorUpdatedAt: result.mirrorUpdatedAt, + }), + input.waitUntil, + ) } diff --git a/packages/worker/src/entitlements/user-meter-client.ts b/packages/worker/src/entitlements/user-meter-client.ts new file mode 100644 index 0000000000..e8e27a2019 --- /dev/null +++ b/packages/worker/src/entitlements/user-meter-client.ts @@ -0,0 +1,28 @@ +import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' +import { type UserMeterRpc } from './user-meter-do.ts' + +export type UserMeterEnv = { + USER_METER?: DurableObjectNamespace +} + +export function userMeterNamespace( + env: UserMeterEnv, +): DurableObjectNamespace | null { + return env.USER_METER ?? null +} + +/** Typed per-user UserMeter RPC stub; throws when `USER_METER` is missing. */ +export function userMeterRpc(input: { + env: UserMeterEnv + userId: string +}): UserMeterRpc { + const namespace = userMeterNamespace(input.env) + if (!namespace) { + throw new Error('USER_METER Durable Object binding is not configured.') + } + return namespace.get( + namespace.idFromName(userMeterDurableObjectName(input.userId)), + ) as unknown as UserMeterRpc +} + +export type { UserMeterRpc } diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts new file mode 100644 index 0000000000..36ac9c5213 --- /dev/null +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -0,0 +1,678 @@ +import * as Sentry from '@sentry/cloudflare' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { DurableObject } from 'cloudflare:workers' +import { buildSentryOptions } from '#worker/sentry-options.ts' +import { type EntitlementResource } from './plans.ts' + +/** Daily rate-style resources stored in the per-user UserMeter (UTC day keys). */ +export const dailyEntitlementResources = [ + 'email_sends_per_day', + 'email_receives_per_day', + 'execute_calls_per_day', + 'outbound_fetches_per_day', +] as const satisfies ReadonlyArray + +export type DailyEntitlementResource = + (typeof dailyEntitlementResources)[number] + +export function isDailyEntitlementResource( + resource: string, +): resource is DailyEntitlementResource { + return (dailyEntitlementResources as ReadonlyArray).includes(resource) +} + +/** Retention window for UserMeter daily rows (enforcement needs today only). */ +export const userMeterDailyCounterRetentionDays = 7 + +const metaSchemaVersionKey = 'schema_version' +/** Bump when initializeSchema DDL changes; warm objects skip DDL. */ +const userMeterSchemaVersion = 3 + +const defaultExportPageSize = 100 +const maxExportPageSize = 500 +const maxInboundDeliveryIdLength = 256 +const inboundReceiveResource = + 'email_receives_per_day' satisfies DailyEntitlementResource +const utcDayKeyPattern = /^\d{4}-\d{2}-\d{2}$/ + +/** + * Legacy D1 mirror `updated_at` token. Lexicographic order matches revision + * order; `r/` sorts after pre-cutover ISO timestamps. + */ +export function userMeterMirrorUpdatedAtToken(revision: number): string { + const safeRevision = + Number.isSafeInteger(revision) && revision > 0 ? revision : 0 + return `r/${String(safeRevision).padStart(20, '0')}` +} + +export type UserMeterCounterRow = { + resource: DailyEntitlementResource + day: string + count: number + revision: number + updatedAt: string + mirrorUpdatedAt: string +} + +export type UserMeterReadyState = { + outcome: 'ready' + count: number + revision: number + mirrorUpdatedAt: string +} + +export type UserMeterBootstrapState = { + outcome: 'needs_bootstrap' +} + +export type UserMeterConsumeResult = + | UserMeterBootstrapState + | (UserMeterReadyState & { consumed: boolean }) + +export type UserMeterReadResult = UserMeterBootstrapState | UserMeterReadyState + +export type UserMeterRefundResult = UserMeterReadyState + +export type UserMeterInitializeResult = UserMeterReadyState & { + created: boolean +} + +export type UserMeterExportResult = { + counters: Array + nextStartAfter: string | null + truncated: boolean +} + +/** + * Inbound delivery claim + receive consume. Retries set `replayed` without + * incrementing; `day`/`resource` come from the original claim on cross-day + * retries. + */ +export type UserMeterInboundDeliveryConsumeResult = + | UserMeterBootstrapState + | (UserMeterReadyState & { + consumed: boolean + replayed: boolean + day: string + resource: typeof inboundReceiveResource + }) + +type ExportCursor = { + day: string + resource: string +} + +function assertDailyResource(resource: string): DailyEntitlementResource { + if (!isDailyEntitlementResource(resource)) { + throw new Error( + `UserMeter resource must be a daily entitlement resource; got ${JSON.stringify(resource)}.`, + ) + } + return resource +} + +function assertInboundReceiveResource( + resource: string, +): typeof inboundReceiveResource { + const daily = assertDailyResource(resource) + if (daily !== inboundReceiveResource) { + throw new Error( + `UserMeter inbound delivery consume requires ${inboundReceiveResource}; got ${JSON.stringify(resource)}.`, + ) + } + return daily +} + +function assertUtcDayKey(day: string): string { + if (!utcDayKeyPattern.test(day)) { + throw new Error( + `UserMeter day must be a UTC YYYY-MM-DD key; got ${JSON.stringify(day)}.`, + ) + } + return day +} + +function assertInboundDeliveryId(deliveryId: string): string { + if ( + typeof deliveryId !== 'string' || + deliveryId.length === 0 || + deliveryId.length > maxInboundDeliveryIdLength + ) { + throw new Error( + `UserMeter inbound deliveryId must be a non-empty string up to ${maxInboundDeliveryIdLength} characters.`, + ) + } + return deliveryId +} + +function retentionCutoffDay(now: Date): string { + const cutoff = new Date(now) + cutoff.setUTCDate( + cutoff.getUTCDate() - (userMeterDailyCounterRetentionDays - 1), + ) + return utcDayKey(cutoff) +} + +function encodeExportCursor(cursor: ExportCursor) { + return JSON.stringify([cursor.day, cursor.resource]) +} + +function decodeExportCursor(startAfter: string): ExportCursor | null { + try { + const parsed = JSON.parse(startAfter) as unknown + if (!Array.isArray(parsed) || parsed.length !== 2) return null + const [day, resource] = parsed + if (typeof day !== 'string' || typeof resource !== 'string') return null + return { day, resource } + } catch { + return null + } +} + +function normalizePageSize(pageSize: number | undefined) { + const requested = + typeof pageSize === 'number' && Number.isFinite(pageSize) + ? Math.trunc(pageSize) + : defaultExportPageSize + return Math.min(Math.max(requested, 1), maxExportPageSize) +} + +function readyState(count: number, revision: number): UserMeterReadyState { + const safeCount = Math.max(0, count) + const safeRevision = Math.max(0, revision) + return { + outcome: 'ready', + count: safeCount, + revision: safeRevision, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(safeRevision), + } +} + +class UserMeterBase extends DurableObject { + constructor(ctx: DurableObjectState, env: Env) { + super(ctx, env) + this.ctx.blockConcurrencyWhile(async () => { + this.initializeSchema() + }) + } + + private initializeSchema() { + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS user_meter_meta ( + key TEXT PRIMARY KEY NOT NULL, + value INTEGER NOT NULL + ) + `) + const versionRow = this.ctx.storage.sql + .exec<{ + value: number + }>( + `SELECT value FROM user_meter_meta WHERE key = ? LIMIT 1`, + metaSchemaVersionKey, + ) + .toArray()[0] + const version = versionRow == null ? null : Number(versionRow.value) || 0 + if (version === userMeterSchemaVersion) return + + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS daily_counters ( + resource TEXT NOT NULL, + day TEXT NOT NULL, + count INTEGER NOT NULL, + revision INTEGER NOT NULL, + updated_at TEXT NOT NULL, + PRIMARY KEY (resource, day) + ) + `) + if (version === 1) { + try { + this.ctx.storage.sql.exec( + `ALTER TABLE daily_counters ADD COLUMN revision INTEGER NOT NULL DEFAULT 0`, + ) + } catch { + // Column already present on a partially migrated object. + } + } + this.ctx.storage.sql.exec( + `CREATE INDEX IF NOT EXISTS idx_daily_counters_day + ON daily_counters (day)`, + ) + // Per-user delivery-id ledger (PK is delivery_id alone); same retention + // window as daily counters so Email Routing retries cannot double-charge. + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS inbound_delivery_claims ( + delivery_id TEXT PRIMARY KEY NOT NULL, + resource TEXT NOT NULL, + day TEXT NOT NULL, + count_after INTEGER NOT NULL, + revision INTEGER NOT NULL, + claimed_at TEXT NOT NULL + ) + `) + this.ctx.storage.sql.exec( + `CREATE INDEX IF NOT EXISTS idx_inbound_delivery_claims_day + ON inbound_delivery_claims (day)`, + ) + this.ctx.storage.sql.exec( + `INSERT INTO user_meter_meta (key, value) VALUES (?, ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value`, + metaSchemaVersionKey, + userMeterSchemaVersion, + ) + } + + private deleteStaleCounters(now: Date) { + const cutoffDay = retentionCutoffDay(now) + this.ctx.storage.sql.exec( + `DELETE FROM daily_counters WHERE day < ?`, + cutoffDay, + ) + this.ctx.storage.sql.exec( + `DELETE FROM inbound_delivery_claims WHERE day < ?`, + cutoffDay, + ) + } + + private readInboundDeliveryClaim(deliveryId: string): { + resource: string + day: string + countAfter: number + revision: number + claimedAt: string + } | null { + const row = this.ctx.storage.sql + .exec<{ + resource: string + day: string + count_after: number + revision: number + claimed_at: string + }>( + `SELECT resource, day, count_after, revision, claimed_at + FROM inbound_delivery_claims + WHERE delivery_id = ?`, + deliveryId, + ) + .toArray()[0] + if (!row) return null + return { + resource: String(row.resource), + day: String(row.day), + countAfter: Math.max(0, Number(row.count_after ?? 0)), + revision: Math.max(0, Number(row.revision ?? 0)), + claimedAt: String(row.claimed_at ?? ''), + } + } + + private readRow( + resource: DailyEntitlementResource, + day: string, + ): { count: number; revision: number; updatedAt: string } | null { + const row = this.ctx.storage.sql + .exec<{ count: number; revision: number; updated_at: string }>( + `SELECT count, revision, updated_at + FROM daily_counters + WHERE resource = ? AND day = ?`, + resource, + day, + ) + .toArray()[0] + if (!row) return null + return { + count: Math.max(0, Number(row.count ?? 0)), + revision: Math.max(0, Number(row.revision ?? 0)), + updatedAt: String(row.updated_at ?? ''), + } + } + + /** Cold-key seed from legacy D1; INSERT OR IGNORE is concurrency-safe. */ + async initialize(input: { + resource: string + day: string + count: number + updatedAt: string + }): Promise { + const resource = assertDailyResource(input.resource) + const day = assertUtcDayKey(input.day) + const now = new Date(input.updatedAt) + this.deleteStaleCounters(Number.isNaN(now.valueOf()) ? new Date() : now) + const count = Math.max(0, Math.trunc(Number(input.count) || 0)) + const cursor = this.ctx.storage.sql.exec( + `INSERT INTO daily_counters (resource, day, count, revision, updated_at) + VALUES (?, ?, ?, 1, ?) + ON CONFLICT(resource, day) DO NOTHING`, + resource, + day, + count, + input.updatedAt, + ) + const created = cursor.rowsWritten > 0 + const row = this.readRow(resource, day) + if (!row) { + throw new Error( + 'UserMeter initialize failed to materialize a counter row.', + ) + } + return { ...readyState(row.count, row.revision), created } + } + + /** Check-and-increment one unit; missing keys return `needs_bootstrap`. */ + async consume(input: { + resource: string + day: string + limit: number + updatedAt: string + }): Promise { + const resource = assertDailyResource(input.resource) + const day = assertUtcDayKey(input.day) + const now = new Date(input.updatedAt) + this.deleteStaleCounters(Number.isNaN(now.valueOf()) ? new Date() : now) + + const existing = this.readRow(resource, day) + if (!existing) { + return { outcome: 'needs_bootstrap' } + } + + if (input.limit < 1 || existing.count + 1 > input.limit) { + return { + ...readyState(existing.count, existing.revision), + consumed: false, + } + } + + const nextCount = existing.count + 1 + const nextRevision = existing.revision + 1 + this.ctx.storage.sql.exec( + `UPDATE daily_counters + SET count = ?, + revision = ?, + updated_at = ? + WHERE resource = ? AND day = ? AND revision = ?`, + nextCount, + nextRevision, + input.updatedAt, + resource, + day, + existing.revision, + ) + const row = this.readRow(resource, day) + if (!row) { + throw new Error('UserMeter consume lost the counter row.') + } + return { + ...readyState(row.count, row.revision), + consumed: row.count === nextCount && row.revision === nextRevision, + } + } + + async read(input: { + resource: string + day: string + now?: string + }): Promise { + const resource = assertDailyResource(input.resource) + const day = assertUtcDayKey(input.day) + const now = input.now ? new Date(input.now) : new Date() + this.deleteStaleCounters(Number.isNaN(now.valueOf()) ? new Date() : now) + const row = this.readRow(resource, day) + if (!row) return { outcome: 'needs_bootstrap' } + return readyState(row.count, row.revision) + } + + /** + * Claim `deliveryId` and consume one receive unit. Increment + claim insert + * run in `storage.transactionSync`. + */ + async consumeInboundDelivery(input: { + deliveryId: string + resource: string + day: string + limit: number + updatedAt: string + }): Promise { + const resource = assertInboundReceiveResource(input.resource) + const day = assertUtcDayKey(input.day) + const deliveryId = assertInboundDeliveryId(input.deliveryId) + const now = new Date(input.updatedAt) + this.deleteStaleCounters(Number.isNaN(now.valueOf()) ? new Date() : now) + + const priorClaim = this.readInboundDeliveryClaim(deliveryId) + if (priorClaim) { + if (priorClaim.resource !== resource) { + throw new Error( + `UserMeter inbound deliveryId was claimed for ${JSON.stringify(priorClaim.resource)}; cannot reuse for ${JSON.stringify(resource)}.`, + ) + } + const claimedResource = assertInboundReceiveResource(priorClaim.resource) + const claimedDay = assertUtcDayKey(priorClaim.day) + const row = this.readRow(claimedResource, claimedDay) + return { + ...(row + ? readyState(row.count, row.revision) + : readyState(priorClaim.countAfter, priorClaim.revision)), + consumed: false, + replayed: true, + day: claimedDay, + resource: claimedResource, + } + } + + const existing = this.readRow(resource, day) + if (!existing) { + return { outcome: 'needs_bootstrap' } + } + + if (input.limit < 1 || existing.count + 1 > input.limit) { + return { + ...readyState(existing.count, existing.revision), + consumed: false, + replayed: false, + day, + resource, + } + } + + const nextCount = existing.count + 1 + const nextRevision = existing.revision + 1 + this.ctx.storage.transactionSync(() => { + this.ctx.storage.sql.exec( + `UPDATE daily_counters + SET count = ?, + revision = ?, + updated_at = ? + WHERE resource = ? AND day = ? AND revision = ?`, + nextCount, + nextRevision, + input.updatedAt, + resource, + day, + existing.revision, + ) + this.ctx.storage.sql.exec( + `INSERT INTO inbound_delivery_claims ( + delivery_id, resource, day, count_after, revision, claimed_at + ) VALUES (?, ?, ?, ?, ?, ?)`, + deliveryId, + resource, + day, + nextCount, + nextRevision, + input.updatedAt, + ) + }) + const row = this.readRow(resource, day) + if (!row) { + throw new Error( + 'UserMeter inbound delivery consume lost the counter row.', + ) + } + return { + ...readyState(row.count, row.revision), + consumed: row.count === nextCount && row.revision === nextRevision, + replayed: false, + day, + resource, + } + } + + /** Decrement one unit (floors at zero); missing keys stay uninitialized. */ + async refund(input: { + resource: string + day: string + updatedAt: string + }): Promise { + const resource = assertDailyResource(input.resource) + const day = assertUtcDayKey(input.day) + const now = new Date(input.updatedAt) + this.deleteStaleCounters(Number.isNaN(now.valueOf()) ? new Date() : now) + + const existing = this.readRow(resource, day) + if (!existing) { + return readyState(0, 0) + } + const nextCount = Math.max(0, existing.count - 1) + const nextRevision = existing.revision + 1 + this.ctx.storage.sql.exec( + `UPDATE daily_counters + SET count = ?, + revision = ?, + updated_at = ? + WHERE resource = ? AND day = ? AND revision = ?`, + nextCount, + nextRevision, + input.updatedAt, + resource, + day, + existing.revision, + ) + const row = this.readRow(resource, day) + if (!row) { + throw new Error('UserMeter refund lost the counter row.') + } + return readyState(row.count, row.revision) + } + + async purge(): Promise<{ ok: true }> { + await this.ctx.storage.deleteAll() + this.initializeSchema() + return { ok: true } + } + + async exportCounters(input: { + pageSize?: number + startAfter?: string | null + }): Promise { + this.deleteStaleCounters(new Date()) + const pageSize = normalizePageSize(input.pageSize) + const cursor = + typeof input.startAfter === 'string' && input.startAfter.length > 0 + ? decodeExportCursor(input.startAfter) + : null + + const rows = ( + cursor + ? this.ctx.storage.sql.exec<{ + resource: string + day: string + count: number + revision: number + updated_at: string + }>( + `SELECT resource, day, count, revision, updated_at + FROM daily_counters + WHERE day > ? + OR (day = ? AND resource > ?) + ORDER BY day ASC, resource ASC + LIMIT ?`, + cursor.day, + cursor.day, + cursor.resource, + pageSize + 1, + ) + : this.ctx.storage.sql.exec<{ + resource: string + day: string + count: number + revision: number + updated_at: string + }>( + `SELECT resource, day, count, revision, updated_at + FROM daily_counters + ORDER BY day ASC, resource ASC + LIMIT ?`, + pageSize + 1, + ) + ).toArray() + + const truncated = rows.length > pageSize + const pageRows = truncated ? rows.slice(0, pageSize) : rows + const counters: Array = [] + for (const row of pageRows) { + const resource = String(row.resource) + if (!isDailyEntitlementResource(resource)) continue + const revision = Math.max(0, Number(row.revision ?? 0)) + counters.push({ + resource, + day: String(row.day), + count: Math.max(0, Number(row.count ?? 0)), + revision, + updatedAt: String(row.updated_at), + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(revision), + }) + } + const last = pageRows[pageRows.length - 1] + return { + counters, + nextStartAfter: + truncated && last + ? encodeExportCursor({ + day: String(last.day), + resource: String(last.resource), + }) + : null, + truncated, + } + } +} + +export const UserMeter = Sentry.instrumentDurableObjectWithSentry( + (env: Env) => buildSentryOptions(env), + UserMeterBase, +) + +export type UserMeterRpc = { + initialize: (input: { + resource: string + day: string + count: number + updatedAt: string + }) => Promise + consume: (input: { + resource: string + day: string + limit: number + updatedAt: string + }) => Promise + consumeInboundDelivery: (input: { + deliveryId: string + resource: string + day: string + limit: number + updatedAt: string + }) => Promise + read: (input: { + resource: string + day: string + now?: string + }) => Promise + refund: (input: { + resource: string + day: string + updatedAt: string + }) => Promise + purge: () => Promise<{ ok: true }> + exportCounters: (input: { + pageSize?: number + startAfter?: string | null + }) => Promise +} diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts new file mode 100644 index 0000000000..4bfec33bfc --- /dev/null +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -0,0 +1,616 @@ +import { runInDurableObject } from 'cloudflare:test' +import { env } from 'cloudflare:workers' +import { expect, test } from 'vitest' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { consoleWarn } from '#worker/test-support/console-spies.ts' +import { seedAccount } from '#worker/test-support/workers-seed.ts' +import { createStableUserIdFromEmail } from '#worker/user-id.ts' +import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' +import { EntitlementLimitError } from './errors.ts' +import { planLimits } from './plans.ts' +import { consumeDailyEntitlement, refundDailyEntitlement } from './service.ts' +import { ensureEntitlementTestSchema } from './test-schema.ts' +import { userMeterRpc } from './user-meter-client.ts' +import { UserMeter, userMeterMirrorUpdatedAtToken } from './user-meter-do.ts' +import { + createWaitUntilDrain, + withPatchedDbPrepare, +} from '#worker/test-support/user-meter.ts' + +async function seedFreeUser(emailPrefix: string) { + await ensureEntitlementTestSchema(env.APP_DB) + const email = `${emailPrefix}-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(email) + await seedAccount({ + db: env.APP_DB, + email, + username: `meter-${crypto.randomUUID().slice(0, 8)}`, + plan: 'free', + stableUserId: userId, + }) + return { email, userId } +} + +async function readD1DailyCount(input: { + userId: string + resource: string + day: string +}) { + const row = await env.APP_DB.prepare( + `SELECT count, updated_at FROM entitlement_daily_counters + WHERE user_id = ? AND resource = ? AND day = ?`, + ) + .bind(input.userId, input.resource, input.day) + .first<{ count: number; updated_at: string }>() + return row + ? { count: Number(row.count), updatedAt: String(row.updated_at) } + : null +} + +async function waitFor( + predicate: () => boolean, + timeoutMs = 5_000, + label = 'condition', +) { + const started = Date.now() + while (!predicate()) { + if (Date.now() - started > timeoutMs) { + throw new Error(`Timed out waiting for ${label}.`) + } + await new Promise((resolve) => setTimeout(resolve, 5)) + } +} + +const mirrorSql = `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(user_id, resource, day) DO UPDATE SET + count = excluded.count, + updated_at = excluded.updated_at + WHERE entitlement_daily_counters.updated_at < excluded.updated_at` + +test('UserMeter bootstraps from legacy D1 once, then warms without APP_DB awaits', async () => { + const now = new Date('2026-07-31T15:00:00.000Z') + const day = utcDayKey(now) + const sendLimit = planLimits.free.maxEmailSendsPerDay + const user = await seedFreeUser('meter-bootstrap') + const drain = createWaitUntilDrain() + const meter = userMeterRpc({ env, userId: user.userId }) + + const legacyCount = sendLimit - 1 + await env.APP_DB.prepare( + `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) + VALUES (?, ?, ?, ?, ?)`, + ) + .bind( + user.userId, + 'email_sends_per_day', + day, + legacyCount, + '2026-07-31T12:00:00.000Z', + ) + .run() + + let d1PointReads = 0 + using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + const statement = originalPrepare(query) + const isPointRead = + query.includes('SELECT count FROM entitlement_daily_counters') && + query.includes('WHERE user_id = ? AND resource = ? AND day = ?') + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: async () => { + if (isPointRead) d1PointReads += 1 + return await bound.first() + }, + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: bound.run.bind(bound), + } + }, + } + }) as D1Database['prepare'] + }) + + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + expect(d1PointReads).toBe(1) + expect( + await meter.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ + outcome: 'ready', + count: sendLimit, + }) + + const denied = await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }).then( + () => null, + (thrown: unknown) => thrown, + ) + expect(denied).toBeInstanceOf(EntitlementLimitError) + expect(denied).toMatchObject({ + details: { + resource: 'email_sends_per_day', + plan: 'free', + limit: sendLimit, + current: sendLimit, + }, + }) + expect(d1PointReads).toBe(1) + + await refundDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + await drain.drain() + d1PointReads = 0 + let nonMirrorAppDbRuns = 0 + let releaseMirror!: () => void + const mirrorGate = new Promise((resolve) => { + releaseMirror = resolve + }) + let mirrorRunStarted = false + using _warmPatch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + const statement = originalPrepare(query) + const isMirrorWrite = + query.includes('INSERT INTO entitlement_daily_counters') && + query.includes('updated_at < excluded.updated_at') + const isPointRead = + query.includes('SELECT count FROM entitlement_daily_counters') && + query.includes('WHERE user_id = ? AND resource = ? AND day = ?') + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: async () => { + if (isPointRead) d1PointReads += 1 + return await bound.first() + }, + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: async () => { + if (isMirrorWrite) { + mirrorRunStarted = true + await mirrorGate + return await bound.run() + } + nonMirrorAppDbRuns += 1 + return await bound.run() + }, + } + }, + } + }) as D1Database['prepare'] + }) + try { + await expect( + consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }), + ).resolves.toBeUndefined() + expect(mirrorRunStarted).toBe(true) + expect(d1PointReads).toBe(0) + expect(nonMirrorAppDbRuns).toBe(0) + } finally { + releaseMirror() + } + await drain.drain() +}, 30_000) + +test('UserMeter mirror tokens reject out-of-order absolute D1 writes', async () => { + const now = new Date('2026-07-31T16:00:00.000Z') + const day = utcDayKey(now) + const user = await seedFreeUser('meter-mirror-order') + const meter = userMeterRpc({ env, userId: user.userId }) + + await meter.initialize({ + resource: 'email_sends_per_day', + day, + count: 0, + updatedAt: now.toISOString(), + }) + const first = await meter.consume({ + resource: 'email_sends_per_day', + day, + limit: 100, + updatedAt: now.toISOString(), + }) + expect(first).toMatchObject({ + outcome: 'ready', + consumed: true, + count: 1, + revision: 2, + }) + const second = await meter.consume({ + resource: 'email_sends_per_day', + day, + limit: 100, + updatedAt: now.toISOString(), + }) + expect(second).toMatchObject({ + outcome: 'ready', + consumed: true, + count: 2, + revision: 3, + }) + const refunded = await meter.refund({ + resource: 'email_sends_per_day', + day, + updatedAt: now.toISOString(), + }) + expect(refunded).toMatchObject({ + outcome: 'ready', + count: 1, + revision: 4, + }) + if (first.outcome !== 'ready') { + throw new Error('Expected ready consume results.') + } + + type PendingMirror = { + count: number + mirrorUpdatedAt: string + resolve: () => void + promise: Promise + } + const pending: Array = [] + function deferMirror(count: number, mirrorUpdatedAt: string) { + let resolve!: () => void + const promise = new Promise((res) => { + resolve = res + }) + pending.push({ count, mirrorUpdatedAt, resolve, promise }) + return promise + } + + using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + const statement = originalPrepare(query) + const isMirrorWrite = + query.includes('INSERT INTO entitlement_daily_counters') && + query.includes('updated_at < excluded.updated_at') + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: bound.first.bind(bound), + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: async () => { + if (!isMirrorWrite) return await bound.run() + await deferMirror(Number(params[3]), String(params[4])) + return await bound.run() + }, + } + }, + } + }) as D1Database['prepare'] + }) + + const older = env.APP_DB.prepare(mirrorSql) + .bind( + user.userId, + 'email_sends_per_day', + day, + first.count, + first.mirrorUpdatedAt, + ) + .run() + const newer = env.APP_DB.prepare(mirrorSql) + .bind( + user.userId, + 'email_sends_per_day', + day, + refunded.count, + refunded.mirrorUpdatedAt, + ) + .run() + + await waitFor(() => pending.length === 2, 5_000, 'mirror gates') + expect(pending.map((entry) => entry.mirrorUpdatedAt)).toEqual([ + userMeterMirrorUpdatedAtToken(2), + userMeterMirrorUpdatedAtToken(4), + ]) + const [olderGate, newerGate] = pending + newerGate!.resolve() + await newer + expect( + await readD1DailyCount({ + userId: user.userId, + resource: 'email_sends_per_day', + day, + }), + ).toEqual({ + count: 1, + updatedAt: userMeterMirrorUpdatedAtToken(4), + }) + + olderGate!.resolve() + await older + expect( + await readD1DailyCount({ + userId: user.userId, + resource: 'email_sends_per_day', + day, + }), + ).toEqual({ + count: 1, + updatedAt: userMeterMirrorUpdatedAtToken(4), + }) + expect( + await meter.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 1, revision: 4 }) +}, 30_000) + +test('UserMeter daily entitlement consume/refund/read/export/purge workflow is per-user and mirror-safe', async () => { + const now = new Date('2026-07-31T15:00:00.000Z') + const day = utcDayKey(now) + const sendLimit = planLimits.free.maxEmailSendsPerDay + const userA = await seedFreeUser('meter-a') + const userB = await seedFreeUser('meter-b') + const drain = createWaitUntilDrain() + const meterA = userMeterRpc({ env, userId: userA.userId }) + const meterB = userMeterRpc({ env, userId: userB.userId }) + + expect(userMeterDurableObjectName(userA.userId)).toBe(userA.userId) + + await meterA.initialize({ + resource: 'email_sends_per_day', + day, + count: 0, + updatedAt: now.toISOString(), + }) + await expect( + meterA.consume({ + resource: 'email_sends_per_day', + day, + limit: sendLimit, + updatedAt: now.toISOString(), + }), + ).resolves.toMatchObject({ + outcome: 'ready', + consumed: true, + count: 1, + }) + + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + email: userA.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + await drain.drain() + expect( + await meterA.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 2 }) + + for (let index = 2; index < sendLimit; index += 1) { + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + email: userA.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + } + const concurrent = await Promise.all( + Array.from({ length: 8 }, async () => + consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + email: userA.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }).then( + () => null, + (thrown: unknown) => thrown, + ), + ), + ) + expect(concurrent.filter((result) => result === null)).toHaveLength(0) + const denials = concurrent.filter( + (result) => result instanceof EntitlementLimitError, + ) + expect(denials).toHaveLength(8) + for (const denial of denials) { + expect(denial).toMatchObject({ + details: { + code: 'entitlement_limit_exceeded', + resource: 'email_sends_per_day', + plan: 'free', + limit: sendLimit, + current: sendLimit, + }, + }) + } + + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userB.userId, + email: userB.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + expect( + await meterB.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 1 }) + expect( + await meterA.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: sendLimit }) + + await refundDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + expect( + await meterA.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: sendLimit - 1 }) + for (let index = 0; index < sendLimit; index += 1) { + await refundDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + } + expect( + await meterA.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 0 }) + + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + email: userA.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }) + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: userA.userId, + email: userA.email, + resource: 'execute_calls_per_day', + now, + waitUntil: drain.waitUntil, + }) + const exportedA = await meterA.exportCounters({}) + expect(exportedA.counters).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + resource: 'email_sends_per_day', + day, + count: 1, + }), + expect.objectContaining({ + resource: 'execute_calls_per_day', + day, + count: 1, + }), + ]), + ) + + await meterA.purge() + expect(await meterA.read({ resource: 'email_sends_per_day', day })).toEqual({ + outcome: 'needs_bootstrap', + }) + expect(await meterA.exportCounters({})).toEqual({ + counters: [], + nextStartAfter: null, + truncated: false, + }) + expect( + await meterB.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 1 }) + + await env.APP_DB.prepare( + `DELETE FROM entitlement_daily_counters + WHERE user_id = ? AND resource = ? AND day = ?`, + ) + .bind(userA.userId, 'email_sends_per_day', day) + .run() + + consoleWarn.mockImplementation(() => {}) + const failingMirrorDb = { + prepare(query: string) { + if ( + query.includes('INSERT INTO entitlement_daily_counters') && + query.includes('updated_at < excluded.updated_at') + ) { + return { + bind() { + return { + async run() { + throw new Error('mirror write failed') + }, + } + }, + } + } + return env.APP_DB.prepare(query) + }, + } as unknown as D1Database + + await expect( + consumeDailyEntitlement({ + db: failingMirrorDb, + env, + userId: userA.userId, + email: userA.email, + resource: 'email_sends_per_day', + now, + waitUntil: drain.waitUntil, + }), + ).resolves.toBeUndefined() + await drain.drain() + expect( + await meterA.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 1 }) + expect(consoleWarn).toHaveBeenCalledWith( + 'entitlement-daily-mirror-failed', + expect.any(Error), + ) +}, 30_000) + +test('UserMeter consume rejects invalid UTC day keys before SQL', async () => { + const user = await seedFreeUser('meter-invalid-day') + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(user.userId)), + ) + await runInDurableObject(stub, async (instance: UserMeter) => { + await expect( + instance.consume({ + resource: 'email_sends_per_day', + day: 'not-a-day', + limit: 1, + updatedAt: '2026-07-31T15:00:00.000Z', + }), + ).rejects.toThrow(/UTC YYYY-MM-DD/) + }) +}, 30_000) diff --git a/packages/worker/src/index.ts b/packages/worker/src/index.ts index 33f7fe0be5..0d50120601 100644 --- a/packages/worker/src/index.ts +++ b/packages/worker/src/index.ts @@ -6,6 +6,7 @@ import { MCP } from './mcp/index.ts' import { JobManager } from './jobs/manager-do.ts' import { StorageRunner } from './storage-runner.ts' import { RunLog } from './run-records/run-log-do.ts' +import { UserMeter } from './entitlements/user-meter-do.ts' import { RepoSession } from './repo/repo-session-do.ts' import { PackageRealtimeSession } from '#worker/package-runtime/realtime-session.ts' import { PackageServiceInstance } from '#worker/package-runtime/package-service.ts' @@ -108,6 +109,7 @@ export { PackageAppRuntimeBridge, StorageRunner, RunLog, + UserMeter, OAuthPurgeCoordinator, } diff --git a/packages/worker/src/mcp/capabilities/email/email-usage-get.ts b/packages/worker/src/mcp/capabilities/email/email-usage-get.ts index 5e3e6436ab..e1db22b96d 100644 --- a/packages/worker/src/mcp/capabilities/email/email-usage-get.ts +++ b/packages/worker/src/mcp/capabilities/email/email-usage-get.ts @@ -8,6 +8,7 @@ import { } from '#worker/entitlements/plans.ts' import { getUserPlan, + readDailyEntitlementResourceUsage, readEntitlementResourceUsage, } from '#worker/entitlements/service.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' @@ -54,14 +55,16 @@ export const emailUsageGetCapability = defineDomainCapability( resource: 'stored_email_messages', now, }), - readEntitlementResourceUsage({ + readDailyEntitlementResourceUsage({ db, + env: ctx.env, userId: user.userId, resource: 'email_sends_per_day', now, }), - readEntitlementResourceUsage({ + readDailyEntitlementResourceUsage({ db, + env: ctx.env, userId: user.userId, resource: 'email_receives_per_day', now, diff --git a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts index 2a59e9abb5..2b71fd5734 100644 --- a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts +++ b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts @@ -2,6 +2,7 @@ import { env } from 'cloudflare:workers' import { expect, test } from 'vitest' import { createMcpCallerContext } from '#mcp/context.ts' import { maxPlanEmailLimits, planLimits } from '#worker/entitlements/plans.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import { ensureEmailTestSchema } from '#worker/email/test-schema.ts' import { seedAccount } from '#worker/test-support/workers-seed.ts' @@ -62,103 +63,157 @@ function buildCallerContext(user: { userId: string; email: string } | null) { }) } -test('email_usage_get returns usage for verified users and enforces auth requirements', async () => { - await ensureEmailTestSchema(env.APP_DB) - await expect( - emailUsageGetCapability.handler( - {}, - { env, callerContext: buildCallerContext(null) }, - ), - ).rejects.toThrow(/Authenticated MCP user/) +test( + 'email_usage_get enforces auth and reads UserMeter-backed daily counts', + { timeout: 30_000 }, + async () => { + await ensureEmailTestSchema(env.APP_DB) + await expect( + emailUsageGetCapability.handler( + {}, + { env, callerContext: buildCallerContext(null) }, + ), + ).rejects.toThrow(/Authenticated MCP user/) - const unverifiedEmail = `usage-unverified-${crypto.randomUUID()}@example.com` - const unverifiedUserId = await createStableUserIdFromEmail(unverifiedEmail) - await seedUsageAccount({ - email: unverifiedEmail, - plan: 'max', - emailVerifiedAt: null, - }) - await expect( - emailUsageGetCapability.handler( + const unverifiedEmail = `usage-unverified-${crypto.randomUUID()}@example.com` + const unverifiedUserId = await createStableUserIdFromEmail(unverifiedEmail) + await seedUsageAccount({ + email: unverifiedEmail, + plan: 'max', + emailVerifiedAt: null, + }) + await expect( + emailUsageGetCapability.handler( + {}, + { + env, + callerContext: buildCallerContext({ + userId: unverifiedUserId, + email: unverifiedEmail, + }), + }, + ), + ).rejects.toThrow(/Account email is not verified/) + + const planEmail = `usage-plan-${crypto.randomUUID()}@example.com` + const planUserId = await createStableUserIdFromEmail(planEmail) + await seedUsageAccount({ email: planEmail, plan: 'pro' }) + await seedDailyCounter({ + userId: planUserId, + resource: 'email_sends_per_day', + count: 3, + }) + await seedDailyCounter({ + userId: planUserId, + resource: 'email_receives_per_day', + count: 5, + }) + await seedStoredMessages(planUserId, 2) + + const planResult = await emailUsageGetCapability.handler( {}, { env, callerContext: buildCallerContext({ - userId: unverifiedUserId, - email: unverifiedEmail, + userId: planUserId, + email: planEmail, }), }, - ), - ).rejects.toThrow(/Account email is not verified/) + ) + expect(planResult).toEqual({ + plan: 'pro', + day: utcDayKey(), + stored_messages: { + count: 2, + limit: planLimits.pro.maxStoredEmailMessages, + }, + sends_today: { count: 3, limit: planLimits.pro.maxEmailSendsPerDay }, + receives_today: { + count: 5, + limit: planLimits.pro.maxEmailReceivesPerDay, + }, + max_message_bytes: planLimits.pro.maxEmailMessageBytes, + }) + const day = utcDayKey() + const planMeter = userMeterRpc({ env, userId: planUserId }) + expect( + await planMeter.read({ resource: 'email_sends_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 3 }) + expect( + await planMeter.read({ resource: 'email_receives_per_day', day }), + ).toMatchObject({ outcome: 'ready', count: 5 }) - const planEmail = `usage-plan-${crypto.randomUUID()}@example.com` - const planUserId = await createStableUserIdFromEmail(planEmail) - await seedUsageAccount({ email: planEmail, plan: 'pro' }) - await seedDailyCounter({ - userId: planUserId, - resource: 'email_sends_per_day', - count: 3, - }) - await seedDailyCounter({ - userId: planUserId, - resource: 'email_receives_per_day', - count: 5, - }) - await seedStoredMessages(planUserId, 2) + const maxEmail = `usage-max-${crypto.randomUUID()}@example.com` + const maxUserId = await createStableUserIdFromEmail(maxEmail) + await seedUsageAccount({ email: maxEmail, plan: 'max' }) + const maxResult = await emailUsageGetCapability.handler( + {}, + { + env, + callerContext: buildCallerContext({ + userId: maxUserId, + email: maxEmail, + }), + }, + ) + expect(maxResult).toEqual({ + plan: 'max', + day: utcDayKey(), + stored_messages: { + count: 0, + limit: maxPlanEmailLimits.stored_email_messages, + }, + sends_today: { + count: 0, + limit: maxPlanEmailLimits.email_sends_per_day, + }, + receives_today: { + count: 0, + limit: maxPlanEmailLimits.email_receives_per_day, + }, + max_message_bytes: maxPlanEmailLimits.email_message_bytes, + }) - const planResult = await emailUsageGetCapability.handler( - {}, - { - env, - callerContext: buildCallerContext({ - userId: planUserId, - email: planEmail, - }), - }, - ) - expect(planResult).toEqual({ - plan: 'pro', - day: utcDayKey(), - stored_messages: { - count: 2, - limit: planLimits.pro.maxStoredEmailMessages, - }, - sends_today: { count: 3, limit: planLimits.pro.maxEmailSendsPerDay }, - receives_today: { + const meterEmail = `usage-meter-${crypto.randomUUID()}@example.com` + const meterUserId = await createStableUserIdFromEmail(meterEmail) + await seedUsageAccount({ email: meterEmail, plan: 'pro' }) + await seedDailyCounter({ + userId: meterUserId, + resource: 'email_sends_per_day', + count: 3, + }) + await seedDailyCounter({ + userId: meterUserId, + resource: 'email_receives_per_day', count: 5, - limit: planLimits.pro.maxEmailReceivesPerDay, - }, - max_message_bytes: planLimits.pro.maxEmailMessageBytes, - }) - - const maxEmail = `usage-max-${crypto.randomUUID()}@example.com` - const maxUserId = await createStableUserIdFromEmail(maxEmail) - await seedUsageAccount({ email: maxEmail, plan: 'max' }) - const maxResult = await emailUsageGetCapability.handler( - {}, - { - env, - callerContext: buildCallerContext({ - userId: maxUserId, - email: maxEmail, - }), - }, - ) - expect(maxResult).toEqual({ - plan: 'max', - day: utcDayKey(), - stored_messages: { - count: 0, - limit: maxPlanEmailLimits.stored_email_messages, - }, - sends_today: { - count: 0, - limit: maxPlanEmailLimits.email_sends_per_day, - }, - receives_today: { - count: 0, - limit: maxPlanEmailLimits.email_receives_per_day, - }, - max_message_bytes: maxPlanEmailLimits.email_message_bytes, - }) -}) + }) + await seedStoredMessages(meterUserId, 2) + const updatedAt = new Date().toISOString() + const meter = userMeterRpc({ env, userId: meterUserId }) + await meter.initialize({ + resource: 'email_sends_per_day', + day, + count: 13, + updatedAt, + }) + await meter.initialize({ + resource: 'email_receives_per_day', + day, + count: 15, + updatedAt, + }) + const meterResult = await emailUsageGetCapability.handler( + {}, + { + env, + callerContext: buildCallerContext({ + userId: meterUserId, + email: meterEmail, + }), + }, + ) + expect(meterResult.sends_today.count).toBe(13) + expect(meterResult.receives_today.count).toBe(15) + expect(meterResult.stored_messages.count).toBe(2) + }, +) diff --git a/packages/worker/src/mcp/fetch-gateway.node.test.ts b/packages/worker/src/mcp/fetch-gateway.node.test.ts index 66cbdf1b10..189235e61e 100644 --- a/packages/worker/src/mcp/fetch-gateway.node.test.ts +++ b/packages/worker/src/mcp/fetch-gateway.node.test.ts @@ -12,15 +12,10 @@ import { import { buildBasicAuthSecretPlaceholder } from '#mcp/secrets/placeholders.ts' import * as secretService from '#mcp/secrets/service.ts' import * as communityRepo from '#worker/community/repo.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import * as packageRepo from '#worker/package-registry/repo.ts' -/** - * Minimal D1 stub for the entitlement reads/writes executeGatewayFetch - * performs (account reverse-resolution `first()` and the conditional - * counter upsert `run()`); secret resolution itself is spied at the - * service layer. Both statements must bind the acting userId — the stub - * fails loudly when a query drops the per-user scoping. - */ +const userMeter = createInMemoryUserMeterEnv() const env = { APP_DB: { prepare(query: string) { @@ -57,9 +52,10 @@ const env = { } }, } as unknown as D1Database, + ...userMeter.env, COOKIE_SECRET: 'test-cookie-secret', SECRET_STORE_KEY: 'test-secret-store-key-32-chars-minimum', -} +} as unknown as Env const props = { baseUrl: 'https://example.com', @@ -626,7 +622,8 @@ test('gateway fetch records outbound_fetch usage metering', async () => { }) expect(successResponse.status).toBe(200) expect(fetchStub).toHaveBeenCalledTimes(1) - expect(waitUntil).toHaveBeenCalledTimes(1) + // waitUntil: entitlement D1 mirror + usage metering + expect(waitUntil).toHaveBeenCalledTimes(2) expect(recordUsageSpy).toHaveBeenCalledTimes(1) expect(recordUsageSpy).toHaveBeenCalledWith(env, { userId: 'user-123', @@ -750,7 +747,8 @@ test('gateway fetch records outbound_fetch usage metering', async () => { waitUntil, }), ).rejects.toThrow('network failed with waitUntil') - expect(waitUntil).toHaveBeenCalledTimes(1) + // waitUntil: entitlement D1 mirror (before fetch) + usage metering (error path) + expect(waitUntil).toHaveBeenCalledTimes(2) expect(recordUsageSpy).toHaveBeenCalledTimes(1) expect(recordUsageSpy.mock.calls[0]?.[1]).toMatchObject({ entityId: 'api.example.com', diff --git a/packages/worker/src/mcp/fetch-gateway.ts b/packages/worker/src/mcp/fetch-gateway.ts index c2450f366f..c76041cd22 100644 --- a/packages/worker/src/mcp/fetch-gateway.ts +++ b/packages/worker/src/mcp/fetch-gateway.ts @@ -87,7 +87,7 @@ function applyOutboundFetchTimeout( } export async function executeGatewayFetch(input: { - env: Pick & UsageEnv + env: Pick & UsageEnv props: FetchGatewayProps request: Request globalFetch?: typeof fetch @@ -131,9 +131,11 @@ export async function executeGatewayFetch(input: { null await consumeDailyEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.props.userId, email, resource: 'outbound_fetches_per_day', + waitUntil: input.waitUntil, }) } const transformed = await expandSecretPlaceholders({ diff --git a/packages/worker/src/mcp/fetch-gateway.workers.test.ts b/packages/worker/src/mcp/fetch-gateway.workers.test.ts index 70a8678c0d..9fa02cee28 100644 --- a/packages/worker/src/mcp/fetch-gateway.workers.test.ts +++ b/packages/worker/src/mcp/fetch-gateway.workers.test.ts @@ -1,15 +1,42 @@ import { env } from 'cloudflare:workers' +import { runInDurableObject } from 'cloudflare:test' import { expect, test } from 'vitest' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import { executeGatewayFetch } from '#mcp/fetch-gateway.ts' import { ensureEntitlementTestSchema } from '#worker/entitlements/test-schema.ts' import { isEntitlementLimitError } from '#worker/entitlements/errors.ts' import { planLimits } from '#worker/entitlements/plans.ts' -import { incrementDailyEntitlementCounter } from '#worker/entitlements/service.ts' +import { UserMeter } from '#worker/entitlements/user-meter-do.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' +import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' import { silenceIncidentalRuntimeWarnings } from '#worker/test-support/incidental-runtime-warnings.ts' import { seedAccount } from '#worker/test-support/workers-seed.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' +async function seedOutboundFetchCounter(userId: string, count: number) { + const day = utcDayKey() + const updatedAt = new Date().toISOString() + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(userId)), + ) + await runInDurableObject(stub, async (instance: UserMeter, state) => { + expect(instance).toBeInstanceOf(UserMeter) + await instance.read({ resource: 'outbound_fetches_per_day', day }) + state.storage.sql.exec( + `INSERT INTO daily_counters (resource, day, count, revision, updated_at) + VALUES (?, ?, ?, 1, ?) + ON CONFLICT(resource, day) DO UPDATE SET + count = excluded.count, + revision = excluded.revision, + updated_at = excluded.updated_at`, + 'outbound_fetches_per_day', + day, + count, + updatedAt, + ) + }) +} + test('gateway fetches consume the daily outbound-fetch entitlement and deny over the plan limit', async () => { silenceIncidentalRuntimeWarnings() await ensureEntitlementTestSchema(env.APP_DB) @@ -30,34 +57,38 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over } const globalFetch = (async () => new Response('ok')) as unknown as typeof fetch + const waitUntilTasks: Array> = [] + const waitUntil = (promise: Promise) => { + waitUntilTasks.push(promise) + } const allowed = await executeGatewayFetch({ env, props, request: new Request('https://api.example.net/data'), globalFetch, + waitUntil, }) expect(await allowed.text()).toBe('ok') - const counter = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'outbound_fetches_per_day' AND day = ?`, - ) - .bind(userId, utcDayKey()) - .first<{ count: number }>() - expect(Number(counter?.count)).toBe(1) + await Promise.all(waitUntilTasks) + expect( + await userMeterRpc({ env, userId }).read({ + resource: 'outbound_fetches_per_day', + day: utcDayKey(), + }), + ).toMatchObject({ outcome: 'ready', count: 1 }) - // Fill the remaining free-plan quota, then expect denial. - await incrementDailyEntitlementCounter({ - db: env.APP_DB, + // Fill the remaining free-plan quota in UserMeter, then expect denial. + await seedOutboundFetchCounter( userId, - resource: 'outbound_fetches_per_day', - amount: planLimits.free.maxOutboundFetchesPerDay - 1, - }) + planLimits.free.maxOutboundFetchesPerDay, + ) const denied = await executeGatewayFetch({ env, props, request: new Request('https://api.example.net/data'), globalFetch, + waitUntil, }).catch((error: unknown) => error) expect(isEntitlementLimitError(denied)).toBe(true) @@ -70,6 +101,7 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over props: { ...props, email: null }, request: new Request('https://api.example.net/data'), globalFetch, + waitUntil, }).catch((error: unknown) => error) expect(isEntitlementLimitError(deniedWithoutEmail)).toBe(true) @@ -80,6 +112,7 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over props: { ...props, userId: null, email: null }, request: new Request('https://api.example.net/data'), globalFetch, + waitUntil, }) expect(await contextless.text()).toBe('ok') -}) +}, 30_000) diff --git a/packages/worker/src/mcp/tools/execute.ts b/packages/worker/src/mcp/tools/execute.ts index 6105080192..22e16ae25c 100644 --- a/packages/worker/src/mcp/tools/execute.ts +++ b/packages/worker/src/mcp/tools/execute.ts @@ -262,9 +262,11 @@ export async function registerExecuteTool(agent: McpRegistrationAgent) { if (callerContext.user?.userId) { await consumeDailyEntitlement({ db: env.APP_DB, + env, userId: callerContext.user.userId, email: callerContext.user.email, resource: 'execute_calls_per_day', + waitUntil: agent.waitUntil?.bind(agent), }) } diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts new file mode 100644 index 0000000000..bf3a390f23 --- /dev/null +++ b/packages/worker/src/test-support/user-meter.ts @@ -0,0 +1,182 @@ +import { + isDailyEntitlementResource, + userMeterMirrorUpdatedAtToken, + type DailyEntitlementResource, +} from '#worker/entitlements/user-meter-do.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' + +type MeterRow = { count: number; revision: number } + +/** + * In-memory UserMeter stub keyed by `idFromName` (stable userId) for node-unit + * coverage of expand-phase service/usage paths. Workers suites exercise the + * real Durable Object binding instead. + */ +export function createInMemoryUserMeterEnv() { + const metersByUser = new Map>() + + function counterKey(resource: string, day: string) { + return `${resource}\0${day}` + } + + function meterFor(userId: string) { + const existingRows = metersByUser.get(userId) + const rows = existingRows ?? new Map() + if (!existingRows) metersByUser.set(userId, rows) + + function readRow(resource: string, day: string) { + return rows.get(counterKey(resource, day)) ?? null + } + + function ready(row: MeterRow) { + return { + outcome: 'ready' as const, + count: row.count, + revision: row.revision, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), + } + } + + return { + async initialize(input: { + resource: string + day: string + count: number + updatedAt: string + }) { + if (!isDailyEntitlementResource(input.resource)) { + throw new Error(`Invalid daily resource: ${input.resource}`) + } + const key = counterKey(input.resource, input.day) + const existing = rows.get(key) + if (existing) return { ...ready(existing), created: false } + const row = { + count: Math.max(0, Math.trunc(input.count) || 0), + revision: 1, + } + rows.set(key, row) + return { ...ready(row), created: true } + }, + async consume(input: { + resource: string + day: string + limit: number + updatedAt: string + }) { + if (!isDailyEntitlementResource(input.resource)) { + throw new Error(`Invalid daily resource: ${input.resource}`) + } + const resource: DailyEntitlementResource = input.resource + const existing = readRow(resource, input.day) + if (!existing) return { outcome: 'needs_bootstrap' as const } + if (input.limit < 1 || existing.count + 1 > input.limit) { + return { ...ready(existing), consumed: false } + } + const next = { + count: existing.count + 1, + revision: existing.revision + 1, + } + rows.set(counterKey(resource, input.day), next) + return { ...ready(next), consumed: true } + }, + async read(input: { resource: string; day: string }) { + const existing = readRow(input.resource, input.day) + if (!existing) return { outcome: 'needs_bootstrap' as const } + return ready(existing) + }, + async refund(input: { + resource: string + day: string + updatedAt: string + }) { + const existing = readRow(input.resource, input.day) + if (!existing) { + return { + outcome: 'ready' as const, + count: 0, + revision: 0, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(0), + } + } + const next = { + count: Math.max(0, existing.count - 1), + revision: existing.revision + 1, + } + rows.set(counterKey(input.resource, input.day), next) + return ready(next) + }, + async purge() { + rows.clear() + return { ok: true as const } + }, + async exportCounters() { + const counters = [...rows.entries()].map(([entryKey, row]) => { + const [resource, day] = entryKey.split('\0') + return { + resource: resource as DailyEntitlementResource, + day: day!, + count: row.count, + revision: row.revision, + updatedAt: new Date().toISOString(), + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), + } + }) + return { counters, nextStartAfter: null, truncated: false } + }, + } + } + + const env = { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: (id: { name: string }) => meterFor(id.name), + }, + } as unknown as UserMeterEnv + + return { + env, + metersByUser, + async seed(input: { + userId: string + resource: DailyEntitlementResource + day: string + count: number + }) { + await meterFor(input.userId).initialize({ + resource: input.resource, + day: input.day, + count: input.count, + updatedAt: new Date().toISOString(), + }) + }, + } +} + +export function createWaitUntilDrain() { + const tasks: Array> = [] + return { + waitUntil(promise: Promise) { + tasks.push(promise) + }, + async drain() { + await Promise.all(tasks) + tasks.length = 0 + }, + } +} + +/** + * Patch `db.prepare` and restore it via `using` even when the body throws. + */ +export function withPatchedDbPrepare( + db: D1Database, + patch: (originalPrepare: D1Database['prepare']) => D1Database['prepare'], +) { + const originalPrepare = db.prepare.bind(db) + db.prepare = patch(originalPrepare) + return { + [Symbol.dispose]() { + db.prepare = originalPrepare + }, + } +} diff --git a/packages/worker/src/user-scoped-durable-object-name.node.test.ts b/packages/worker/src/user-scoped-durable-object-name.node.test.ts index 12525e9749..c21190f1f2 100644 --- a/packages/worker/src/user-scoped-durable-object-name.node.test.ts +++ b/packages/worker/src/user-scoped-durable-object-name.node.test.ts @@ -9,7 +9,9 @@ import { packageRealtimeSessionDurableObjectName, packageServiceInstanceDurableObjectName, repoSessionDurableObjectName, + runLogDurableObjectName, storageRunnerDurableObjectName, + userMeterDurableObjectName, } from './user-scoped-durable-object-name.ts' const identityModuleRelativePath = @@ -22,6 +24,13 @@ test('user-scoped Durable Object name helpers preserve frozen idFromName contrac // JobManager historically does not trim; keep that wire format frozen. expect(jobManagerDurableObjectName(' user-aaa ')).toBe(' user-aaa ') + expect(runLogDurableObjectName('user-aaa')).toBe('user-aaa') + expect(runLogDurableObjectName(' user-aaa ')).toBe(' user-aaa ') + + expect(userMeterDurableObjectName('user-aaa')).toBe('user-aaa') + // UserMeter matches RunLog: untrimmed userId is the frozen idFromName. + expect(userMeterDurableObjectName(' user-aaa ')).toBe(' user-aaa ') + expect(mcpClientHubDurableObjectName(' user-aaa ')).toBe('user-aaa') expect(storageRunnerDurableObjectName('user-aaa', 'job:1')).toBe( diff --git a/packages/worker/src/user-scoped-durable-object-name.ts b/packages/worker/src/user-scoped-durable-object-name.ts index 69383ea8eb..6fa7e52a20 100644 --- a/packages/worker/src/user-scoped-durable-object-name.ts +++ b/packages/worker/src/user-scoped-durable-object-name.ts @@ -30,6 +30,11 @@ export function runLogDurableObjectName(userId: string) { return userId } +/** UserMeter — one daily-entitlement meter DO per user (untrimmed, like RunLog). */ +export function userMeterDurableObjectName(userId: string) { + return userId +} + /** McpClientHub — one client-hub DO per user (trimmed). */ export function mcpClientHubDurableObjectName(userId: string) { return userId.trim() diff --git a/packages/worker/worker-configuration.d.ts b/packages/worker/worker-configuration.d.ts index f1818a4947..4b6cb5eaae 100644 --- a/packages/worker/worker-configuration.d.ts +++ b/packages/worker/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types --config=packages/worker/wrangler.jsonc --env=production ./packages/worker/worker-configuration.d.ts` (hash: c94d594932ed070991b8178999e5f795) +// Generated by Wrangler by running `wrangler types --config=packages/worker/wrangler.jsonc --env=production ./packages/worker/worker-configuration.d.ts` (hash: b7b4d993601fc627f26550ebd7901b3a) // Runtime types generated with workerd@1.20260701.1 2026-04-16 global_fetch_strictly_public,nodejs_compat interface __BaseEnv_Env { OAUTH_KV: KVNamespace; @@ -13,11 +13,13 @@ interface __BaseEnv_Env { FLAG_EXPOSURES: AnalyticsEngineDataset; PLATFORM_FEEDBACK_DISPATCH_QUEUE: Queue; COMMUNITY_ACTIVITY_DISPATCH_QUEUE: Queue; + AUTH_RATE_LIMITER: RateLimit; LOADER: WorkerLoader; APP_LOADER: WorkerLoader; AI: Ai; ASSETS: Fetcher; SENTRY_ENVIRONMENT: "production"; + SIGNUP_MODE: "invite"; SENTRY_TRACES_SAMPLE_RATE: 0; ARTIFACTS_NAMESPACE: "production"; PACKAGE_APP_BASE_URL: "https://kodyapps.dev"; @@ -42,6 +44,7 @@ interface __BaseEnv_Env { JOB_MANAGER: DurableObjectNamespace; STORAGE_RUNNER: DurableObjectNamespace; RUN_LOG: DurableObjectNamespace; + USER_METER: DurableObjectNamespace; REPO_SESSION: DurableObjectNamespace; PACKAGE_REALTIME_SESSION: DurableObjectNamespace; PACKAGE_SERVICE_INSTANCE: DurableObjectNamespace; @@ -50,7 +53,7 @@ interface __BaseEnv_Env { declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/index"); - durableNamespaces: "MCP" | "JobManager" | "StorageRunner" | "RepoSession" | "PackageRealtimeSession" | "PackageServiceInstance" | "RemoteConnectorSession" | "McpClientHub" | "OAuthPurgeCoordinator" | "RunLog"; + durableNamespaces: "MCP" | "JobManager" | "StorageRunner" | "RepoSession" | "PackageRealtimeSession" | "PackageServiceInstance" | "RemoteConnectorSession" | "McpClientHub" | "OAuthPurgeCoordinator" | "RunLog" | "UserMeter"; } interface Env extends __BaseEnv_Env {} } @@ -59,7 +62,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } declare module "*.md" { const value: string; diff --git a/packages/worker/wrangler.jsonc b/packages/worker/wrangler.jsonc index c3d1959989..36bc9f852c 100644 --- a/packages/worker/wrangler.jsonc +++ b/packages/worker/wrangler.jsonc @@ -91,6 +91,10 @@ "new_sqlite_classes": ["RunLog"], "tag": "v20", }, + { + "new_sqlite_classes": ["UserMeter"], + "tag": "v21", + }, ], "observability": { @@ -193,6 +197,10 @@ "class_name": "RunLog", "name": "RUN_LOG", }, + { + "class_name": "UserMeter", + "name": "USER_METER", + }, { "class_name": "RepoSession", "name": "REPO_SESSION", @@ -407,6 +415,10 @@ "class_name": "RunLog", "name": "RUN_LOG", }, + { + "class_name": "UserMeter", + "name": "USER_METER", + }, { "class_name": "RepoSession", "name": "REPO_SESSION", @@ -538,6 +550,10 @@ "class_name": "RunLog", "name": "RUN_LOG", }, + { + "class_name": "UserMeter", + "name": "USER_METER", + }, { "class_name": "RepoSession", "name": "REPO_SESSION", From 8ebf663dff1ef9f0da4006e4d1487c492bf3be96 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:11:39 +0000 Subject: [PATCH 04/39] feat(account): inventory UserMeter lifecycle Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 94 ++++++++++-- .../contributing/architecture/entitlements.md | 120 ++++++++++++--- .../contributing/architecture/primitives.yaml | 14 ++ .../worker/src/account/export.node.test.ts | 124 +++++++++++++++ packages/worker/src/account/export.ts | 145 ++++++++++++++---- .../account/user-owned-surfaces.node.test.ts | 14 ++ .../worker/src/account/user-owned-surfaces.ts | 9 ++ .../src/app/account-deletion.node.test.ts | 41 +++++ packages/worker/src/app/account-deletion.ts | 30 ++++ 9 files changed, 530 insertions(+), 61 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 8b2b37a6cd..353aa888ca 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -78,12 +78,13 @@ Deletion must cover these user-owned surfaces: references a stale column. - **Durable Objects:** `JobManager`, `StorageRunner`, `RepoSession`, `RemoteConnectorSession`, `PackageRealtimeSession`, `PackageServiceInstance`, - `McpClientHub`, and `RunLog` are purged through account-deletion RPCs after - their D1 identifiers are collected (`RunLog` is one object per user and needs - no D1 id scan). `MCP` objects remain SDK session-keyed, while - `mcp_agent_sessions` indexes each Durable Object id by authenticated stable - user id so account deletion can purge stored props, conversation state, - raw-fetch state, and transport storage before revoking OAuth grants. + `McpClientHub`, `RunLog`, and `UserMeter` are purged through account-deletion + RPCs after their D1 identifiers are collected (`RunLog` and `UserMeter` are + one object per user and need no D1 id scan). `MCP` objects remain SDK + session-keyed, while `mcp_agent_sessions` indexes each Durable Object id by + authenticated stable user id so account deletion can purge stored props, + conversation state, raw-fetch state, and transport storage before revoking + OAuth grants. - **Vectorize:** memory, job, and saved-package vector ids are derived from D1 rows and removed with `deleteByIds`. - **R2:** raw email MIME blobs in `EMAIL_BLOBS` are enumerated from @@ -164,13 +165,14 @@ migration-safe chunked interface: with `section: "d1_table"` and a table name. Durable storage buckets are read with `section: "storage_runner"` and a `storage_id`, using the same StorageRunner `exportStorage({ pageSize, startAfter })` RPC as the dedicated - storage export capability. R2 raw MIME, attachment, avatar, and icon objects - use `section: "r2_object"`; each response contains at most one 256 KiB base64 - chunk and an opaque cursor. Each request uses bounded `LIMIT 1` ownership - queries rather than reconstructing inventory. Continuation cursors bind the - source row, object key, size, and ETag; ownership/key mutations and object - overwrites are reported instead of mixing generations. Missing objects are - represented explicitly. + storage export capability. User meter counters use `section: "user_meter"` and + the `UserMeter.exportCounters` RPC. R2 raw MIME, attachment, avatar, and icon + objects use `section: "r2_object"`; each response contains at most one 256 KiB + base64 chunk and an opaque cursor. Each request uses bounded `LIMIT 1` + ownership queries rather than reconstructing inventory. Continuation cursors + bind the source row, object key, size, and ETag; ownership/key mutations and + object overwrites are reported instead of mixing generations. Missing objects + are represented explicitly. D1 manifest counts use bounded SQL `COUNT(*)` queries. D1 section rows are read with SQL-level keyset pagination: every query orders by the table's `rowid`, @@ -189,6 +191,11 @@ Durable Object export behavior: section (`exportRuns` RPC; one cursor pages runs first, then ledger rows). Retention is self-enforced inside the DO (~30 days / 2,000 runs; ledger terminal rows 90 days); see [Run records](./run-records.md). +- `UserMeter` exports daily entitlement counter rows through the `user_meter` + section (`exportCounters` RPC; keyset pagination by UTC `day` and `resource`). + Retention is self-enforced inside the DO (seven UTC days of counter and + inbound-delivery-claim rows); see + [Entitlements](./entitlements.md#user-meter-expand-phase). - `RemoteConnectorSession` exposes persisted connector metadata and tool descriptors through an export RPC. - `PackageServiceInstance` uses its status RPC as the stable persisted service @@ -445,6 +452,55 @@ Storage split: due jobs - `StorageRunner` SQLite: isolated durable state addressed by `storageId` +## Durable Objects (`UserMeter`) + +Daily rate-style entitlement counters and inbound email delivery-id idempotency +live in a per-user `UserMeter` Durable Object with SQLite +(`packages/worker/src/entitlements/user-meter-do.ts`). The Worker binding is +`USER_METER` (class `UserMeter`; Wrangler SQLite migration tag `v21` via +`new_sqlite_classes` in `packages/worker/wrangler.jsonc`). + +Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP +`userId` via `userMeterDurableObjectName(userId)` → `idFromName(userId)` in +`packages/worker/src/user-scoped-durable-object-name.ts`. There is no `user_id` +column inside the DO because the object identity is the user. + +SQLite ownership (schema version tracked in `user_meter_meta`): + +- `daily_counters` — authoritative UTC-day counters for `email_sends_per_day`, + `email_receives_per_day`, `execute_calls_per_day`, and + `outbound_fetches_per_day` (`resource`, `day`, `count`, monotonic `revision`, + `updated_at`). +- `inbound_delivery_claims` — idempotency ledger keyed by inbound `delivery_id` + (scoped by DO identity, so the primary key is delivery id alone). Records the + claim's resource/day, post-charge counter, revision, and `claimed_at` so + Cloudflare Email Routing retries inside the 48-hour inbound dedupe window + cannot double-charge `email_receives_per_day`. + +Retention is self-enforced inside the DO: every read/write path +opportunistically deletes counter and claim rows older than seven UTC days +(`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; +the window covers timezone edge cases, recent account exports, and inbound +retries. + +**Expand-phase D1 mirror:** enforcement and point reads are authoritative in +UserMeter. D1 `entitlement_daily_counters` is **not** dropped — it remains a +best-effort mirror for existing readers and reporting. After each DO +consume/refund/inbound claim, the entitlements service schedules a non-awaited +absolute mirror write keyed by `(user_id, resource, day)` with a +revision-ordered `updated_at` token (`r/` + zero-padded revision from +`userMeterMirrorUpdatedAtToken`) so late writes cannot overwrite newer state. +See [Entitlements](./entitlements.md#user-meter-expand-phase). + +**Cold bootstrap:** a missing `(resource, day)` row returns `needs_bootstrap`. +The service performs one legacy D1 point read on `entitlement_daily_counters`, +then `initialize()` seeds the DO row with `INSERT OR IGNORE` (concurrent callers +cannot double-apply the baseline). Warm paths never read D1 for enforcement. + +Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan). +Account export pages `UserMeter.exportCounters` through the `user_meter` +manifest section / `account_export_section`. + ### Package state model Saved packages are the only top-level persisted primitive. Their state maps onto @@ -478,6 +534,9 @@ via `durableObjectNameFromParts`); domain helpers such as - `RunLog` — `runLogDurableObjectName(userId)` → `idFromName(userId)`. One execution-history DO per user; there is no `user_id` column inside it because the DO identity is the user. See [Run records](./run-records.md). +- `UserMeter` — `userMeterDurableObjectName(userId)` → `idFromName(userId)`. One + daily-entitlement meter DO per user (untrimmed stable id, same as `RunLog`). + See [Entitlements](./entitlements.md#user-meter-expand-phase). - `McpClientHub` — `mcpClientHubDurableObjectName(userId)` → `idFromName(userId.trim())`. - `StorageRunner` — `storageRunnerDurableObjectName(userId, storageId)` → @@ -542,6 +601,8 @@ Bindings are configured per environment in `packages/worker/wrangler.jsonc` - `JOB_MANAGER` (Durable Objects) - `RUN_LOG` (Durable Objects; per-user run records — see [Run records](./run-records.md)) +- `USER_METER` (Durable Objects; per-user daily entitlement counters — see + [Entitlements](./entitlements.md#user-meter-expand-phase)) - `STORAGE_RUNNER` (Durable Objects) - `REPO_SESSION` (Durable Objects) - `PACKAGE_REALTIME_SESSION` (Durable Objects) @@ -786,6 +847,8 @@ to `durableObjectNameFromParts`). - `JobManager`: `idFromName(userId)` (no trim). - `RunLog`: `idFromName(userId)` (no trim); one execution-history DO per user. +- `UserMeter`: `idFromName(userId)` (no trim); one daily-entitlement meter DO + per user. - `McpClientHub`: `idFromName(userId.trim())`. - `StorageRunner`: `idFromName(JSON.stringify([userId, storageId]))`. - `RepoSession`: `idFromName(repo_sessions.id)`; the key is not user-prefixed, @@ -953,7 +1016,10 @@ Current retention policies: blob delete fails, those rows are skipped and still selected on later runs so cleanup can retry. Dependent `email_attachments` rows are deleted before their messages, and threads left with no messages are pruned for the affected users. -- `entitlement_daily_counters`: daily rate counters keep 400 days by `day` key. +- `entitlement_daily_counters`: expand-phase **mirror** of UserMeter daily + counters (authoritative state lives in the per-user `UserMeter` DO). Rows keep + 400 days by `day` key until mirror retirement is verified after + reporting-off-D1 merges; the table is not dropped in this phase. - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; raw Analytics Engine usage events follow platform retention. - `feature_flag_exposure_rollups`: local-dev/test flag exposure rollups keep 90 diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 5b66773a99..5bee0a0a6f 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -107,7 +107,68 @@ two compute surfaces `usage-metering.md` already observes: cannot grant elevated quotas. Both consume only when the context has a `userId`, matching the usage-metering -rule that events without an owning user are skipped. +rule that events without an owning user are skipped. Daily consumption is +authoritative in the per-user `UserMeter` Durable Object; see +[UserMeter (expand phase)](#user-meter-expand-phase). + +## UserMeter (expand phase) + +Daily rate-style resources (`email_sends_per_day`, `email_receives_per_day`, +`execute_calls_per_day`, `outbound_fetches_per_day`) are **authoritative in the +per-user `UserMeter` Durable Object** (`USER_METER` binding). Code lives in +`packages/worker/src/entitlements/user-meter-do.ts` and `user-meter-client.ts`; +storage layout and naming are documented in [Data storage](./data-storage.md). + +**Strong enforcement:** `consumeDailyEntitlement` and inbound +`consumeInboundDelivery` RPCs check the plan limit and increment inside the DO. +The Durable Object request model serializes mutations per user; counter updates +use optimistic concurrency on monotonic `revision` so concurrent consumes cannot +overshoot. Missing `(resource, day)` rows return `needs_bootstrap` rather than +silently starting at zero on a warm account. + +**Cold bootstrap:** on `needs_bootstrap`, the service performs one legacy D1 +point read on `entitlement_daily_counters`, then `UserMeter.initialize()` seeds +the row with `INSERT OR IGNORE` (concurrent callers cannot double-apply the +baseline). The warm enforcement path awaits only the DO RPC — never APP_DB. + +**Non-awaited D1 mirror:** after each successful consume, refund, or inbound +delivery claim, the service schedules a best-effort absolute mirror write to +`entitlement_daily_counters` via `waitUntil` when available (otherwise a caught +void promise). Mirror ordering uses the DO-minted `mirrorUpdatedAt` token +(`r/` + zero-padded revision) in the existing `updated_at` TEXT column so late +writes cannot overwrite newer state, including refunds that lower `count`. +Mirror failures are logged and never affect enforcement. The D1 table is **not** +dropped in this phase — it remains for existing readers and reporting. + +**Point-read surfaces** now call `readDailyEntitlementResourceUsage` (UserMeter +with the same cold-bootstrap path) instead of reading D1 directly: + +- Account usage UI — `packages/worker/src/app/account-usage-data.ts` +- `email_usage_get` MCP capability +- Admin per-user usage drill-down — + `packages/worker/src/admin/user-usage-data.ts` + +Non-daily resources and contexts without `USER_METER` still use +`readEntitlementResourceUsage` against D1. + +**Inbound retry idempotency:** inbound receive quota uses +`UserMeter.consumeInboundDelivery`, which atomically claims `delivery_id` and +consumes one `email_receives_per_day` unit inside a SQLite transaction. Retries +return the accepted counter without incrementing (`replayed: true`). +Cross-UTC-day retries use the original claim's resource/day. The legacy D1 +mirror is scheduled from the email path via +`scheduleAbsoluteDailyEntitlementMirror` so the email subsystem does not +duplicate mirror SQL. + +**Remaining migration dependency:** retiring the D1 mirror and +`entitlement_daily_counters` table waits until reporting-off-D1 work merges and +mirror parity is verified in production. This expand phase intentionally does +**not** move `storage_bytes` enforcement, `package_service_states` running +counts, or account-deletion write fencing to UserMeter. + +Module wiring: `consumeDailyEntitlement`, `refundDailyEntitlement`, and +`readDailyEntitlementResourceUsage` in `service.ts` require `env.USER_METER` and +fail closed when the binding is missing. ## Schema history @@ -245,21 +306,35 @@ Rules: package services are counted from `package_service_states` (status `running` and freshly heartbeaten), not from run-history rows — see [Run records](./run-records.md) (`state-vs-history`). -- **Rate-style limits** (email sends and receives per day) use the - `entitlement_daily_counters` table keyed by `(user_id, resource, day)` with - UTC day keys. Call `consumeDailyEntitlement` on every attempt: it checks the - plan limit from `resolvePlanLimit` and increments the counter in one - conditional D1 upsert (no check-then-increment race). Every resolved plan has - a finite numeric limit. Counting attempts rather than successes keeps the - limit abuse-resistant for permanent rejects (parse failures, entitlement/quota - rejects). Only typed pre-commit `RetryableInboundStorageError` failures - (thread prework, R2 put, D1 message/attachment storage after successful - cleanup) refund exactly one `email_receives_per_day` unit via - `refundDailyEntitlement` for the same UTC day that was charged, so Cloudflare - Email Routing retries do not burn the daily receive quota. Post-commit - bookkeeping failures do not refund or retry. - `incrementDailyEntitlementCounter` remains for raw counter writes (tests, - backfills). +- **Rate-style limits** (email sends/receives per day, execute calls per day, + outbound fetches per day) are **authoritative in the per-user UserMeter + Durable Object** (UTC day keys). Call `consumeDailyEntitlement` on every + attempt: it resolves the plan limit, atomically checks and increments inside + the DO, and throws `EntitlementLimitError` when over limit. Every resolved + plan has a finite numeric limit. Counting attempts rather than successes keeps + the limit abuse-resistant for permanent rejects (parse failures, + entitlement/quota rejects). + + **Cold bootstrap:** missing `(resource, day)` rows trigger one legacy D1 point + read and a single `UserMeter.initialize()` before retrying the consume. + + **D1 mirror (expand phase):** after each successful consume/refund/inbound + claim, a best-effort, non-awaited absolute mirror write updates + `entitlement_daily_counters` with revision-ordered `updated_at` tokens. The + table is not dropped — it remains for existing readers and reporting until + reporting-off-D1 retirement is verified. + + Only typed pre-commit `RetryableInboundStorageError` failures (thread prework, + R2 put, D1 message/attachment storage after successful cleanup) refund exactly + one `email_receives_per_day` unit via `refundDailyEntitlement` for the same + UTC day that was charged, so Cloudflare Email Routing retries do not burn the + daily receive quota. Post-commit bookkeeping failures do not refund or retry. + Inbound routing retries that replay the same `delivery_id` are idempotent via + `UserMeter.consumeInboundDelivery` and do not consume an additional unit. + + `incrementDailyEntitlementCounter` remains for raw D1 counter writes (tests, + backfills, and legacy paths). + - **Boolean allowances** (persistent package services) are modeled as limit `0` (not allowed) vs `1` (allowed) so the numeric contract stays uniform. - **Per-unit size limits** (`email_message_bytes`) compare one candidate value @@ -328,9 +403,9 @@ separate statements, so a burst of concurrent creates can overshoot a limit by a few rows before the next check sees the new count. That is an accepted trade-off — these limits are order-of-magnitude denial-of-wallet caps, not billing-grade accounting, and folding every insert into a conditional statement would couple -the entitlements module to each resource's write path. The rate-style path does -not share this window: `consumeDailyEntitlement` checks and increments in a -single conditional upsert. +the entitlements module to each resource's write path. Daily rate-style limits +do not share this window: UserMeter consumes are serialized per user inside the +DO with revision-checked updates. ## How to add an enforcement point @@ -459,6 +534,9 @@ in [`../environment-variables.md`](../environment-variables.md). `0066-stripe-billing.sql`; owned by `packages/worker/src/billing/`, read by `getUserPlan` via `resolveEffectivePlan`. `stripe_plan` stays nullable because it is Stripe-derived; `max` is manual-only. -- `entitlement_daily_counters` — rate counters, created by migration - `0048-user-plans-and-entitlement-counters.sql`; included in the +- `entitlement_daily_counters` — expand-phase **mirror** of UserMeter daily + counters (authoritative state in the per-user `UserMeter` DO), created by + migration `0048-user-plans-and-entitlement-counters.sql`; included in the account-deletion cascade (`packages/worker/src/app/account-deletion.ts`). + Table retirement waits until reporting-off-D1 merges and mirror parity is + verified. diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index 21d6e45f63..e559d43738 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -487,6 +487,19 @@ primitives: docs: - docs/contributing/architecture/run-records.md + - id: user-meter + group: storage + name: User meter + summary: + Per-user UserMeter DO SQLite for daily entitlement counters and inbound + delivery idempotency. + code: + - packages/worker/src/entitlements/user-meter-do.ts + - packages/worker/src/entitlements/user-meter-client.ts + docs: + - docs/contributing/architecture/data-storage.md + - docs/contributing/architecture/entitlements.md + - id: email-blobs-r2 group: storage name: Email blob storage (R2) @@ -605,3 +618,4 @@ invariants: docs: - docs/contributing/architecture/usage-metering.md - docs/contributing/architecture/run-records.md + - docs/contributing/architecture/entitlements.md diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index d1d741058d..bdcaad8d14 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1474,6 +1474,130 @@ test('account export includes run_records section with runs and log lines', asyn ]) }) +test('account export includes user_meter counters, pages them, and warns on truncation', async () => { + const { sqlite, db } = createMigratedDb() + sqlite.exec(` + INSERT INTO users ( + id, username, email, password_hash, created_at, updated_at, + email_verified_at, stable_user_id + ) + VALUES ( + 1, 'user-a', 'a@example.com', 'password-hash-a', '2026-07-05', + '2026-07-05', '2026-07-05', 'user-aaa' + ); + `) + const counters = [ + { + resource: 'email_sends_per_day' as const, + day: '2026-07-30', + count: 2, + revision: 2, + updatedAt: '2026-07-30T01:00:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000002', + }, + { + resource: 'execute_calls_per_day' as const, + day: '2026-07-30', + count: 5, + revision: 5, + updatedAt: '2026-07-30T02:00:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000005', + }, + { + resource: 'outbound_fetches_per_day' as const, + day: '2026-07-31', + count: 1, + revision: 1, + updatedAt: '2026-07-31T00:00:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000001', + }, + ] + const exportCounters = vi.fn( + async (input: { pageSize?: number; startAfter?: string | null }) => { + const pageSize = input.pageSize ?? 100 + const startIndex = input.startAfter + ? counters.findIndex( + (row) => `${row.day}:${row.resource}` === input.startAfter, + ) + 1 + : 0 + const page = counters.slice(startIndex, startIndex + pageSize) + const truncated = startIndex + pageSize < counters.length + return { + counters: page, + nextStartAfter: truncated + ? `${page.at(-1)!.day}:${page.at(-1)!.resource}` + : null, + truncated, + } + }, + ) + const idFromName = vi.fn((name: string) => name as unknown as DurableObjectId) + const env = { + APP_DB: db, + USER_METER: { + idFromName, + get: () => ({ exportCounters }), + }, + } as unknown as Env + + const accountExport = await createAccountExport({ + env, + dbUserId: 1, + mcpUserId: 'user-aaa', + }) + expect(idFromName).toHaveBeenCalledWith('user-aaa') + expect(accountExport.manifest.sections.user_meter?.count).toBe(3) + expect(accountExport.durableObjects.userMeter).toEqual({ + counters, + nextStartAfter: null, + truncated: false, + }) + expect(accountExport.manifest.warnings).not.toEqual( + expect.arrayContaining([ + expect.stringContaining('entitlement_daily_counters'), + ]), + ) + + const first = await readAccountExportSection({ + env, + dbUserId: 1, + mcpUserId: 'user-aaa', + section: 'user_meter', + pageSize: 2, + }) + expect(first.items).toEqual(counters.slice(0, 2)) + expect(first.truncated).toBe(true) + expect(first.nextStartAfter).toBe('2026-07-30:execute_calls_per_day') + + const second = await readAccountExportSection({ + env, + dbUserId: 1, + mcpUserId: 'user-aaa', + section: 'user_meter', + pageSize: 2, + startAfter: first.nextStartAfter ?? undefined, + }) + expect(second.items).toEqual(counters.slice(2)) + expect(second.truncated).toBe(false) + expect(second.nextStartAfter).toBeNull() + expect(exportCounters).toHaveBeenCalled() + + exportCounters.mockImplementation(async () => ({ + counters: [counters[0]!], + nextStartAfter: 'cursor-more', + truncated: true, + })) + const truncatedExport = await createAccountExport({ + env, + dbUserId: 1, + mcpUserId: 'user-aaa', + }) + expect(truncatedExport.durableObjects.userMeter?.truncated).toBe(true) + expect(truncatedExport.manifest.warnings).toContain( + 'User meter counters were truncated in the full export; use account_export_section with section "user_meter" to retrieve additional pages.', + ) +}) + test('account export includes a package service known only via package_service_states', async () => { consoleWarn.mockImplementation(() => {}) try { diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index 6fe7f08622..a5874f0911 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -34,6 +34,11 @@ import { listRunRecordStorageIds, summarizeRunRecords, } from '#worker/run-records/service.ts' +import { + userMeterNamespace, + userMeterRpc, +} from '#worker/entitlements/user-meter-client.ts' +import { type UserMeterExportResult } from '#worker/entitlements/user-meter-do.ts' import { resolveUserStableId } from '#worker/user-id.ts' import { listAccountUserPackageServices, @@ -55,6 +60,7 @@ export const accountExportSectionNames = [ 'storage_runner', 'job_manager', 'run_records', + 'user_meter', 'remote_connector_session', 'package_service', 'oauth_grants', @@ -126,6 +132,7 @@ type UserExportInventory = { type ManifestInventoryCounts = { storageRunners: number runRecords: number + userMeterCounters: number remoteConnectorSessions: number packageServices: number artifactRepos: number @@ -198,6 +205,7 @@ export type AccountExportArtifactRepo = { type AccountExportDurableObjects = { jobManager: unknown | null runRecords: Awaited> | null + userMeter: UserMeterExportResult | null remoteConnectorSessions: Array<{ instanceId: string export: RemoteConnectorSessionExport | null @@ -977,6 +985,7 @@ async function collectManifestInventoryCounts(input: { const [ storageRunners, runRecords, + userMeterCounters, remoteConnectorSessions, packageServices, artifactRepos, @@ -993,6 +1002,16 @@ async function collectManifestInventoryCounts(input: { }) return summary.total }), + safeCount('user meter counters', async () => { + if (!userMeterNamespace(input.env)) return 0 + const page = await userMeterRpc({ + env: input.env, + userId: input.userId, + }).exportCounters({ + pageSize: maxExportPageSize, + }) + return page.counters.length + }), safeCount('remote connectors', async () => countScalar( input.env, @@ -1030,6 +1049,7 @@ async function collectManifestInventoryCounts(input: { return { storageRunners, runRecords, + userMeterCounters, remoteConnectorSessions, packageServices, artifactRepos, @@ -1272,6 +1292,36 @@ async function exportUserRunRecords(input: { } } +async function exportUserMeterCounters(input: { + env: AccountExportEnv + userId: string + warnings: Array +}): Promise { + try { + if (!userMeterNamespace(input.env)) { + input.warnings.push( + 'USER_METER binding was unavailable; user meter counters were not exported.', + ) + return null + } + const page = await userMeterRpc({ + env: input.env, + userId: input.userId, + }).exportCounters({ + pageSize: maxExportPageSize, + }) + if (page.truncated) { + input.warnings.push( + `User meter counters were truncated in the full export; use account_export_section with section "user_meter" to retrieve additional pages.`, + ) + } + return page + } catch (error) { + input.warnings.push(`User meter export failed: ${getErrorMessage(error)}`) + return null + } +} + async function exportRemoteConnectorSessions(input: { env: AccountExportEnv userId: string @@ -1364,32 +1414,42 @@ async function exportDurableObjects(input: { inventory: UserExportInventory warnings: Array }): Promise { - const [storageRunners, remoteConnectorSessions, packageServices, runRecords] = - await Promise.all([ - exportStorageRunners({ - env: input.env, - userId: input.userId, - storageIds: input.inventory.storageIds, - warnings: input.warnings, - }), - exportRemoteConnectorSessions({ - env: input.env, - userId: input.userId, - connectors: input.inventory.remoteConnectors, - warnings: input.warnings, - }), - exportPackageServices({ - env: input.env, - userId: input.userId, - services: input.inventory.packageServices, - warnings: input.warnings, - }), - exportUserRunRecords({ - env: input.env, - userId: input.userId, - warnings: input.warnings, - }), - ]) + const [ + storageRunners, + remoteConnectorSessions, + packageServices, + runRecords, + userMeter, + ] = await Promise.all([ + exportStorageRunners({ + env: input.env, + userId: input.userId, + storageIds: input.inventory.storageIds, + warnings: input.warnings, + }), + exportRemoteConnectorSessions({ + env: input.env, + userId: input.userId, + connectors: input.inventory.remoteConnectors, + warnings: input.warnings, + }), + exportPackageServices({ + env: input.env, + userId: input.userId, + services: input.inventory.packageServices, + warnings: input.warnings, + }), + exportUserRunRecords({ + env: input.env, + userId: input.userId, + warnings: input.warnings, + }), + exportUserMeterCounters({ + env: input.env, + userId: input.userId, + warnings: input.warnings, + }), + ]) let jobManager: unknown | null = null try { jobManager = await exportJobManagerForUser({ @@ -1402,6 +1462,7 @@ async function exportDurableObjects(input: { return { jobManager, runRecords, + userMeter, remoteConnectorSessions, packageServices, storageRunners, @@ -1467,6 +1528,17 @@ function buildManifest(input: { ), discovery: { section: 'run_records' }, } + sections.user_meter = { + count: + input.durableObjects?.userMeter == null + ? (input.inventoryCounts?.userMeterCounters ?? 0) + : input.durableObjects.userMeter.counters.length, + warnings: input.warnings.filter( + (warning) => + warning.startsWith('User meter ') || warning.startsWith('USER_METER '), + ), + discovery: { section: 'user_meter' }, + } sections.remote_connector_sessions = { count: input.inventoryCounts?.remoteConnectorSessions ?? @@ -1792,6 +1864,27 @@ export async function readAccountExportSection(input: { warnings, } } + if (input.section === 'user_meter') { + if (!userMeterNamespace(input.env)) { + throw new Error('USER_METER binding was unavailable.') + } + const pageSize = normalizePageSize(input.pageSize) + const page = await userMeterRpc({ + env: input.env, + userId: input.mcpUserId, + }).exportCounters({ + pageSize, + startAfter: input.startAfter ?? null, + }) + return { + section: input.section, + items: page.counters, + truncated: page.truncated, + nextStartAfter: page.nextStartAfter, + pageSize, + warnings, + } + } if (input.section === 'remote_connector_session') { if (!input.instanceId) { throw new Error( diff --git a/packages/worker/src/account/user-owned-surfaces.node.test.ts b/packages/worker/src/account/user-owned-surfaces.node.test.ts index c1a6dff39c..3772f2dc93 100644 --- a/packages/worker/src/account/user-owned-surfaces.node.test.ts +++ b/packages/worker/src/account/user-owned-surfaces.node.test.ts @@ -59,11 +59,25 @@ test('account deletion and export consume the out-of-band surface registry', () ).toBe(true) expect(coverage.r2SurfaceIds.has('email_raw_mime')).toBe(true) expect(coverage.r2SurfaceIds.has('user_avatar')).toBe(true) + expect(coverage.durableObjectIds.has('user_meter')).toBe(true) + expect( + accountUserOwnedDurableObjectSurfaces.find( + (surface) => surface.id === 'user_meter', + ), + ).toMatchObject({ + binding: 'USER_METER', + deletionResultKey: 'userMeters', + export: 'include', + }) expect( accountUserOwnedKvKeySchemes.some((scheme) => scheme.prefixTemplate?.includes('source-snapshot:v1:'), ), ).toBe(true) + expect(accountExportSource).toContain("'user_meter'") + expect(accountExportSource).toContain('userMeterRpc') + expect(accountDeletionSource).toContain('userMeterRpc') + expect(accountDeletionSource).toContain('userMeters') expect(accountDeletionSource).toContain( "from '#worker/account/user-owned-surfaces.ts'", diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 38cce17acb..9975b02533 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -8,6 +8,7 @@ export type UserOwnedDurableObjectSurface = { | 'package_realtime_session' | 'package_service_instance' | 'run_log' + | 'user_meter' | 'mcp' binding: string /** Result key used in AccountDeletionResult.clearedDurableObjects when purged */ @@ -123,6 +124,14 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray [] as Array, }), }, + USER_METER: { + idFromName: durableObjectId, + get: () => ({ purge: async () => ({ ok: true as const }) }), + }, JOB_MANAGER: { idFromName: durableObjectId, get: () => ({ purgeUser: async () => ({ ok: true as const }) }), @@ -1106,6 +1110,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async const clearStorageMock = vi.fn(async () => ({ ok: true as const })) const clearRunLogMock = vi.fn(async () => ({ ok: true as const })) + const purgeUserMeterMock = vi.fn(async () => ({ ok: true as const })) const purgeJobManagerMock = vi.fn(async () => ({ ok: true as const })) const purgeRepoSessionMock = vi.fn(async () => ({ ok: true as const })) const purgeRemoteConnectorMock = vi.fn(async () => ({ ok: true as const })) @@ -1131,6 +1136,10 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async listStorageIds: async () => [] as Array, }), }, + USER_METER: { + idFromName: (name: string) => name as unknown as DurableObjectId, + get: () => ({ purge: purgeUserMeterMock }), + }, JOB_MANAGER: { idFromName: (name: string) => name as unknown as DurableObjectId, get: () => ({ purgeUser: purgeJobManagerMock }), @@ -1431,6 +1440,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async expect(result.clearedDurableObjects).toMatchObject({ storageRunners: 6, runLogs: 1, + userMeters: 1, jobManagers: 1, repoSessions: 1, remoteConnectorSessions: 1, @@ -1443,6 +1453,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async packageServiceInstances: 3, }) expect(clearRunLogMock).toHaveBeenCalledTimes(1) + expect(purgeUserMeterMock).toHaveBeenCalledTimes(1) expect(purgeMcpClientHubMock).toHaveBeenCalledTimes(1) expect(purgeMcpAgentSessionMock).toHaveBeenCalledWith({ userId: userAaa, @@ -1990,6 +2001,36 @@ test('account deletion empties the user RunLog DO and leaves other users untouch expect(clearedStorageIds.some((id) => id.includes('bbb-bucket'))).toBe(false) }) +test('account deletion reports a missing USER_METER binding and remains retryable', async () => { + const { db, rows } = createTestDb({ + users: [{ id: 1, email: 'a@example.com', stable_user_id: 'user-aaa' }], + }) + await expect( + deleteUserAccount({ + env: createSuccessfulDeletionEnv(db, { + USER_METER: undefined, + }), + dbUserId: 1, + mcpUserId: 'user-aaa', + }), + ).rejects.toMatchObject({ + cleanupErrors: expect.arrayContaining([ + 'USER_METER binding was unavailable; the user meter Durable Object was not purged.', + ]), + partialResult: expect.objectContaining({ + clearedDurableObjects: expect.objectContaining({ + userMeters: 0, + }), + }), + }) + expect(rows.users).toEqual([ + expect.objectContaining({ + id: 1, + deleting_at: expect.any(String), + }), + ]) +}) + test('account deletion purges a StorageRunner known only via user_storage_buckets', async () => { const userId = 'user-bucket-only' const clearStorage = vi.fn(async () => ({ ok: true as const })) diff --git a/packages/worker/src/app/account-deletion.ts b/packages/worker/src/app/account-deletion.ts index 50ed4395be..bb2d825d7d 100644 --- a/packages/worker/src/app/account-deletion.ts +++ b/packages/worker/src/app/account-deletion.ts @@ -17,6 +17,10 @@ import { mcpClientHubDurableObjectName } from '#worker/user-scoped-durable-objec import { packageServiceRpc } from '#worker/package-runtime/package-service.ts' import { packageRealtimeSessionRpc } from '#worker/package-runtime/realtime-session.ts' import { clearRunRecords } from '#worker/run-records/service.ts' +import { + userMeterNamespace, + userMeterRpc, +} from '#worker/entitlements/user-meter-client.ts' import { listAccountUserPackageServices, listAccountUserStorageIds, @@ -596,6 +600,27 @@ async function clearRunLog(input: { } } +async function purgeUserMeter(input: { + env: Env + userId: string + warnings: Array +}): Promise { + try { + if (!userMeterNamespace(input.env)) { + input.warnings.push( + 'USER_METER binding was unavailable; the user meter Durable Object was not purged.', + ) + return 0 + } + await userMeterRpc({ env: input.env, userId: input.userId }).purge() + return 1 + } catch (error) { + const message = getErrorMessage(error) + input.warnings.push(`User meter purge failed: ${message}`) + return 0 + } +} + async function purgeJobManager(input: { env: Env userId: string @@ -1092,6 +1117,11 @@ export async function deleteUserAccount(input: { userId: input.mcpUserId, warnings, }) + result.clearedDurableObjects.userMeters = await purgeUserMeter({ + env: input.env, + userId: input.mcpUserId, + warnings, + }) if (input.env.BUNDLE_ARTIFACTS_KV) { const sourceSnapshotKeys = await listKvKeysByPrefix({ From 4a60dd992fe172d797fc2084469042efc4cbbbef Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:23:10 +0000 Subject: [PATCH 05/39] test(entitlements): provide UserMeter fixtures Co-authored-by: Kent C. Dodds --- .../worker/src/email/inbound.workers.test.ts | 309 +++++++++--------- .../admin/admin-capabilities.node.test.ts | 4 + .../operation-request.node.test.ts | 4 + .../worker/src/mcp/observability.node.test.ts | 4 + .../worker/src/mcp/tools/execute.node.test.ts | 4 + 5 files changed, 174 insertions(+), 151 deletions(-) diff --git a/packages/worker/src/email/inbound.workers.test.ts b/packages/worker/src/email/inbound.workers.test.ts index 7d9ff38520..feadc77b7b 100644 --- a/packages/worker/src/email/inbound.workers.test.ts +++ b/packages/worker/src/email/inbound.workers.test.ts @@ -119,164 +119,171 @@ async function seedVerifiedAccount(input: { // budget these explicitly like the other sandbox-executing suites rather than // letting them flake under a loaded `npm run validate`. const subscriptionDispatchTimeoutMs = 60_000 +// handleInboundEmail touches the real USER_METER DO; isolated runs finish in +// a few seconds but the default 5s budget flakes under full-suite contention. +const inboundRoutingTimeoutMs = 15_000 -test('inbound email routes {username}@platform-domain and auto-provisions the default inbox', async () => { - // Usage recording degrades with a warn when the usage_rollups table is - // not part of this test's schema; that is incidental to routing. - silenceIncidentalRuntimeWarnings() - await ensureEmailTestSchema(env.APP_DB) - const username = `inbound-${crypto.randomUUID().slice(0, 8)}` - const accountEmail = `account-${crypto.randomUUID()}@example.com` - const userId = await createStableUserIdFromEmail(accountEmail) - const address = `${username}@${platformDomain}` - await seedVerifiedAccount({ - db: env.APP_DB, - email: accountEmail, - username, - }) - - const firstMessage = createForwardableEmailMessage({ - from: 'stranger@example.net', - to: address, - raw: [ - 'From: Stranger ', - `To: ${address}`, - 'Subject: Unknown sender', - 'Message-ID: ', - '', - 'Please help.', - ].join('\r\n'), - }) - await handleInboundEmail(firstMessage, createInboundEnv()) - expect(firstMessage.rejectedReason).toBeNull() +test( + 'inbound email routes {username}@platform-domain and auto-provisions the default inbox', + async () => { + // Usage recording degrades with a warn when the usage_rollups table is + // not part of this test's schema; that is incidental to routing. + silenceIncidentalRuntimeWarnings() + await ensureEmailTestSchema(env.APP_DB) + const username = `inbound-${crypto.randomUUID().slice(0, 8)}` + const accountEmail = `account-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(accountEmail) + const address = `${username}@${platformDomain}` + await seedVerifiedAccount({ + db: env.APP_DB, + email: accountEmail, + username, + }) - // First inbound provisioned the default inbox and address. - const inboxes = await listEmailInboxesForUser({ db: env.APP_DB, userId }) - expect(inboxes).toHaveLength(1) - expect(inboxes[0]).toMatchObject({ name: defaultEmailInboxName }) - const addresses = await listEmailInboxAddressesForUser({ - db: env.APP_DB, - userId, - }) - expect(addresses).toHaveLength(1) - expect(addresses[0]).toMatchObject({ - address, - localPart: username, - domain: platformDomain, - }) - // Provisioning also created the platform-assigned verified sender - // identity for the same address. - const identity = await env.APP_DB.prepare( - `SELECT email, domain, status FROM email_sender_identities + const firstMessage = createForwardableEmailMessage({ + from: 'stranger@example.net', + to: address, + raw: [ + 'From: Stranger ', + `To: ${address}`, + 'Subject: Unknown sender', + 'Message-ID: ', + '', + 'Please help.', + ].join('\r\n'), + }) + await handleInboundEmail(firstMessage, createInboundEnv()) + expect(firstMessage.rejectedReason).toBeNull() + + // First inbound provisioned the default inbox and address. + const inboxes = await listEmailInboxesForUser({ db: env.APP_DB, userId }) + expect(inboxes).toHaveLength(1) + expect(inboxes[0]).toMatchObject({ name: defaultEmailInboxName }) + const addresses = await listEmailInboxAddressesForUser({ + db: env.APP_DB, + userId, + }) + expect(addresses).toHaveLength(1) + expect(addresses[0]).toMatchObject({ + address, + localPart: username, + domain: platformDomain, + }) + // Provisioning also created the platform-assigned verified sender + // identity for the same address. + const identity = await env.APP_DB.prepare( + `SELECT email, domain, status FROM email_sender_identities WHERE user_id = ? AND email = ?`, - ) - .bind(userId, address) - .first>() - expect(identity).toEqual({ - email: address, - domain: platformDomain, - status: 'verified', - }) - const inbox = inboxes[0]! - - const secondMessage = createForwardableEmailMessage({ - from: 'agent@trusted.example', - to: `${username.toUpperCase()}@${platformDomain}`, - raw: [ - 'From: Agent ', - `To: ${address}`, - 'Subject: Approved sender', - 'Message-ID: ', - '', - 'Approved body.', - ].join('\r\n'), - }) - await handleInboundEmail(secondMessage, createInboundEnv()) - expect(secondMessage.rejectedReason).toBeNull() + ) + .bind(userId, address) + .first>() + expect(identity).toEqual({ + email: address, + domain: platformDomain, + status: 'verified', + }) + const inbox = inboxes[0]! - // The second delivery reuses the provisioned inbox instead of creating - // another one. - expect( - await listEmailInboxesForUser({ db: env.APP_DB, userId }), - ).toHaveLength(1) + const secondMessage = createForwardableEmailMessage({ + from: 'agent@trusted.example', + to: `${username.toUpperCase()}@${platformDomain}`, + raw: [ + 'From: Agent ', + `To: ${address}`, + 'Subject: Approved sender', + 'Message-ID: ', + '', + 'Approved body.', + ].join('\r\n'), + }) + await handleInboundEmail(secondMessage, createInboundEnv()) + expect(secondMessage.rejectedReason).toBeNull() - const messages = await listEmailMessages({ - db: env.APP_DB, - userId, - inboxId: inbox.id, - limit: 10, - }) - expect(messages).toHaveLength(2) - expect(messages[0]).toMatchObject({ - fromAddress: 'agent@trusted.example', - subject: 'Approved sender', - processingStatus: 'stored', - }) - expect(messages[1]).toMatchObject({ - fromAddress: 'stranger@example.net', - subject: 'Unknown sender', - error: null, - }) + // The second delivery reuses the provisioned inbox instead of creating + // another one. + expect( + await listEmailInboxesForUser({ db: env.APP_DB, userId }), + ).toHaveLength(1) - const normalizedExistingThread = await createEmailThread({ - db: env.APP_DB, - userId, - inboxId: inbox.id, - subjectNormalized: 'normalized subject', - }) - const subjectOnlyMessage = createForwardableEmailMessage({ - from: 'sender@example.net', - to: address, - raw: [ - 'From: Sender ', - `To: ${address}`, - 'Subject: Re: Normalized Subject', - '', - 'Subject-only body.', - ].join('\r\n'), - }) - await handleInboundEmail(subjectOnlyMessage, createInboundEnv()) - const subjectOnly = await listEmailMessages({ - db: env.APP_DB, - userId, - inboxId: inbox.id, - limit: 1, - }) - expect(subjectOnly[0]?.threadId).not.toBe(normalizedExistingThread.id) + const messages = await listEmailMessages({ + db: env.APP_DB, + userId, + inboxId: inbox.id, + limit: 10, + }) + expect(messages).toHaveLength(2) + expect(messages[0]).toMatchObject({ + fromAddress: 'agent@trusted.example', + subject: 'Approved sender', + processingStatus: 'stored', + }) + expect(messages[1]).toMatchObject({ + fromAddress: 'stranger@example.net', + subject: 'Unknown sender', + error: null, + }) - // RFC 5233 subaddressing routes {username}+{tag} to the same inbox and - // keeps the full tagged address in the stored to_addresses so package - // handlers can dispatch on the tag. - const taggedAddress = `${username}+billing@${platformDomain}` - const taggedMessage = createForwardableEmailMessage({ - from: 'invoices@example.net', - to: taggedAddress, - raw: [ - 'From: Invoices ', - `To: ${taggedAddress}`, - 'Subject: Subaddressed mail', - 'Message-ID: ', - '', - 'Tagged body.', - ].join('\r\n'), - }) - await handleInboundEmail(taggedMessage, createInboundEnv()) - expect(taggedMessage.rejectedReason).toBeNull() - const taggedStored = await listEmailMessages({ - db: env.APP_DB, - userId, - inboxId: inbox.id, - limit: 1, - }) - expect(taggedStored[0]).toMatchObject({ - subject: 'Subaddressed mail', - toAddresses: [taggedAddress], - }) - // Still the same single auto-provisioned inbox after the tagged delivery. - expect( - await listEmailInboxesForUser({ db: env.APP_DB, userId }), - ).toHaveLength(1) -}) + const normalizedExistingThread = await createEmailThread({ + db: env.APP_DB, + userId, + inboxId: inbox.id, + subjectNormalized: 'normalized subject', + }) + const subjectOnlyMessage = createForwardableEmailMessage({ + from: 'sender@example.net', + to: address, + raw: [ + 'From: Sender ', + `To: ${address}`, + 'Subject: Re: Normalized Subject', + '', + 'Subject-only body.', + ].join('\r\n'), + }) + await handleInboundEmail(subjectOnlyMessage, createInboundEnv()) + const subjectOnly = await listEmailMessages({ + db: env.APP_DB, + userId, + inboxId: inbox.id, + limit: 1, + }) + expect(subjectOnly[0]?.threadId).not.toBe(normalizedExistingThread.id) + + // RFC 5233 subaddressing routes {username}+{tag} to the same inbox and + // keeps the full tagged address in the stored to_addresses so package + // handlers can dispatch on the tag. + const taggedAddress = `${username}+billing@${platformDomain}` + const taggedMessage = createForwardableEmailMessage({ + from: 'invoices@example.net', + to: taggedAddress, + raw: [ + 'From: Invoices ', + `To: ${taggedAddress}`, + 'Subject: Subaddressed mail', + 'Message-ID: ', + '', + 'Tagged body.', + ].join('\r\n'), + }) + await handleInboundEmail(taggedMessage, createInboundEnv()) + expect(taggedMessage.rejectedReason).toBeNull() + const taggedStored = await listEmailMessages({ + db: env.APP_DB, + userId, + inboxId: inbox.id, + limit: 1, + }) + expect(taggedStored[0]).toMatchObject({ + subject: 'Subaddressed mail', + toAddresses: [taggedAddress], + }) + // Still the same single auto-provisioned inbox after the tagged delivery. + expect( + await listEmailInboxesForUser({ db: env.APP_DB, userId }), + ).toHaveLength(1) + }, + inboundRoutingTimeoutMs, +) test('inbound email rejects unknown usernames, reserved locals, and foreign domains', async () => { // Usage recording degrades with a warn when the usage_rollups table is diff --git a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts index 95451cd0f8..2532aaaf91 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts @@ -13,6 +13,7 @@ import { adminUserGetCapability } from './admin-user-get.ts' import { adminUserListCapability } from './admin-user-list.ts' import { adminUserUpdateCapability } from './admin-user-update.ts' import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import { loadAdminUserByTarget } from '#worker/admin/users-data.ts' type UserRow = { @@ -504,12 +505,15 @@ function createAdminCapabilityTestDb(input: { return { db, auditEvents, passwordResets, userRoles, users } } +const userMeter = createInMemoryUserMeterEnv() + function createAdminCapabilityContext( db: D1Database, blobs?: Pick, ) { return { env: { + ...userMeter.env, APP_DB: db, EMAIL_BLOBS: blobs ?? { get: async () => null, diff --git a/packages/worker/src/mcp/capabilities/openapi-provider/operation-request.node.test.ts b/packages/worker/src/mcp/capabilities/openapi-provider/operation-request.node.test.ts index 7ebd2a30d5..5f1fd83e70 100644 --- a/packages/worker/src/mcp/capabilities/openapi-provider/operation-request.node.test.ts +++ b/packages/worker/src/mcp/capabilities/openapi-provider/operation-request.node.test.ts @@ -7,8 +7,11 @@ import * as integrationsService from '#worker/integrations/service.ts' import * as packageRepo from '#worker/package-registry/repo.ts' import * as usageModule from '#worker/usage/record-usage.ts' import { type OpenApiBinding } from '#worker/openapi/binding-shared.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import { executeOpenApiOperationRequest } from './operation-request.ts' +const userMeter = createInMemoryUserMeterEnv() + const spotifyIntegrationConfig = { name: 'spotify', tokenUrl: 'https://accounts.spotify.com/api/token', @@ -137,6 +140,7 @@ function createEnv(entries: Map = new Map()) { } as unknown as D1Database return { + ...userMeter.env, APP_DB: db, SECRET_STORE_KEY: 'test-secret-store-key-32-chars-minimum', } as unknown as Env diff --git a/packages/worker/src/mcp/observability.node.test.ts b/packages/worker/src/mcp/observability.node.test.ts index 3383b2d96c..304f0b63ef 100644 --- a/packages/worker/src/mcp/observability.node.test.ts +++ b/packages/worker/src/mcp/observability.node.test.ts @@ -26,6 +26,10 @@ vi.mock('@sentry/cloudflare', () => ({ sentryMock.captureException(...args), captureMessage: (...args: Array) => sentryMock.captureMessage(...args), + instrumentDurableObjectWithSentry: ( + _getOptions: unknown, + durableObjectClass: unknown, + ) => durableObjectClass, })) const { logMcpEvent } = await import('./observability.ts') diff --git a/packages/worker/src/mcp/tools/execute.node.test.ts b/packages/worker/src/mcp/tools/execute.node.test.ts index bd331a22b1..67df4897a3 100644 --- a/packages/worker/src/mcp/tools/execute.node.test.ts +++ b/packages/worker/src/mcp/tools/execute.node.test.ts @@ -7,6 +7,7 @@ import { } from '#mcp/downstream-mcp-result.ts' import { formatRawFetchHostNudge } from '#mcp/raw-fetch-host-nudge.ts' import type * as RunRecordsServiceModule from '#worker/run-records/service.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' const mockModule = vi.hoisted(() => ({ runModuleWithRegistry: vi.fn(), @@ -59,6 +60,8 @@ vi.mock('#worker/run-records/service.ts', async () => { const { registerExecuteTool } = await import('./execute.ts') +const userMeter = createInMemoryUserMeterEnv() + /** * Minimal env stub: the daily execute entitlement consumed at the top of * the tool handler issues one conditional upsert (allowed when @@ -66,6 +69,7 @@ const { registerExecuteTool } = await import('./execute.ts') * contexts carry no account email (resolves to `max`). */ const stubEnv = { + ...userMeter.env, APP_DB: { prepare() { return { From 73a8e3712c4b337b841b633804585496f60ba58b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:25:44 +0000 Subject: [PATCH 06/39] test(entitlements): clean meter test imports Co-authored-by: Kent C. Dodds --- packages/worker/src/entitlements/entitlements.node.test.ts | 1 - packages/worker/src/entitlements/user-meter.workers.test.ts | 5 ++++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/worker/src/entitlements/entitlements.node.test.ts b/packages/worker/src/entitlements/entitlements.node.test.ts index 568a352fdc..024fab0a29 100644 --- a/packages/worker/src/entitlements/entitlements.node.test.ts +++ b/packages/worker/src/entitlements/entitlements.node.test.ts @@ -29,7 +29,6 @@ import { findCachedUserAccountByStableUserId, getCachedUserPlan, getUserPlan, - incrementDailyEntitlementCounter, packageServiceStateStaleMs, refundDailyEntitlement, } from './service.ts' diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index 4bfec33bfc..a9f11debc1 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -11,7 +11,10 @@ import { planLimits } from './plans.ts' import { consumeDailyEntitlement, refundDailyEntitlement } from './service.ts' import { ensureEntitlementTestSchema } from './test-schema.ts' import { userMeterRpc } from './user-meter-client.ts' -import { UserMeter, userMeterMirrorUpdatedAtToken } from './user-meter-do.ts' +import { + type UserMeter, + userMeterMirrorUpdatedAtToken, +} from './user-meter-do.ts' import { createWaitUntilDrain, withPatchedDbPrepare, From 8b001bf8f24850c2fa5d9cc29912e5ece3d1d30c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:41:45 +0000 Subject: [PATCH 07/39] style(mcp): format meter gateway tests Co-authored-by: Kent C. Dodds --- .../worker/src/mcp/fetch-gateway.workers.test.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/packages/worker/src/mcp/fetch-gateway.workers.test.ts b/packages/worker/src/mcp/fetch-gateway.workers.test.ts index 9fa02cee28..dafd062dd9 100644 --- a/packages/worker/src/mcp/fetch-gateway.workers.test.ts +++ b/packages/worker/src/mcp/fetch-gateway.workers.test.ts @@ -71,12 +71,12 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over }) expect(await allowed.text()).toBe('ok') await Promise.all(waitUntilTasks) - expect( - await userMeterRpc({ env, userId }).read({ - resource: 'outbound_fetches_per_day', - day: utcDayKey(), - }), - ).toMatchObject({ outcome: 'ready', count: 1 }) + expect( + await userMeterRpc({ env, userId }).read({ + resource: 'outbound_fetches_per_day', + day: utcDayKey(), + }), + ).toMatchObject({ outcome: 'ready', count: 1 }) // Fill the remaining free-plan quota in UserMeter, then expect denial. await seedOutboundFetchCounter( From df32168fdd3e7a70ded564c67a87849203c7f1e2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 00:53:50 +0000 Subject: [PATCH 08/39] fix(email): read daily usage from UserMeter Co-authored-by: Kent C. Dodds --- packages/worker/src/app/account-email-data.ts | 10 +- .../app/handlers/account-email.node.test.ts | 174 +++++++++++++++--- 2 files changed, 156 insertions(+), 28 deletions(-) diff --git a/packages/worker/src/app/account-email-data.ts b/packages/worker/src/app/account-email-data.ts index 6bacbb4af5..8a6ef3a4ff 100644 --- a/packages/worker/src/app/account-email-data.ts +++ b/packages/worker/src/app/account-email-data.ts @@ -12,8 +12,10 @@ import { } from '#worker/entitlements/plans.ts' import { getUserPlan, + readDailyEntitlementResourceUsage, readEntitlementResourceUsage, } from '#worker/entitlements/service.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' import { buildPlatformEmailAddress, getPlatformEmailDomain, @@ -369,6 +371,7 @@ async function countAndListMessages(input: { async function loadUsage(input: { db: D1Database + env: UserMeterEnv userId: string email: string }): Promise { @@ -384,14 +387,16 @@ async function loadUsage(input: { resource: 'stored_email_messages', now, }), - readEntitlementResourceUsage({ + readDailyEntitlementResourceUsage({ db: input.db, + env: input.env, userId: input.userId, resource: 'email_sends_per_day', now, }), - readEntitlementResourceUsage({ + readDailyEntitlementResourceUsage({ db: input.db, + env: input.env, userId: input.userId, resource: 'email_receives_per_day', now, @@ -578,6 +583,7 @@ export async function loadAccountEmailData(input: { loadInboxes({ db: input.env.APP_DB, userId }), loadUsage({ db: input.env.APP_DB, + env: input.env, userId, email: input.user.email, }), diff --git a/packages/worker/src/app/handlers/account-email.node.test.ts b/packages/worker/src/app/handlers/account-email.node.test.ts index d9cc2cdba2..14660553ee 100644 --- a/packages/worker/src/app/handlers/account-email.node.test.ts +++ b/packages/worker/src/app/handlers/account-email.node.test.ts @@ -1,6 +1,9 @@ import { expect, test, vi } from 'vitest' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import type * as EmailPlatformAddress from '#worker/email/platform-address.ts' import type * as EntitlementPlans from '#worker/entitlements/plans.ts' +import type * as EntitlementService from '#worker/entitlements/service.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' const messageRow = { id: 'msg-1', @@ -180,11 +183,15 @@ vi.mock('#app/email-verification.ts', () => ({ mockModule.isAccountEmailVerified(...args), })) -vi.mock('#worker/entitlements/service.ts', () => ({ - getUserPlan: (...args: Array) => mockModule.getUserPlan(...args), - readEntitlementResourceUsage: (...args: Array) => - mockModule.readEntitlementResourceUsage(...args), -})) +vi.mock('#worker/entitlements/service.ts', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + getUserPlan: (...args: Array) => mockModule.getUserPlan(...args), + readEntitlementResourceUsage: (...args: Array) => + mockModule.readEntitlementResourceUsage(...args), + } +}) vi.mock('#worker/entitlements/plans.ts', async (importOriginal) => { const actual = await importOriginal() @@ -239,13 +246,47 @@ function createListResult(rows: Array>) { } } -function createEnv() { - let prepareCall = 0 - mockModule.prepare.mockImplementation(() => { - prepareCall += 1 - // countAndListMessages issues COUNT then SELECT in Promise.all order. - if (prepareCall % 2 === 1) return createCountResult(1) - return createListResult([messageRow]) +function createEnv(input?: { + meter?: ReturnType + dailyCounters?: Array<{ + user_id: string + resource: string + day: string + count: number + }> + messageRows?: Array> + messageTotal?: number +}) { + const meter = input?.meter ?? createInMemoryUserMeterEnv() + const dailyCounters = input?.dailyCounters?.map((row) => ({ ...row })) ?? [] + const messageRows = input?.messageRows ?? [messageRow] + const messageTotal = input?.messageTotal ?? messageRows.length + mockModule.prepare.mockImplementation((query: string) => { + const normalized = String(query).replace(/\s+/g, ' ').trim().toLowerCase() + if (normalized.includes('from entitlement_daily_counters')) { + return { + bind(...params: Array) { + return { + async first() { + const row = dailyCounters.find( + (counter) => + counter.user_id === params[0] && + counter.resource === params[1] && + counter.day === params[2], + ) + return (row ? { count: row.count } : { count: 0 }) as T + }, + async all() { + return { results: [] } + }, + } + }, + } + } + if (normalized.includes('count(*)')) { + return createCountResult(messageTotal) + } + return createListResult(messageRows) }) return { APP_DB: { @@ -253,6 +294,7 @@ function createEnv() { } as unknown as D1Database, APP_BASE_URL: 'https://example.com', COOKIE_SECRET: 'secret', + ...meter.env, } as Env } @@ -260,7 +302,26 @@ test('email API lists messages with pagination, usage, and selected detail', asy mockModule.prepare.mockClear() mockModule.getEmailMessageById.mockClear() mockModule.listEmailDeliveryEvents.mockClear() - const env = createEnv() + const day = utcDayKey() + const userId = 'stable-user-1' + const bootstrapMeter = createInMemoryUserMeterEnv() + const env = createEnv({ + meter: bootstrapMeter, + dailyCounters: [ + { + user_id: userId, + resource: 'email_sends_per_day', + day, + count: 2, + }, + { + user_id: userId, + resource: 'email_receives_per_day', + day, + count: 2, + }, + ], + }) const handler = createAccountEmailApiHandler(env) const listResponse = await handler.handler({ @@ -299,6 +360,7 @@ test('email API lists messages with pagination, usage, and selected detail', asy selectedMessage: null, usage: expect.objectContaining({ plan: 'free', + day, stored_messages: { count: 2, limit: 100 }, sends_today: { count: 2, limit: 20 }, receives_today: { count: 2, limit: 50 }, @@ -313,7 +375,22 @@ test('email API lists messages with pagination, usage, and selected detail', asy mockModule.prepare.mockClear() mockModule.getEmailMessageById.mockClear() - const envWithSelection = createEnv() + const envWithSelection = createEnv({ + dailyCounters: [ + { + user_id: userId, + resource: 'email_sends_per_day', + day, + count: 2, + }, + { + user_id: userId, + resource: 'email_receives_per_day', + day, + count: 2, + }, + ], + }) const selectedResponse = await createAccountEmailApiHandler( envWithSelection, ).handler({ @@ -374,6 +451,57 @@ test('email API lists messages with pagination, usage, and selected detail', asy expect(unauthorizedResponse.status).toBe(401) }) +test('email API usage prefers initialized UserMeter daily counts over D1', async () => { + const day = utcDayKey() + const userId = 'stable-user-1' + const meter = createInMemoryUserMeterEnv() + await meter.seed({ + userId, + resource: 'email_sends_per_day', + day, + count: 13, + }) + await meter.seed({ + userId, + resource: 'email_receives_per_day', + day, + count: 15, + }) + const env = createEnv({ + meter, + dailyCounters: [ + { + user_id: userId, + resource: 'email_sends_per_day', + day, + count: 3, + }, + { + user_id: userId, + resource: 'email_receives_per_day', + day, + count: 5, + }, + ], + }) + const response = await createAccountEmailApiHandler(env).handler({ + request: new Request('https://example.com/account/email.json'), + }) + expect(response.status).toBe(200) + const body = await response.json() + expect(body).toMatchObject({ + ok: true, + usage: expect.objectContaining({ + day, + stored_messages: { count: 2, limit: 100 }, + sends_today: { count: 13, limit: 20 }, + receives_today: { count: 15, limit: 50 }, + }), + }) + expect(body.usage.sends_today.count).not.toBe(3) + expect(body.usage.receives_today.count).not.toBe(5) +}) + test('email API lists classification filters and classifies inbound messages', async () => { const quarantinedRow = { ...messageRow, @@ -381,12 +509,10 @@ test('email API lists classification filters and classifies inbound messages', a classification: 'quarantined', classification_reason: 'DMARC failed.', } - const env = createEnv() - let prepareCall = 0 - mockModule.prepare.mockImplementation(() => { - prepareCall += 1 - if (prepareCall % 2 === 1) return createCountResult(1) - return createListResult([quarantinedRow]) + const meter = createInMemoryUserMeterEnv() + const env = createEnv({ + meter, + messageRows: [quarantinedRow], }) mockModule.getEmailMessageById.mockResolvedValueOnce({ ...messageRecord, @@ -420,12 +546,8 @@ test('email API lists classification filters and classifies inbound messages', a }) mockModule.setEmailMessageClassification.mockClear() - let classifyPrepareCall = 0 - mockModule.prepare.mockImplementation(() => { - classifyPrepareCall += 1 - if (classifyPrepareCall % 2 === 1) return createCountResult(1) - return createListResult([messageRow]) - }) + // Reload after classify uses the default message list again. + createEnv({ meter, messageRows: [messageRow] }) const quarantineResponse = await handler.handler({ request: new Request('https://example.com/account/email.json', { method: 'POST', From e6622a42de6744486dcd6a94cbfb275f87a552e9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 01:15:58 +0000 Subject: [PATCH 09/39] docs(entitlements): clarify meter cutover semantics Co-authored-by: Kent C. Dodds --- .../contributing/architecture/entitlements.md | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 5bee0a0a6f..9f8fc15d46 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -144,6 +144,7 @@ dropped in this phase — it remains for existing readers and reporting. with the same cold-bootstrap path) instead of reading D1 directly: - Account usage UI — `packages/worker/src/app/account-usage-data.ts` +- Account email usage panel — `packages/worker/src/app/account-email-data.ts` - `email_usage_get` MCP capability - Admin per-user usage drill-down — `packages/worker/src/admin/user-usage-data.ts` @@ -151,6 +152,12 @@ with the same cold-bootstrap path) instead of reading D1 directly: Non-daily resources and contexts without `USER_METER` still use `readEntitlementResourceUsage` against D1. +During a rolling deployment, requests already running on the previous Worker +version may still increment D1 after a new-version request bootstraps its DO +row. Cloudflare activation bounds that overlap to in-flight requests, but +operators should treat mirror parity during the deploy window as approximate; +post-deploy requests have one authority in UserMeter. + **Inbound retry idempotency:** inbound receive quota uses `UserMeter.consumeInboundDelivery`, which atomically claims `delivery_id` and consumes one `email_receives_per_day` unit inside a SQLite transaction. Retries @@ -324,13 +331,11 @@ Rules: table is not dropped — it remains for existing readers and reporting until reporting-off-D1 retirement is verified. - Only typed pre-commit `RetryableInboundStorageError` failures (thread prework, - R2 put, D1 message/attachment storage after successful cleanup) refund exactly - one `email_receives_per_day` unit via `refundDailyEntitlement` for the same - UTC day that was charged, so Cloudflare Email Routing retries do not burn the - daily receive quota. Post-commit bookkeeping failures do not refund or retry. - Inbound routing retries that replay the same `delivery_id` are idempotent via - `UserMeter.consumeInboundDelivery` and do not consume an additional unit. + A delivery claim remains charged when later storage fails. Cloudflare Email + Routing retries replay that same `delivery_id` through + `UserMeter.consumeInboundDelivery` without incrementing again, including + across a UTC-day boundary. The retained claim is the idempotency boundary; + production inbound handling does not call `refundDailyEntitlement`. `incrementDailyEntitlementCounter` remains for raw D1 counter writes (tests, backfills, and legacy paths). From 7628aaf31365e6142f1698d1f486089b68bb9231 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 01:35:48 +0000 Subject: [PATCH 10/39] fix(entitlements): address UserMeter review feedback Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 8 +- .../contributing/architecture/entitlements.md | 2 +- .../app/handlers/account-email.node.test.ts | 12 +++ .../inbound-entitlements.workers.test.ts | 7 +- .../worker/src/email/inbound.workers.test.ts | 13 ++- .../worker/src/entitlements/user-meter-do.ts | 6 +- .../entitlements/user-meter.workers.test.ts | 83 +++++++++++++++++-- .../email/email-usage-get.workers.test.ts | 13 ++- .../src/mcp/fetch-gateway.workers.test.ts | 17 ++-- 9 files changed, 129 insertions(+), 32 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 25fece6776..f9ef5113b5 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -195,7 +195,7 @@ Durable Object export behavior: section (`exportCounters` RPC; keyset pagination by UTC `day` and `resource`). Retention is self-enforced inside the DO (seven UTC days of counter and inbound-delivery-claim rows); see - [Entitlements](./entitlements.md#user-meter-expand-phase). + [Entitlements](./entitlements.md#usermeter-expand-phase). - `RemoteConnectorSession` exposes persisted connector metadata and tool descriptors through an export RPC. - `PackageServiceInstance` uses its status RPC as the stable persisted service @@ -515,7 +515,7 @@ consume/refund/inbound claim, the entitlements service schedules a non-awaited absolute mirror write keyed by `(user_id, resource, day)` with a revision-ordered `updated_at` token (`r/` + zero-padded revision from `userMeterMirrorUpdatedAtToken`) so late writes cannot overwrite newer state. -See [Entitlements](./entitlements.md#user-meter-expand-phase). +See [Entitlements](./entitlements.md#usermeter-expand-phase). **Cold bootstrap:** a missing `(resource, day)` row returns `needs_bootstrap`. The service performs one legacy D1 point read on `entitlement_daily_counters`, @@ -561,7 +561,7 @@ via `durableObjectNameFromParts`); domain helpers such as the DO identity is the user. See [Run records](./run-records.md). - `UserMeter` — `userMeterDurableObjectName(userId)` → `idFromName(userId)`. One daily-entitlement meter DO per user (untrimmed stable id, same as `RunLog`). - See [Entitlements](./entitlements.md#user-meter-expand-phase). + See [Entitlements](./entitlements.md#usermeter-expand-phase). - `McpClientHub` — `mcpClientHubDurableObjectName(userId)` → `idFromName(userId.trim())`. - `StorageRunner` — `storageRunnerDurableObjectName(userId, storageId)` → @@ -628,7 +628,7 @@ Bindings are configured per environment in `packages/worker/wrangler.jsonc` - `RUN_LOG` (Durable Objects; per-user run records — see [Run records](./run-records.md)) - `USER_METER` (Durable Objects; per-user daily entitlement counters — see - [Entitlements](./entitlements.md#user-meter-expand-phase)) + [Entitlements](./entitlements.md#usermeter-expand-phase)) - `STORAGE_RUNNER` (Durable Objects) - `REPO_SESSION` (Durable Objects) - `PACKAGE_REALTIME_SESSION` (Durable Objects) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 9f8fc15d46..44c202e89b 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -109,7 +109,7 @@ two compute surfaces `usage-metering.md` already observes: Both consume only when the context has a `userId`, matching the usage-metering rule that events without an owning user are skipped. Daily consumption is authoritative in the per-user `UserMeter` Durable Object; see -[UserMeter (expand phase)](#user-meter-expand-phase). +[UserMeter (expand phase)](#usermeter-expand-phase). ## UserMeter (expand phase) diff --git a/packages/worker/src/app/handlers/account-email.node.test.ts b/packages/worker/src/app/handlers/account-email.node.test.ts index 14660553ee..fa111e18b0 100644 --- a/packages/worker/src/app/handlers/account-email.node.test.ts +++ b/packages/worker/src/app/handlers/account-email.node.test.ts @@ -298,10 +298,21 @@ function createEnv(input?: { } as Env } +function useFrozenUtcTime(iso: string) { + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(new Date(iso)) + return { + [Symbol.dispose]() { + vi.useRealTimers() + }, + } +} + test('email API lists messages with pagination, usage, and selected detail', async () => { mockModule.prepare.mockClear() mockModule.getEmailMessageById.mockClear() mockModule.listEmailDeliveryEvents.mockClear() + using _frozenTime = useFrozenUtcTime('2026-07-31T15:00:00.000Z') const day = utcDayKey() const userId = 'stable-user-1' const bootstrapMeter = createInMemoryUserMeterEnv() @@ -452,6 +463,7 @@ test('email API lists messages with pagination, usage, and selected detail', asy }) test('email API usage prefers initialized UserMeter daily counts over D1', async () => { + using _frozenTime = useFrozenUtcTime('2026-07-31T15:00:00.000Z') const day = utcDayKey() const userId = 'stable-user-1' const meter = createInMemoryUserMeterEnv() diff --git a/packages/worker/src/email/inbound-entitlements.workers.test.ts b/packages/worker/src/email/inbound-entitlements.workers.test.ts index 41757b39cd..53a9de2f46 100644 --- a/packages/worker/src/email/inbound-entitlements.workers.test.ts +++ b/packages/worker/src/email/inbound-entitlements.workers.test.ts @@ -215,6 +215,7 @@ function withFailingReceiveStartedInsert(input: { const isDeliveryInsert = query.includes('INSERT OR IGNORE INTO email_delivery_events') && query.includes("'receive_started'") + if (!isDeliveryInsert) return statement const originalBind = statement.bind.bind(statement) return { bind(...params: Array) { @@ -224,11 +225,7 @@ function withFailingReceiveStartedInsert(input: { all: bound.all.bind(bound), raw: bound.raw?.bind(bound), run: async () => { - if ( - isDeliveryInsert && - failNext && - params[0] === input.deliveryId - ) { + if (failNext && params[0] === input.deliveryId) { failNext = false throw new Error(input.message) } diff --git a/packages/worker/src/email/inbound.workers.test.ts b/packages/worker/src/email/inbound.workers.test.ts index feadc77b7b..2505321ec9 100644 --- a/packages/worker/src/email/inbound.workers.test.ts +++ b/packages/worker/src/email/inbound.workers.test.ts @@ -864,8 +864,17 @@ async function readUserDailyReceiveCount(userId: string) { } async function readUserTotalReceiveCount(userId: string) { - const exported = await userMeterRpc({ env, userId }).exportCounters({}) - return exported.counters + const meter = userMeterRpc({ env, userId }) + const counters: Array<{ resource: string; count: number }> = [] + let startAfter: string | null = null + for (;;) { + const page = await meter.exportCounters({ startAfter }) + counters.push(...page.counters) + if (!page.truncated) break + if (!page.nextStartAfter) break + startAfter = page.nextStartAfter + } + return counters .filter((row) => row.resource === 'email_receives_per_day') .reduce((total, row) => total + row.count, 0) } diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 36ac9c5213..8616078585 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -553,8 +553,10 @@ class UserMeterBase extends DurableObject { } async purge(): Promise<{ ok: true }> { - await this.ctx.storage.deleteAll() - this.initializeSchema() + await this.ctx.blockConcurrencyWhile(async () => { + await this.ctx.storage.deleteAll() + this.initializeSchema() + }) return { ok: true } } diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index a9f11debc1..d974b3fd63 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -11,10 +11,7 @@ import { planLimits } from './plans.ts' import { consumeDailyEntitlement, refundDailyEntitlement } from './service.ts' import { ensureEntitlementTestSchema } from './test-schema.ts' import { userMeterRpc } from './user-meter-client.ts' -import { - type UserMeter, - userMeterMirrorUpdatedAtToken, -} from './user-meter-do.ts' +import { UserMeter, userMeterMirrorUpdatedAtToken } from './user-meter-do.ts' import { createWaitUntilDrain, withPatchedDbPrepare, @@ -539,7 +536,20 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p ]), ) - await meterA.purge() + const [purgeResult, readDuringPurge, exportDuringPurge] = await Promise.all([ + meterA.purge(), + meterA.read({ resource: 'email_sends_per_day', day }).then( + (value) => value, + (error: unknown) => error, + ), + meterA.exportCounters({}).then( + (value) => value, + (error: unknown) => error, + ), + ]) + expect(purgeResult).toEqual({ ok: true }) + expect(readDuringPurge).not.toBeInstanceOf(Error) + expect(exportDuringPurge).not.toBeInstanceOf(Error) expect(await meterA.read({ resource: 'email_sends_per_day', day })).toEqual({ outcome: 'needs_bootstrap', }) @@ -617,3 +627,66 @@ test('UserMeter consume rejects invalid UTC day keys before SQL', async () => { ).rejects.toThrow(/UTC YYYY-MM-DD/) }) }, 30_000) + +test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore', async () => { + const now = new Date('2026-07-31T15:00:00.000Z') + const day = utcDayKey(now) + const user = await seedFreeUser('meter-purge-concurrency') + const meter = userMeterRpc({ env, userId: user.userId }) + await meter.initialize({ + resource: 'email_sends_per_day', + day, + count: 4, + updatedAt: now.toISOString(), + }) + + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(user.userId)), + ) + let releaseDelete: (() => void) | undefined + const deletePaused = new Promise((resolve) => { + releaseDelete = resolve + }) + let deleteAllReached = false + await runInDurableObject(stub, async (instance: UserMeter, state) => { + expect(instance).toBeInstanceOf(UserMeter) + const originalDeleteAll = state.storage.deleteAll.bind(state.storage) + state.storage.deleteAll = async () => { + await originalDeleteAll() + deleteAllReached = true + await deletePaused + } + }) + + const purgePromise = meter.purge() + await waitFor(() => deleteAllReached, 5_000, 'purge deleteAll') + + const readPromise = meter.read({ resource: 'email_sends_per_day', day }).then( + (value) => value, + (error: unknown) => error, + ) + const exportPromise = meter.exportCounters({}).then( + (value) => value, + (error: unknown) => error, + ) + // Give queued RPCs a chance to enter the wiped-schema window if + // blockConcurrencyWhile is missing around deleteAll+initializeSchema. + await new Promise((resolve) => setTimeout(resolve, 25)) + releaseDelete!() + + const [purgeResult, readDuringPurge, exportDuringPurge] = await Promise.all([ + purgePromise, + readPromise, + exportPromise, + ]) + expect(purgeResult).toEqual({ ok: true }) + expect(readDuringPurge).toEqual({ outcome: 'needs_bootstrap' }) + expect(exportDuringPurge).toEqual({ + counters: [], + nextStartAfter: null, + truncated: false, + }) + expect(await meter.read({ resource: 'email_sends_per_day', day })).toEqual({ + outcome: 'needs_bootstrap', + }) +}, 30_000) diff --git a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts index 2b71fd5734..4b4b6a211b 100644 --- a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts +++ b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts @@ -27,6 +27,7 @@ async function seedDailyCounter(input: { userId: string resource: 'email_sends_per_day' | 'email_receives_per_day' count: number + day: string }) { await env.APP_DB.prepare( `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) @@ -35,7 +36,7 @@ async function seedDailyCounter(input: { .bind( input.userId, input.resource, - utcDayKey(), + input.day, input.count, new Date().toISOString(), ) @@ -68,6 +69,7 @@ test( { timeout: 30_000 }, async () => { await ensureEmailTestSchema(env.APP_DB) + const day = utcDayKey() await expect( emailUsageGetCapability.handler( {}, @@ -102,11 +104,13 @@ test( userId: planUserId, resource: 'email_sends_per_day', count: 3, + day, }) await seedDailyCounter({ userId: planUserId, resource: 'email_receives_per_day', count: 5, + day, }) await seedStoredMessages(planUserId, 2) @@ -122,7 +126,7 @@ test( ) expect(planResult).toEqual({ plan: 'pro', - day: utcDayKey(), + day, stored_messages: { count: 2, limit: planLimits.pro.maxStoredEmailMessages, @@ -134,7 +138,6 @@ test( }, max_message_bytes: planLimits.pro.maxEmailMessageBytes, }) - const day = utcDayKey() const planMeter = userMeterRpc({ env, userId: planUserId }) expect( await planMeter.read({ resource: 'email_sends_per_day', day }), @@ -158,7 +161,7 @@ test( ) expect(maxResult).toEqual({ plan: 'max', - day: utcDayKey(), + day, stored_messages: { count: 0, limit: maxPlanEmailLimits.stored_email_messages, @@ -181,11 +184,13 @@ test( userId: meterUserId, resource: 'email_sends_per_day', count: 3, + day, }) await seedDailyCounter({ userId: meterUserId, resource: 'email_receives_per_day', count: 5, + day, }) await seedStoredMessages(meterUserId, 2) const updatedAt = new Date().toISOString() diff --git a/packages/worker/src/mcp/fetch-gateway.workers.test.ts b/packages/worker/src/mcp/fetch-gateway.workers.test.ts index dafd062dd9..4c8bbfcecb 100644 --- a/packages/worker/src/mcp/fetch-gateway.workers.test.ts +++ b/packages/worker/src/mcp/fetch-gateway.workers.test.ts @@ -10,6 +10,7 @@ import { UserMeter } from '#worker/entitlements/user-meter-do.ts' import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' import { silenceIncidentalRuntimeWarnings } from '#worker/test-support/incidental-runtime-warnings.ts' +import { createWaitUntilDrain } from '#worker/test-support/user-meter.ts' import { seedAccount } from '#worker/test-support/workers-seed.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' @@ -57,20 +58,17 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over } const globalFetch = (async () => new Response('ok')) as unknown as typeof fetch - const waitUntilTasks: Array> = [] - const waitUntil = (promise: Promise) => { - waitUntilTasks.push(promise) - } + const drain = createWaitUntilDrain() const allowed = await executeGatewayFetch({ env, props, request: new Request('https://api.example.net/data'), globalFetch, - waitUntil, + waitUntil: drain.waitUntil, }) expect(await allowed.text()).toBe('ok') - await Promise.all(waitUntilTasks) + await drain.drain() expect( await userMeterRpc({ env, userId }).read({ resource: 'outbound_fetches_per_day', @@ -88,7 +86,7 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over props, request: new Request('https://api.example.net/data'), globalFetch, - waitUntil, + waitUntil: drain.waitUntil, }).catch((error: unknown) => error) expect(isEntitlementLimitError(denied)).toBe(true) @@ -101,7 +99,7 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over props: { ...props, email: null }, request: new Request('https://api.example.net/data'), globalFetch, - waitUntil, + waitUntil: drain.waitUntil, }).catch((error: unknown) => error) expect(isEntitlementLimitError(deniedWithoutEmail)).toBe(true) @@ -112,7 +110,8 @@ test('gateway fetches consume the daily outbound-fetch entitlement and deny over props: { ...props, userId: null, email: null }, request: new Request('https://api.example.net/data'), globalFetch, - waitUntil, + waitUntil: drain.waitUntil, }) expect(await contextless.text()).toBe('ok') + await drain.drain() }, 30_000) From a8cba2b23ca50b89007d1c3ba36230486a02bc31 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 05:43:46 +0000 Subject: [PATCH 11/39] test(entitlements): add RunLog admin fixture Co-authored-by: Kent C. Dodds --- .../src/mcp/capabilities/admin/admin-capabilities.node.test.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts index 2532aaaf91..7fa3dabbc9 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts @@ -12,6 +12,7 @@ import { adminUserCreateCapability } from './admin-user-create.ts' import { adminUserGetCapability } from './admin-user-get.ts' import { adminUserListCapability } from './admin-user-list.ts' import { adminUserUpdateCapability } from './admin-user-update.ts' +import { createInMemoryRunLogUsageEnv } from '#worker/test-support/run-log-usage.ts' import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' import { loadAdminUserByTarget } from '#worker/admin/users-data.ts' @@ -506,6 +507,7 @@ function createAdminCapabilityTestDb(input: { } const userMeter = createInMemoryUserMeterEnv() +const runLog = createInMemoryRunLogUsageEnv() function createAdminCapabilityContext( db: D1Database, @@ -513,6 +515,7 @@ function createAdminCapabilityContext( ) { return { env: { + ...runLog.env, ...userMeter.env, APP_DB: db, EMAIL_BLOBS: blobs ?? { From edfd973adb94436d39c87df42b5b802ae3d4aaf9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 05:57:57 +0000 Subject: [PATCH 12/39] test(run-log): isolate usage counts per user Co-authored-by: Kent C. Dodds --- .../src/admin/user-usage-data.node.test.ts | 101 ++++++++---------- .../src/app/account-usage-data.node.test.ts | 3 +- .../worker/src/test-support/run-log-usage.ts | 14 ++- 3 files changed, 55 insertions(+), 63 deletions(-) diff --git a/packages/worker/src/admin/user-usage-data.node.test.ts b/packages/worker/src/admin/user-usage-data.node.test.ts index a0c549fba0..4cf563112b 100644 --- a/packages/worker/src/admin/user-usage-data.node.test.ts +++ b/packages/worker/src/admin/user-usage-data.node.test.ts @@ -6,14 +6,9 @@ import { createStableUserIdFromEmail } from '#worker/user-id.ts' import { type AdminUsageRollup } from '#app/loader-data.ts' import { loadAdminUserUsageData } from './user-usage-data.ts' -function withUserMeter( - env: { APP_DB: D1Database } & Record, - options?: { activeWorkflowCount?: number }, -) { +function withUserMeter(env: { APP_DB: D1Database } & Record) { const meter = createInMemoryUserMeterEnv() - const runLog = createInMemoryRunLogUsageEnv({ - activeWorkflowCount: options?.activeWorkflowCount ?? 0, - }) + const runLog = createInMemoryRunLogUsageEnv() return { ...env, ...meter.env, ...runLog.env, meter, runLog } } @@ -502,54 +497,52 @@ test('loadAdminUserUsageData reads daily counts from UserMeter (bootstrap then w const meterEmail = 'meter-drilldown@example.com' const meterUserId = await createStableUserIdFromEmail(meterEmail) - const warmEnv = withUserMeter( - { - APP_DB: createAdminUserUsageTestDb({ - users: [ - { - id: 4, - username: 'meter', - email: meterEmail, - plan: 'pro', - stable_user_id: meterUserId, - }, - ], - dailyCounters: [ - { - user_id: meterUserId, - resource: 'email_sends_per_day', - day, - count: 11, - }, - { - user_id: meterUserId, - resource: 'email_receives_per_day', - day, - count: 22, - }, - { - user_id: meterUserId, - resource: 'execute_calls_per_day', - day, - count: 33, - }, - { - user_id: meterUserId, - resource: 'outbound_fetches_per_day', - day, - count: 44, - }, - ], - resourceCounts: { - [meterUserId]: { - saved_packages: 6, - stored_email_messages: 9, - }, + const warmEnv = withUserMeter({ + APP_DB: createAdminUserUsageTestDb({ + users: [ + { + id: 4, + username: 'meter', + email: meterEmail, + plan: 'pro', + stable_user_id: meterUserId, }, - }), - }, - { activeWorkflowCount: 2 }, - ) + ], + dailyCounters: [ + { + user_id: meterUserId, + resource: 'email_sends_per_day', + day, + count: 11, + }, + { + user_id: meterUserId, + resource: 'email_receives_per_day', + day, + count: 22, + }, + { + user_id: meterUserId, + resource: 'execute_calls_per_day', + day, + count: 33, + }, + { + user_id: meterUserId, + resource: 'outbound_fetches_per_day', + day, + count: 44, + }, + ], + resourceCounts: { + [meterUserId]: { + saved_packages: 6, + stored_email_messages: 9, + }, + }, + }), + }) + warmEnv.runLog.setActiveWorkflowCount(meterUserId, 2) await warmEnv.meter.seed({ userId: meterUserId, resource: 'email_sends_per_day', diff --git a/packages/worker/src/app/account-usage-data.node.test.ts b/packages/worker/src/app/account-usage-data.node.test.ts index a26c0b5e5d..e672a0c141 100644 --- a/packages/worker/src/app/account-usage-data.node.test.ts +++ b/packages/worker/src/app/account-usage-data.node.test.ts @@ -178,7 +178,8 @@ test('loadAccountUsageData returns plan rows and authoritative UserMeter daily c ], }) const warmEnv = withUsageEnv({ APP_DB: warmDb }) - warmEnv.runLog.setActiveWorkflowCount(3) + warmEnv.runLog.setActiveWorkflowCount(meterUserId, 3) + warmEnv.runLog.setActiveWorkflowCount('other-user', 99) await warmEnv.meter.seed({ userId: meterUserId, resource: 'email_sends_per_day', diff --git a/packages/worker/src/test-support/run-log-usage.ts b/packages/worker/src/test-support/run-log-usage.ts index 39876fce38..2742a17b64 100644 --- a/packages/worker/src/test-support/run-log-usage.ts +++ b/packages/worker/src/test-support/run-log-usage.ts @@ -2,23 +2,21 @@ * Minimal in-memory RUN_LOG stub for node-unit usage-reader coverage of * concurrent_workflows (authoritative countActiveWorkflowProjections). */ -export function createInMemoryRunLogUsageEnv(input?: { - activeWorkflowCount?: number -}) { - let activeWorkflowCount = input?.activeWorkflowCount ?? 0 +export function createInMemoryRunLogUsageEnv() { + const activeWorkflowCounts = new Map() return { env: { RUN_LOG: { idFromName: (name: string) => name as unknown as DurableObjectId, - get: () => ({ + get: (id: DurableObjectId) => ({ countActiveWorkflowProjections: async () => ({ - count: activeWorkflowCount, + count: activeWorkflowCounts.get(String(id)) ?? 0, }), }), }, }, - setActiveWorkflowCount(count: number) { - activeWorkflowCount = Math.max(0, Math.trunc(count) || 0) + setActiveWorkflowCount(id: string, count: number) { + activeWorkflowCounts.set(id, Math.max(0, Math.trunc(count) || 0)) }, } } From b18246f79e66c0172e6b9f07a327ddb4288cdf2b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 06:56:32 +0000 Subject: [PATCH 13/39] feat(entitlements): shadow storage bytes in UserMeter Co-authored-by: Kent C. Dodds --- .../entitlements/d1-storage-reconciliation.ts | 6 + .../d1-storage-reconciliation.workers.test.ts | 114 ++++++- packages/worker/src/entitlements/service.ts | 88 ++++- .../worker/src/entitlements/user-meter-do.ts | 245 +++++++++++++- .../entitlements/user-meter.workers.test.ts | 308 +++++++++++++++++- packages/worker/src/mcp/memory/service.ts | 7 +- packages/worker/src/mcp/secrets/service.ts | 27 +- .../src/mcp/values/service.node.test.ts | 126 +++++++ packages/worker/src/mcp/values/service.ts | 6 +- .../src/package-registry/service.node.test.ts | 3 + .../worker/src/package-registry/service.ts | 3 + .../storage-runner.entitlement.node.test.ts | 49 +++ .../worker/src/storage-runner.workers.test.ts | 1 + .../worker/src/test-support/user-meter.ts | 100 +++++- 14 files changed, 1068 insertions(+), 15 deletions(-) diff --git a/packages/worker/src/entitlements/d1-storage-reconciliation.ts b/packages/worker/src/entitlements/d1-storage-reconciliation.ts index 7428ed8496..0759de7b6c 100644 --- a/packages/worker/src/entitlements/d1-storage-reconciliation.ts +++ b/packages/worker/src/entitlements/d1-storage-reconciliation.ts @@ -1,3 +1,4 @@ +import { type UserMeterEnv } from './user-meter-client.ts' import { listUsersForD1StorageReconciliation, markUserD1StorageReconciliationAttempt, @@ -9,6 +10,8 @@ import { * Reservations keep the hot path conservative between sweeps; this lane * corrects drift from deletes, failed writes after reservation, and D1 * surfaces whose writes do not pass through storage-byte enforcement. + * Optional `env` best-effort shadows each absolute into UserMeter after the + * D1 update; shadow failures never fail or poison the lane. */ export const d1StorageReconciliationBatchSize = 8 @@ -16,6 +19,8 @@ export async function reconcileD1StorageBytes(input: { db: D1Database now?: Date batchSize?: number + /** Optional expand-phase UserMeter shadow after each D1 absolute update. */ + env?: UserMeterEnv }): Promise<{ scanned: number; updated: number; failed: number }> { const now = input.now ?? new Date() const rows = await listUsersForD1StorageReconciliation({ @@ -28,6 +33,7 @@ export async function reconcileD1StorageBytes(input: { try { const result = await reconcileUserD1StorageBytes({ db: input.db, + env: input.env, userId: row.userId, now, }) diff --git a/packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts b/packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts index d9daeacd68..a690d38491 100644 --- a/packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts +++ b/packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts @@ -1,11 +1,13 @@ import { env } from 'cloudflare:test' import { expect, test } from 'vitest' import { ensureEmailTestSchema } from '#worker/email/test-schema.ts' +import { consoleWarn } from '#worker/test-support/console-spies.ts' import { calculateUserD1StorageBytes, readUserD1StorageBytes, } from './service.ts' import { reconcileD1StorageBytes } from './d1-storage-reconciliation.ts' +import { userMeterRpc } from './user-meter-client.ts' async function insertReconciliationUser(input: { userId: string @@ -29,7 +31,7 @@ async function insertReconciliationUser(input: { .run() } -test('D1 storage reconciliation is bounded and replaces stored drift with authoritative bytes', async () => { +test('D1 storage reconciliation is bounded and remains D1-authoritative with {db}', async () => { await ensureEmailTestSchema(env.APP_DB) await env.APP_DB.prepare( `UPDATE users @@ -71,6 +73,11 @@ test('D1 storage reconciliation is bounded and replaces stored drift with author }) expect(expectedFirstBytes).toBeGreaterThan(128) + const firstMeter = userMeterRpc({ env, userId: firstUserId }) + expect(await firstMeter.readStorageBytes()).toEqual({ + outcome: 'needs_bootstrap', + }) + await expect( reconcileD1StorageBytes({ db: env.APP_DB, @@ -78,9 +85,13 @@ test('D1 storage reconciliation is bounded and replaces stored drift with author batchSize: 1, }), ).resolves.toEqual({ scanned: 1, updated: 1, failed: 0 }) + await expect( readUserD1StorageBytes({ db: env.APP_DB, userId: firstUserId }), ).resolves.toBe(expectedFirstBytes) + expect(await firstMeter.readStorageBytes()).toEqual({ + outcome: 'needs_bootstrap', + }) await expect( readUserD1StorageBytes({ db: env.APP_DB, userId: secondUserId }), ).resolves.toBe(99) @@ -96,3 +107,104 @@ test('D1 storage reconciliation is bounded and replaces stored drift with author readUserD1StorageBytes({ db: env.APP_DB, userId: secondUserId }), ).resolves.toBe(0) }) + +test('optional reconciliation env shadows UserMeter without becoming authoritative', async () => { + await ensureEmailTestSchema(env.APP_DB) + const userId = '3'.repeat(64) + await insertReconciliationUser({ + userId, + email: `${crypto.randomUUID()}@example.test`, + storedBytes: 50, + updatedAt: '2019-01-01T00:00:00.000Z', + }) + await env.APP_DB.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, from_address, subject, text_body, raw_size, + processing_status, created_at, updated_at + ) VALUES (?, 'inbound', ?, 'sender@example.test', 'subject', 'body', 200, + 'stored', ?, ?)`, + ) + .bind( + crypto.randomUUID(), + userId, + '2026-07-31T00:00:00.000Z', + '2026-07-31T00:00:00.000Z', + ) + .run() + + const meter = userMeterRpc({ env, userId }) + await meter.setStorageBytes({ + bytes: 999_999, + updatedAt: '2026-07-31T00:00:00.000Z', + }) + + const expectedBytes = await calculateUserD1StorageBytes({ + db: env.APP_DB, + userId, + }) + await expect( + reconcileD1StorageBytes({ + db: env.APP_DB, + env, + now: new Date('2026-07-31T01:00:00.000Z'), + batchSize: 8, + }), + ).resolves.toMatchObject({ failed: 0 }) + + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId }), + ).resolves.toBe(expectedBytes) + expect(await meter.readStorageBytes()).toMatchObject({ + outcome: 'ready', + bytes: expectedBytes, + }) +}) + +test('reconciliation shadow failure does not fail or poison the D1 lane', async () => { + await ensureEmailTestSchema(env.APP_DB) + await env.APP_DB.prepare( + `UPDATE users + SET d1_storage_bytes_updated_at = '9999-12-31T23:59:59.999Z'`, + ).run() + const userId = '4'.repeat(64) + await insertReconciliationUser({ + userId, + email: `${crypto.randomUUID()}@example.test`, + storedBytes: 40, + updatedAt: '2018-01-01T00:00:00.000Z', + }) + + consoleWarn.mockImplementation(() => {}) + const failingMeterEnv = { + USER_METER: { + idFromName: env.USER_METER.idFromName.bind(env.USER_METER), + get() { + return { + async setStorageBytes() { + throw new Error('shadow write failed') + }, + } + }, + }, + } as unknown as typeof env + + const expectedBytes = await calculateUserD1StorageBytes({ + db: env.APP_DB, + userId, + }) + await expect( + reconcileD1StorageBytes({ + db: env.APP_DB, + env: failingMeterEnv, + now: new Date('2026-07-31T01:00:00.000Z'), + batchSize: 1, + }), + ).resolves.toEqual({ scanned: 1, updated: 1, failed: 0 }) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId }), + ).resolves.toBe(expectedBytes) + expect(consoleWarn).toHaveBeenCalledWith( + 'entitlement-storage-bytes-shadow-failed', + expect.any(Error), + ) +}) diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index 93d8e72435..ac3d13183b 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -13,7 +13,11 @@ import { isDailyEntitlementResource, type DailyEntitlementResource, } from './user-meter-do.ts' -import { userMeterRpc, type UserMeterEnv } from './user-meter-client.ts' +import { + userMeterNamespace, + userMeterRpc, + type UserMeterEnv, +} from './user-meter-client.ts' /** Env surface for authoritative entitlement usage readers (UserMeter + RunLog). */ export type EntitlementUsageEnv = UserMeterEnv & Pick @@ -682,10 +686,49 @@ export async function calculateUserD1StorageBytes(input: { return sums.reduce((total, value) => total + value, 0) } -export async function readUserD1StorageBytes(input: { +/** Best-effort D1→UserMeter absolute shadow; re-reads D1 at task time. */ +async function shadowUserMeterStorageBytesFromD1(input: { db: D1Database + env: UserMeterEnv userId: string }) { + if (!userMeterNamespace(input.env)) return + const bytes = await readUserD1StorageBytes({ + db: input.db, + userId: input.userId, + }) + const meter = userMeterRpc({ env: input.env, userId: input.userId }) + await meter.setStorageBytes({ + bytes, + updatedAt: new Date().toISOString(), + }) +} + +function scheduleUserMeterStorageBytesShadow(input: { + db: D1Database + env: UserMeterEnv | undefined + userId: string + waitUntil?: (promise: Promise) => void +}) { + if (!input.env || !userMeterNamespace(input.env)) return + const tracked = shadowUserMeterStorageBytesFromD1({ + db: input.db, + env: input.env, + userId: input.userId, + }).catch((error: unknown) => { + console.warn('entitlement-storage-bytes-shadow-failed', error) + }) + if (input.waitUntil) { + input.waitUntil(tracked) + return + } + void tracked +} + +export async function readUserD1StorageBytes(input: { + db: D1Database + userId: string +}): Promise { const row = await input.db .prepare( `SELECT d1_storage_bytes AS bytes @@ -697,10 +740,17 @@ export async function readUserD1StorageBytes(input: { return Math.max(0, Number(row?.bytes ?? 0)) } +/** + * Replace one user's D1 storage-byte counter with the cross-surface D1 payload + * sum. Optional `env` best-effort shadows into UserMeter after the D1 write; + * shadow failures never fail or poison this reconcile. + */ export async function reconcileUserD1StorageBytes(input: { db: D1Database userId: string now?: Date + /** Optional expand-phase UserMeter shadow after the D1 absolute update. */ + env?: UserMeterEnv }) { const bytes = await calculateUserD1StorageBytes(input) const result = await input.db @@ -712,7 +762,19 @@ export async function reconcileUserD1StorageBytes(input: { ) .bind(bytes, (input.now ?? new Date()).toISOString(), input.userId) .run() - return { bytes, updated: (result.meta.changes ?? 0) > 0 } + const updated = (result.meta.changes ?? 0) > 0 + if (input.env) { + try { + await shadowUserMeterStorageBytesFromD1({ + db: input.db, + env: input.env, + userId: input.userId, + }) + } catch (error) { + console.warn('entitlement-storage-bytes-shadow-failed', error) + } + } + return { bytes, updated } } export async function listUsersForD1StorageReconciliation(input: { @@ -923,12 +985,22 @@ export async function readCurrentEntitlementResourceUsage(input: { }) } +/** + * Atomically reserve D1 payload storage bytes against the plan limit. + * Optional `env` schedules a best-effort UserMeter shadow after a successful + * reserve (never awaited; failures logged only). `getCurrent` is check-only + * and does not reserve in D1 or UserMeter. + */ export async function assertWithinStorageBytesEntitlement(input: { db: D1Database userId: string email: string | null | undefined requested?: number getCurrent?: () => Promise + /** Optional expand-phase UserMeter shadow after a successful D1 reserve. */ + env?: UserMeterEnv + /** Prefer `ctx.waitUntil` for the UserMeter shadow; never awaited here. */ + waitUntil?: (promise: Promise) => void }) { if (input.getCurrent) { await assertWithinEntitlement({ @@ -959,7 +1031,15 @@ export async function assertWithinStorageBytesEntitlement(input: { ) .bind(requested, now, input.userId, requested, limit) .run() - if ((result.meta.changes ?? 0) > 0) return + if ((result.meta.changes ?? 0) > 0) { + scheduleUserMeterStorageBytesShadow({ + db: input.db, + env: input.env, + userId: input.userId, + waitUntil: input.waitUntil, + }) + return + } const current = await readUserD1StorageBytes({ db: input.db, diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 8616078585..adc12de48d 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -26,7 +26,9 @@ export const userMeterDailyCounterRetentionDays = 7 const metaSchemaVersionKey = 'schema_version' /** Bump when initializeSchema DDL changes; warm objects skip DDL. */ -const userMeterSchemaVersion = 3 +const userMeterSchemaVersion = 4 +/** Singleton row id for expand-phase D1 storage-byte shadow (schema v4). */ +const storageBytesStateRowId = 1 const defaultExportPageSize = 100 const maxExportPageSize = 500 @@ -77,8 +79,41 @@ export type UserMeterInitializeResult = UserMeterReadyState & { created: boolean } +export type UserMeterStorageBytesState = { + bytes: number + revision: number + updatedAt: string + mirrorUpdatedAt: string +} + +export type UserMeterStorageBytesReadyState = { + outcome: 'ready' + bytes: number + revision: number + mirrorUpdatedAt: string +} + +export type UserMeterStorageBytesReadResult = + | UserMeterBootstrapState + | UserMeterStorageBytesReadyState + +export type UserMeterStorageBytesReserveResult = + | UserMeterBootstrapState + | (UserMeterStorageBytesReadyState & { reserved: boolean }) + +export type UserMeterStorageBytesInitializeResult = + UserMeterStorageBytesReadyState & { + created: boolean + } + +export type UserMeterStorageBytesSetResult = UserMeterStorageBytesReadyState & { + created: boolean +} + export type UserMeterExportResult = { counters: Array + /** Non-authoritative D1 storage-byte shadow; `null` when never synced. */ + storageBytesShadow: UserMeterStorageBytesState | null nextStartAfter: string | null truncated: boolean } @@ -253,6 +288,15 @@ class UserMeterBase extends DurableObject { `CREATE INDEX IF NOT EXISTS idx_inbound_delivery_claims_day ON inbound_delivery_claims (day)`, ) + // Expand-phase shadow of D1 `users.d1_storage_bytes` (schema v4). + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS storage_bytes_state ( + id INTEGER PRIMARY KEY NOT NULL CHECK (id = 1), + bytes INTEGER NOT NULL, + revision INTEGER NOT NULL, + updated_at TEXT NOT NULL + ) + `) this.ctx.storage.sql.exec( `INSERT INTO user_meter_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, @@ -261,6 +305,41 @@ class UserMeterBase extends DurableObject { ) } + private storageReadyState( + bytes: number, + revision: number, + ): UserMeterStorageBytesReadyState { + const safeBytes = Math.max(0, bytes) + const safeRevision = Math.max(0, revision) + return { + outcome: 'ready', + bytes: safeBytes, + revision: safeRevision, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(safeRevision), + } + } + + private readStorageRow(): { + bytes: number + revision: number + updatedAt: string + } | null { + const row = this.ctx.storage.sql + .exec<{ bytes: number; revision: number; updated_at: string }>( + `SELECT bytes, revision, updated_at + FROM storage_bytes_state + WHERE id = ?`, + storageBytesStateRowId, + ) + .toArray()[0] + if (!row) return null + return { + bytes: Math.max(0, Number(row.bytes ?? 0)), + revision: Math.max(0, Number(row.revision ?? 0)), + updatedAt: String(row.updated_at ?? ''), + } + } + private deleteStaleCounters(now: Date) { const cutoffDay = retentionCutoffDay(now) this.ctx.storage.sql.exec( @@ -552,6 +631,143 @@ class UserMeterBase extends DurableObject { return readyState(row.count, row.revision) } + /** Cutover-support cold seed; expand-phase shadow sync uses {@link setStorageBytes}. */ + async initializeStorageBytes(input: { + bytes: number + updatedAt: string + }): Promise { + const bytes = Math.max(0, Math.trunc(Number(input.bytes) || 0)) + const cursor = this.ctx.storage.sql.exec( + `INSERT INTO storage_bytes_state (id, bytes, revision, updated_at) + VALUES (?, ?, 1, ?) + ON CONFLICT(id) DO NOTHING`, + storageBytesStateRowId, + bytes, + input.updatedAt, + ) + const created = cursor.rowsWritten > 0 + const row = this.readStorageRow() + if (!row) { + throw new Error( + 'UserMeter initializeStorageBytes failed to materialize storage state.', + ) + } + return { ...this.storageReadyState(row.bytes, row.revision), created } + } + + /** Cutover-support / shadow read; expand-phase usage and enforcement use D1. */ + async readStorageBytes(): Promise { + const row = this.readStorageRow() + if (!row) return { outcome: 'needs_bootstrap' } + return this.storageReadyState(row.bytes, row.revision) + } + + /** + * Cutover-support atomic reserve. Expand-phase enforcement uses D1 only; + * missing state returns `needs_bootstrap`. + */ + async reserveStorageBytes(input: { + requested: number + limit: number + updatedAt: string + }): Promise { + const requested = Math.max(0, Math.trunc(Number(input.requested) || 0)) + const existing = this.readStorageRow() + if (!existing) { + return { outcome: 'needs_bootstrap' } + } + if ( + requested > 0 && + (input.limit < 1 || existing.bytes + requested > input.limit) + ) { + return { + ...this.storageReadyState(existing.bytes, existing.revision), + reserved: false, + } + } + if (requested === 0) { + return { + ...this.storageReadyState(existing.bytes, existing.revision), + reserved: true, + } + } + const nextBytes = existing.bytes + requested + const nextRevision = existing.revision + 1 + this.ctx.storage.sql.exec( + `UPDATE storage_bytes_state + SET bytes = ?, + revision = ?, + updated_at = ? + WHERE id = ? AND revision = ?`, + nextBytes, + nextRevision, + input.updatedAt, + storageBytesStateRowId, + existing.revision, + ) + const row = this.readStorageRow() + if (!row) { + throw new Error('UserMeter reserveStorageBytes lost storage state.') + } + return { + ...this.storageReadyState(row.bytes, row.revision), + reserved: row.bytes === nextBytes && row.revision === nextRevision, + } + } + + /** + * Absolute shadow-set from D1. Materializes the singleton and bumps + * revision. Callers should re-read latest D1 before invoking so delayed + * shadows cannot leave the DO behind D1. + */ + async setStorageBytes(input: { + bytes: number + updatedAt: string + }): Promise { + const bytes = Math.max(0, Math.trunc(Number(input.bytes) || 0)) + const existing = this.readStorageRow() + if (!existing) { + this.ctx.storage.sql.exec( + `INSERT INTO storage_bytes_state (id, bytes, revision, updated_at) + VALUES (?, ?, 1, ?)`, + storageBytesStateRowId, + bytes, + input.updatedAt, + ) + const row = this.readStorageRow() + if (!row) { + throw new Error( + 'UserMeter setStorageBytes failed to materialize storage state.', + ) + } + return { + ...this.storageReadyState(row.bytes, row.revision), + created: true, + } + } + const nextRevision = existing.revision + 1 + this.ctx.storage.sql.exec( + `UPDATE storage_bytes_state + SET bytes = ?, + revision = ?, + updated_at = ? + WHERE id = ? AND revision = ?`, + bytes, + nextRevision, + input.updatedAt, + storageBytesStateRowId, + existing.revision, + ) + const row = this.readStorageRow() + if (!row) { + throw new Error('UserMeter setStorageBytes lost storage state.') + } + return { + ...this.storageReadyState(row.bytes, row.revision), + created: false, + } + } + async purge(): Promise<{ ok: true }> { await this.ctx.blockConcurrencyWhile(async () => { await this.ctx.storage.deleteAll() @@ -622,9 +838,18 @@ class UserMeterBase extends DurableObject { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(revision), }) } + const storageRow = this.readStorageRow() const last = pageRows[pageRows.length - 1] return { counters, + storageBytesShadow: storageRow + ? { + bytes: storageRow.bytes, + revision: storageRow.revision, + updatedAt: storageRow.updatedAt, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(storageRow.revision), + } + : null, nextStartAfter: truncated && last ? encodeExportCursor({ @@ -672,6 +897,24 @@ export type UserMeterRpc = { day: string updatedAt: string }) => Promise + /** Cutover-support cold seed; expand-phase shadow uses setStorageBytes. */ + initializeStorageBytes: (input: { + bytes: number + updatedAt: string + }) => Promise + /** Cutover-support / shadow read; expand-phase usage reads D1. */ + readStorageBytes: () => Promise + /** Cutover-support reserve; expand-phase enforcement uses D1. */ + reserveStorageBytes: (input: { + requested: number + limit: number + updatedAt: string + }) => Promise + /** Expand-phase absolute shadow-set from D1. */ + setStorageBytes: (input: { + bytes: number + updatedAt: string + }) => Promise purge: () => Promise<{ ok: true }> exportCounters: (input: { pageSize?: number diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index d974b3fd63..3722d4b156 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -8,7 +8,12 @@ import { createStableUserIdFromEmail } from '#worker/user-id.ts' import { userMeterDurableObjectName } from '#worker/user-scoped-durable-object-name.ts' import { EntitlementLimitError } from './errors.ts' import { planLimits } from './plans.ts' -import { consumeDailyEntitlement, refundDailyEntitlement } from './service.ts' +import { + assertWithinStorageBytesEntitlement, + consumeDailyEntitlement, + readUserD1StorageBytes, + refundDailyEntitlement, +} from './service.ts' import { ensureEntitlementTestSchema } from './test-schema.ts' import { userMeterRpc } from './user-meter-client.ts' import { UserMeter, userMeterMirrorUpdatedAtToken } from './user-meter-do.ts' @@ -555,6 +560,7 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p }) expect(await meterA.exportCounters({})).toEqual({ counters: [], + storageBytesShadow: null, nextStartAfter: null, truncated: false, }) @@ -683,6 +689,7 @@ test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore expect(readDuringPurge).toEqual({ outcome: 'needs_bootstrap' }) expect(exportDuringPurge).toEqual({ counters: [], + storageBytesShadow: null, nextStartAfter: null, truncated: false, }) @@ -690,3 +697,302 @@ test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore outcome: 'needs_bootstrap', }) }, 30_000) + +test('storage bytes enforcement remains D1-authoritative even when env is passed', async () => { + const user = await seedFreeUser('meter-storage-d1-authority') + const drain = createWaitUntilDrain() + const meter = userMeterRpc({ env, userId: user.userId }) + const storageLimit = planLimits.free.maxStorageBytes + const legacyBytes = storageLimit - 10 + + await env.APP_DB.prepare( + `UPDATE users + SET d1_storage_bytes = ?, + d1_storage_bytes_updated_at = ? + WHERE stable_user_id = ?`, + ) + .bind(legacyBytes, '2026-07-31T12:00:00.000Z', user.userId) + .run() + + await meter.setStorageBytes({ + bytes: storageLimit, + updatedAt: '2026-07-31T11:00:00.000Z', + }) + + await assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + requested: 5, + waitUntil: drain.waitUntil, + }) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(legacyBytes + 5) + + const denied = await assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + requested: 6, + waitUntil: drain.waitUntil, + }).then( + () => null, + (thrown: unknown) => thrown, + ) + expect(denied).toBeInstanceOf(EntitlementLimitError) + expect(denied).toMatchObject({ + details: { + resource: 'storage_bytes', + plan: 'free', + limit: storageLimit, + current: legacyBytes + 5, + }, + }) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(legacyBytes + 5) + + await drain.drain() + expect(await meter.readStorageBytes()).toMatchObject({ + outcome: 'ready', + bytes: legacyBytes + 5, + }) +}, 30_000) + +test('UserMeter storage shadow failure cannot affect successful D1 reserves', async () => { + const user = await seedFreeUser('meter-storage-shadow-fail') + const drain = createWaitUntilDrain() + const storageLimit = planLimits.free.maxStorageBytes + + await env.APP_DB.prepare( + `UPDATE users + SET d1_storage_bytes = ?, + d1_storage_bytes_updated_at = ? + WHERE stable_user_id = ?`, + ) + .bind(storageLimit - 20, '2026-07-31T12:00:00.000Z', user.userId) + .run() + + consoleWarn.mockImplementation(() => {}) + const failingMeterEnv = { + USER_METER: { + idFromName: env.USER_METER.idFromName.bind(env.USER_METER), + get() { + return { + async setStorageBytes() { + throw new Error('shadow write failed') + }, + } + }, + }, + } as unknown as typeof env + + await expect( + assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env: failingMeterEnv, + userId: user.userId, + email: user.email, + requested: 5, + waitUntil: drain.waitUntil, + }), + ).resolves.toBeUndefined() + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(storageLimit - 15) + await drain.drain() + expect(consoleWarn).toHaveBeenCalledWith( + 'entitlement-storage-bytes-shadow-failed', + expect.any(Error), + ) +}, 30_000) + +test('concurrent D1 storage reserves cannot overshoot the plan limit', async () => { + const user = await seedFreeUser('meter-storage-concurrent-d1') + const limit = planLimits.free.maxStorageBytes + await env.APP_DB.prepare( + `UPDATE users + SET d1_storage_bytes = ?, + d1_storage_bytes_updated_at = ? + WHERE stable_user_id = ?`, + ) + .bind(limit - 10, '2026-07-31T15:00:00.000Z', user.userId) + .run() + + const attempts = await Promise.all( + Array.from({ length: 20 }, () => + assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + userId: user.userId, + email: user.email, + requested: 5, + }).then( + () => 'reserved' as const, + (error: unknown) => error, + ), + ), + ) + const reserved = attempts.filter((result) => result === 'reserved') + const denied = attempts.filter( + (result) => result instanceof EntitlementLimitError, + ) + expect(reserved).toHaveLength(2) + expect(denied).toHaveLength(18) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(limit) +}, 30_000) + +test('delayed UserMeter storage shadow re-reads latest D1 and cannot regress D1', async () => { + const user = await seedFreeUser('meter-storage-shadow-latest') + const drain = createWaitUntilDrain() + const meter = userMeterRpc({ env, userId: user.userId }) + + await env.APP_DB.prepare( + `UPDATE users + SET d1_storage_bytes = ?, + d1_storage_bytes_updated_at = ? + WHERE stable_user_id = ?`, + ) + .bind(10, '2026-07-31T16:00:00.000Z', user.userId) + .run() + + let releaseShadow!: () => void + const shadowGate = new Promise((resolve) => { + releaseShadow = resolve + }) + let shadowReads = 0 + using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + const statement = originalPrepare(query) + const isPointRead = + query.includes('SELECT d1_storage_bytes AS bytes') && + query.includes('FROM users') + const isReserve = query.includes( + 'SET d1_storage_bytes = d1_storage_bytes + ?', + ) + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: async () => { + if (isPointRead) { + shadowReads += 1 + await shadowGate + } + return await bound.first() + }, + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: bound.run.bind(bound), + } + }, + } + }) as D1Database['prepare'] + }) + + await assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + requested: 5, + waitUntil: drain.waitUntil, + }) + await assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + userId: user.userId, + email: user.email, + requested: 7, + }) + releaseShadow() + await drain.drain() + expect(shadowReads).toBeGreaterThanOrEqual(1) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(22) + expect(await meter.readStorageBytes()).toMatchObject({ + outcome: 'ready', + bytes: 22, + }) +}, 30_000) + +test('UserMeter storage RPCs, export shadow field, and purge work additively', async () => { + const user = await seedFreeUser('meter-storage-export-purge') + const meter = userMeterRpc({ env, userId: user.userId }) + await meter.initializeStorageBytes({ + bytes: 42, + updatedAt: '2026-07-31T17:00:00.000Z', + }) + await meter.reserveStorageBytes({ + requested: 8, + limit: 1_000, + updatedAt: '2026-07-31T17:01:00.000Z', + }) + await meter.setStorageBytes({ + bytes: 11, + updatedAt: '2026-07-31T17:02:00.000Z', + }) + + const exported = await meter.exportCounters({}) + expect(exported.storageBytesShadow).toEqual({ + bytes: 11, + revision: 3, + updatedAt: '2026-07-31T17:02:00.000Z', + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(3), + }) + await expect( + readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), + ).resolves.toBe(0) + + await expect(meter.purge()).resolves.toEqual({ ok: true }) + expect(await meter.readStorageBytes()).toEqual({ + outcome: 'needs_bootstrap', + }) + expect(await meter.exportCounters({})).toEqual({ + counters: [], + storageBytesShadow: null, + nextStartAfter: null, + truncated: false, + }) +}, 30_000) + +test('missing-user storage reserve preserves EntitlementLimitError semantics on D1 path', async () => { + await ensureEntitlementTestSchema(env.APP_DB) + const missingUserId = 'a'.repeat(64) + const freeLimit = planLimits.free.maxStorageBytes + + await expect( + assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env, + userId: missingUserId, + email: null, + requested: 1, + }), + ).resolves.toBeUndefined() + + const denied = await assertWithinStorageBytesEntitlement({ + db: env.APP_DB, + env, + userId: missingUserId, + email: null, + requested: freeLimit + 1, + }).then( + () => null, + (thrown: unknown) => thrown, + ) + expect(denied).toBeInstanceOf(EntitlementLimitError) + expect(denied).toMatchObject({ + details: { + resource: 'storage_bytes', + plan: 'free', + limit: freeLimit, + current: 0, + }, + }) +}, 30_000) diff --git a/packages/worker/src/mcp/memory/service.ts b/packages/worker/src/mcp/memory/service.ts index 993b24b0c9..846afba466 100644 --- a/packages/worker/src/mcp/memory/service.ts +++ b/packages/worker/src/mcp/memory/service.ts @@ -21,6 +21,7 @@ import { assertWithinStorageBytesEntitlement, estimateEntitlementStorageEntryByteDelta, } from '#worker/entitlements/service.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' import { type McpMemoryRow, type MemoryRecord, @@ -72,7 +73,8 @@ type MemoryOwnerContext = { } type MemoryEnv = Pick & - Partial> + Partial> & + UserMeterEnv type MemoryUpsertInput = MemoryOwnerContext & { env: MemoryEnv @@ -86,6 +88,7 @@ type MemoryUpsertInput = MemoryOwnerContext & { dedupeKey?: string | null status?: 'active' | 'archived' verificationReference?: string | null + waitUntil?: (promise: Promise) => void } type MemoryDeleteInput = MemoryOwnerContext & { @@ -215,6 +218,7 @@ export async function upsertMemory(input: MemoryUpsertInput): Promise<{ await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: estimateEntitlementStorageEntryByteDelta({ @@ -247,6 +251,7 @@ export async function upsertMemory(input: MemoryUpsertInput): Promise<{ } : null, }), + waitUntil: input.waitUntil, }) if (existing) { diff --git a/packages/worker/src/mcp/secrets/service.ts b/packages/worker/src/mcp/secrets/service.ts index e57cd35455..8440cf7ed1 100644 --- a/packages/worker/src/mcp/secrets/service.ts +++ b/packages/worker/src/mcp/secrets/service.ts @@ -39,6 +39,7 @@ import { assertWithinStorageBytesEntitlement, estimateEntitlementStorageEntryByteDelta, } from '#worker/entitlements/service.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' import { type SecretMetadata, type SecretScope } from './types.ts' type SecretOwnerContext = { @@ -46,14 +47,17 @@ type SecretOwnerContext = { storageContext?: StorageContext | null } +type SecretWriteEnv = Pick & UserMeterEnv + type SaveSecretInput = SecretOwnerContext & { - env: Pick + env: SecretWriteEnv scope: SecretScope name: string value: string description?: string | null sessionExpiresAt?: string | null userEmail?: string | null + waitUntil?: (promise: Promise) => void } type ListSecretsInput = SecretOwnerContext & { @@ -126,6 +130,7 @@ export async function saveSecret( const encryptedValue = await encryptSecretValue(input.env, value) await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: estimateEntitlementStorageEntryByteDelta({ @@ -152,6 +157,7 @@ export async function saveSecret( } : null, }), + waitUntil: input.waitUntil, }) const now = new Date().toISOString() await upsertSecretEntry({ @@ -189,13 +195,14 @@ export async function saveSecret( } export async function updateUserSecretForPackage(input: { - env: Pick + env: SecretWriteEnv userId: string userEmail?: string | null packageId: string name: string value: string description?: string | null + waitUntil?: (promise: Promise) => void }) { const [saved] = await updateUserSecretsForPackageAtomically({ env: input.env, @@ -209,6 +216,7 @@ export async function updateUserSecretForPackage(input: { description: input.description, }, ], + waitUntil: input.waitUntil, }) if (!saved) { throw new Error( @@ -226,7 +234,7 @@ export async function updateUserSecretForPackage(input: { * fail-closes in SQL if any secret lacks the grant. */ export async function setSecretsAtomically(input: { - env: Pick + env: SecretWriteEnv userId: string userEmail?: string | null secrets: Array<{ @@ -236,6 +244,7 @@ export async function setSecretsAtomically(input: { description?: string | null }> storageContext?: StorageContext | null + waitUntil?: (promise: Promise) => void }): Promise> { if (input.secrets.length === 0) { throw new Error('At least one secret is required.') @@ -279,6 +288,7 @@ export async function setSecretsAtomically(input: { value: secret.value, description: secret.description, })), + waitUntil: input.waitUntil, }) } @@ -293,11 +303,12 @@ export async function setSecretsAtomically(input: { description: secret.description, })), storageContext, + waitUntil: input.waitUntil, }) } export async function updateUserSecretsForPackageAtomically(input: { - env: Pick + env: SecretWriteEnv userId: string userEmail?: string | null packageId: string @@ -306,6 +317,7 @@ export async function updateUserSecretsForPackageAtomically(input: { value: string description?: string | null }> + waitUntil?: (promise: Promise) => void }): Promise> { const packageId = input.packageId.trim() if (!packageId) throw new Error('Package id is required.') @@ -371,9 +383,11 @@ export async function updateUserSecretsForPackageAtomically(input: { await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: requestedStorageBytes, + waitUntil: input.waitUntil, }) try { @@ -420,7 +434,7 @@ export async function updateUserSecretsForPackageAtomically(input: { } async function saveSecretsAtomically(input: { - env: Pick + env: SecretWriteEnv userId: string userEmail?: string | null scope: SecretScope @@ -430,6 +444,7 @@ async function saveSecretsAtomically(input: { description?: string | null }> storageContext?: StorageContext | null + waitUntil?: (promise: Promise) => void }): Promise> { const bucket = await getOrCreateSecretBucket({ db: input.env.APP_DB, @@ -502,9 +517,11 @@ async function saveSecretsAtomically(input: { } await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: requestedStorageBytes, + waitUntil: input.waitUntil, }) await upsertSecretEntriesAtomically({ diff --git a/packages/worker/src/mcp/values/service.node.test.ts b/packages/worker/src/mcp/values/service.node.test.ts index db0223efa3..fdf7dd1724 100644 --- a/packages/worker/src/mcp/values/service.node.test.ts +++ b/packages/worker/src/mcp/values/service.node.test.ts @@ -6,6 +6,13 @@ import { resolveStorageScopeOrder, } from '#mcp/storage-bindings.ts' import { deleteAllAppScopedValues } from '#worker/package-config-cleanup.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' +import { consoleWarn } from '#worker/test-support/console-spies.ts' +import { + createInMemoryUserMeterEnv, + createWaitUntilDrain, + withPatchedDbPrepare, +} from '#worker/test-support/user-meter.ts' import { deleteValue, getValue, listValues, saveValue } from './service.ts' import { type ValueBucketRow, type ValueEntryRow } from './types.ts' @@ -663,3 +670,122 @@ test('saveValue permits underscore-prefixed ordinary value names', async () => { value: '{"provider":"widgets"}', }) }) + +test('saveValue awaits atomic D1 storage reserve; D1→UserMeter shadow is deferred and non-blocking', async () => { + const testDb = createValueTestDb() + const meter = createInMemoryUserMeterEnv() + const drain = createWaitUntilDrain() + const userId = 'a'.repeat(64) + const email = 'value-storage@example.com' + let d1StorageBytes = 100 + let failShadowPointRead = false + await meter.seedStorageBytes({ userId, bytes: 7 }) + + let releaseShadowRead!: () => void + const shadowReadGate = new Promise((resolve) => { + releaseShadowRead = resolve + }) + let shadowPointReadStarted = false + let atomicReserveCompleted = false + using _patch = withPatchedDbPrepare(testDb.db, (originalPrepare) => { + return ((query: string) => { + const normalized = query.replace(/\s+/g, ' ').trim().toLowerCase() + const statement = originalPrepare(query) + const originalBind = statement.bind.bind(statement) + return { + bind(...params: Array) { + const bound = originalBind(...params) + return { + first: async () => { + if ( + normalized.includes('select plan, stripe_plan from users') && + normalized.includes('stable_user_id') + ) { + return { plan: 'pro', stripe_plan: null } as T + } + if ( + normalized.includes('select d1_storage_bytes as bytes') && + normalized.includes('from users') + ) { + if (failShadowPointRead) { + throw new Error('forced D1→UserMeter shadow read failure') + } + shadowPointReadStarted = true + await shadowReadGate + return { bytes: d1StorageBytes } as T + } + return await bound.first() + }, + all: bound.all.bind(bound), + raw: bound.raw?.bind(bound), + run: async () => { + if ( + normalized.includes( + 'set d1_storage_bytes = d1_storage_bytes + ?', + ) + ) { + const requested = Number(params[0]) + const limit = Number(params[4]) + if (d1StorageBytes + requested <= limit) { + d1StorageBytes += requested + atomicReserveCompleted = true + return { meta: { changes: 1 } } + } + return { meta: { changes: 0 } } + } + return await bound.run() + }, + } + }, + } + }) as D1Database['prepare'] + }) + + const env = { APP_DB: testDb.db, ...meter.env } + const saved = await saveValue({ + env, + userId, + userEmail: email, + scope: 'user', + name: 'metered-value', + value: 'hello-storage', + waitUntil: drain.waitUntil, + }) + expect(saved).toMatchObject({ name: 'metered-value' }) + expect(atomicReserveCompleted).toBe(true) + expect(d1StorageBytes).toBeGreaterThan(100) + expect(shadowPointReadStarted).toBe(true) + expect(await userMeterRpc({ env, userId }).readStorageBytes()).toMatchObject({ + outcome: 'ready', + bytes: 7, + }) + + releaseShadowRead() + await drain.drain() + expect(await userMeterRpc({ env, userId }).readStorageBytes()).toMatchObject({ + outcome: 'ready', + bytes: d1StorageBytes, + }) + + failShadowPointRead = true + consoleWarn.mockImplementation(() => {}) + const failingDrain = createWaitUntilDrain() + const bytesBeforeFailingSave = d1StorageBytes + await expect( + saveValue({ + env, + userId, + userEmail: email, + scope: 'user', + name: 'metered-value-2', + value: 'still-ok', + waitUntil: failingDrain.waitUntil, + }), + ).resolves.toMatchObject({ name: 'metered-value-2' }) + expect(d1StorageBytes).toBeGreaterThan(bytesBeforeFailingSave) + await failingDrain.drain() + expect(consoleWarn).toHaveBeenCalledWith( + 'entitlement-storage-bytes-shadow-failed', + expect.any(Error), + ) +}) diff --git a/packages/worker/src/mcp/values/service.ts b/packages/worker/src/mcp/values/service.ts index e2becb8bfc..8dc3d81b49 100644 --- a/packages/worker/src/mcp/values/service.ts +++ b/packages/worker/src/mcp/values/service.ts @@ -21,6 +21,7 @@ import { assertWithinStorageBytesEntitlement, estimateEntitlementStorageEntryByteDelta, } from '#worker/entitlements/service.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' type ValueOwnerContext = { userId: string @@ -28,13 +29,14 @@ type ValueOwnerContext = { } type SaveValueInput = ValueOwnerContext & { - env: Pick + env: Pick & UserMeterEnv scope: ValueScope name: string value: string description?: string | null sessionExpiresAt?: string | null userEmail?: string | null + waitUntil?: (promise: Promise) => void } type ListValuesInput = ValueOwnerContext & { @@ -87,6 +89,7 @@ export async function saveValue(input: SaveValueInput): Promise { }) await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: estimateEntitlementStorageEntryByteDelta({ @@ -104,6 +107,7 @@ export async function saveValue(input: SaveValueInput): Promise { } : null, }), + waitUntil: input.waitUntil, }) const now = new Date().toISOString() await upsertValueEntry({ diff --git a/packages/worker/src/package-registry/service.node.test.ts b/packages/worker/src/package-registry/service.node.test.ts index 14b86f6dc6..83ba2c744e 100644 --- a/packages/worker/src/package-registry/service.node.test.ts +++ b/packages/worker/src/package-registry/service.node.test.ts @@ -1095,6 +1095,9 @@ function createEntitlementsDatabase(input: { const user = users.find((row) => row.email === params[0]) return (user ? { plan: user.plan } : null) as T | null } + if (query.includes('SELECT d1_storage_bytes AS bytes')) { + return null + } if ( query.includes('SELECT COUNT(*) AS count FROM saved_packages') ) { diff --git a/packages/worker/src/package-registry/service.ts b/packages/worker/src/package-registry/service.ts index c0e4790e48..673946208f 100644 --- a/packages/worker/src/package-registry/service.ts +++ b/packages/worker/src/package-registry/service.ts @@ -217,6 +217,7 @@ export async function refreshSavedPackageProjection(input: { packageId: string sourceId: string rebuildArtifacts?: boolean + waitUntil?: (promise: Promise) => void }) { return await withAccountWriteLease({ db: input.env.APP_DB, @@ -248,6 +249,7 @@ export async function refreshSavedPackageProjection(input: { }) await assertWithinStorageBytesEntitlement({ db: input.env.APP_DB, + env: input.env, userId: input.userId, email: input.userEmail, requested: estimateEntitlementStorageEntryByteDelta({ @@ -276,6 +278,7 @@ export async function refreshSavedPackageProjection(input: { } : null, }), + waitUntil: input.waitUntil, }) if (existing) { await updateSavedPackage(input.env.APP_DB, { diff --git a/packages/worker/src/storage-runner.entitlement.node.test.ts b/packages/worker/src/storage-runner.entitlement.node.test.ts index 5b591fea55..d7c51bb55a 100644 --- a/packages/worker/src/storage-runner.entitlement.node.test.ts +++ b/packages/worker/src/storage-runner.entitlement.node.test.ts @@ -229,3 +229,52 @@ test('an unreachable peer bucket with a stored estimate never blocks a write', a // Only the write target was measured live; no peer fan-out happened. expect(probedStorageIds).toEqual(['package:target']) }) + +test('StorageRunner baseline composes D1 payload with bucket estimates via check-only getCurrent', async () => { + const { planLimits } = await import('#worker/entitlements/plans.ts') + const { isEntitlementLimitError } = + await import('#worker/entitlements/errors.ts') + const userId = 'user-1' + const limit = planLimits.free.maxStorageBytes + mockModule.readUserD1StorageBytes.mockImplementation( + async (input: { userId: string }) => { + expect(input.userId).toBe(userId) + return limit - 100 + }, + ) + mockModule.listUserStorageBucketEstimates.mockResolvedValue([ + { storageId: 'package:peer', estimatedBytes: 50 }, + { storageId: 'package:target', estimatedBytes: 200 }, + ]) + const probedStorageIds: Array = [] + const getEstimatedBytes = async (storageId: string) => { + probedStorageIds.push(storageId) + return { estimatedBytes: storageId === 'package:target' ? 60 : 50 } + } + + const denied = await assertStorageRunnerWriteWithinEntitlement({ + env: createEstimateEnv(getEstimatedBytes), + userId, + email: null, + storageId: 'package:target', + requested: 1, + }).then( + () => null, + (thrown: unknown) => thrown, + ) + expect(isEntitlementLimitError(denied)).toBe(true) + expect(denied).toMatchObject({ + details: { + resource: 'storage_bytes', + current: limit - 100 + 50 + 60, + limit, + }, + }) + expect(probedStorageIds).toEqual(['package:target']) + expect(mockModule.readUserD1StorageBytes).toHaveBeenCalledWith( + expect.objectContaining({ + userId, + db: expect.anything(), + }), + ) +}) diff --git a/packages/worker/src/storage-runner.workers.test.ts b/packages/worker/src/storage-runner.workers.test.ts index 5f9859b213..c606da2f18 100644 --- a/packages/worker/src/storage-runner.workers.test.ts +++ b/packages/worker/src/storage-runner.workers.test.ts @@ -285,6 +285,7 @@ test('storage runner storage byte entitlement aggregates only inventoried user b .bind(targetD1Bytes, new Date().toISOString(), userId) .run() expect(initialD1Bytes).toBe(0) + // D1 remains the sole payload authority for composed getCurrent. await expect( readUserD1StorageBytes({ db: env.APP_DB, userId }), ).resolves.toBe(targetD1Bytes) diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index bf3a390f23..c02fe7f3b5 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -2,10 +2,12 @@ import { isDailyEntitlementResource, userMeterMirrorUpdatedAtToken, type DailyEntitlementResource, + type UserMeterStorageBytesState, } from '#worker/entitlements/user-meter-do.ts' import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' type MeterRow = { count: number; revision: number } +type StorageRow = { bytes: number; revision: number; updatedAt: string } /** * In-memory UserMeter stub keyed by `idFromName` (stable userId) for node-unit @@ -14,6 +16,7 @@ type MeterRow = { count: number; revision: number } */ export function createInMemoryUserMeterEnv() { const metersByUser = new Map>() + const storageByUser = new Map() function counterKey(resource: string, day: string) { return `${resource}\0${day}` @@ -37,6 +40,15 @@ export function createInMemoryUserMeterEnv() { } } + function storageReady(row: StorageRow) { + return { + outcome: 'ready' as const, + bytes: row.bytes, + revision: row.revision, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), + } + } + return { async initialize(input: { resource: string @@ -105,8 +117,69 @@ export function createInMemoryUserMeterEnv() { rows.set(counterKey(input.resource, input.day), next) return ready(next) }, + async initializeStorageBytes(input: { + bytes: number + updatedAt: string + }) { + const existing = storageByUser.get(userId) + if (existing) return { ...storageReady(existing), created: false } + const row = { + bytes: Math.max(0, Math.trunc(input.bytes) || 0), + revision: 1, + updatedAt: input.updatedAt, + } + storageByUser.set(userId, row) + return { ...storageReady(row), created: true } + }, + async readStorageBytes() { + const existing = storageByUser.get(userId) + if (!existing) return { outcome: 'needs_bootstrap' as const } + return storageReady(existing) + }, + async reserveStorageBytes(input: { + requested: number + limit: number + updatedAt: string + }) { + const requested = Math.max(0, Math.trunc(input.requested) || 0) + const existing = storageByUser.get(userId) + if (!existing) return { outcome: 'needs_bootstrap' as const } + if ( + requested > 0 && + (input.limit < 1 || existing.bytes + requested > input.limit) + ) { + return { ...storageReady(existing), reserved: false } + } + if (requested === 0) { + return { ...storageReady(existing), reserved: true } + } + const next = { + bytes: existing.bytes + requested, + revision: existing.revision + 1, + updatedAt: input.updatedAt, + } + storageByUser.set(userId, next) + return { ...storageReady(next), reserved: true } + }, + async setStorageBytes(input: { bytes: number; updatedAt: string }) { + const bytes = Math.max(0, Math.trunc(input.bytes) || 0) + const existing = storageByUser.get(userId) + if (!existing) { + const row = { bytes, revision: 1, updatedAt: input.updatedAt } + storageByUser.set(userId, row) + return { ...storageReady(row), created: true } + } + const next = { + bytes, + revision: existing.revision + 1, + updatedAt: input.updatedAt, + } + storageByUser.set(userId, next) + return { ...storageReady(next), created: false } + }, async purge() { rows.clear() + storageByUser.delete(userId) return { ok: true as const } }, async exportCounters() { @@ -121,7 +194,21 @@ export function createInMemoryUserMeterEnv() { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), } }) - return { counters, nextStartAfter: null, truncated: false } + const storage = storageByUser.get(userId) + const storageBytesShadow: UserMeterStorageBytesState | null = storage + ? { + bytes: storage.bytes, + revision: storage.revision, + updatedAt: storage.updatedAt, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(storage.revision), + } + : null + return { + counters, + storageBytesShadow, + nextStartAfter: null, + truncated: false, + } }, } } @@ -136,6 +223,7 @@ export function createInMemoryUserMeterEnv() { return { env, metersByUser, + storageByUser, async seed(input: { userId: string resource: DailyEntitlementResource @@ -149,6 +237,16 @@ export function createInMemoryUserMeterEnv() { updatedAt: new Date().toISOString(), }) }, + async seedStorageBytes(input: { + userId: string + bytes: number + updatedAt?: string + }) { + await meterFor(input.userId).initializeStorageBytes({ + bytes: input.bytes, + updatedAt: input.updatedAt ?? new Date().toISOString(), + }) + }, } } From eaaa50906609ef9abf196ff61a49603d97ac5140 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 06:56:37 +0000 Subject: [PATCH 14/39] docs(entitlements): inventory storage shadow Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 104 ++++---- .../contributing/architecture/entitlements.md | 222 ++++++++++++------ .../contributing/architecture/primitives.yaml | 4 +- .../worker/src/account/export.node.test.ts | 14 +- packages/worker/src/account/export.ts | 11 +- .../worker/src/account/user-owned-surfaces.ts | 2 +- .../src/app/account-usage-data.node.test.ts | 29 ++- 7 files changed, 263 insertions(+), 123 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 504de9bbdc..8e74c6138e 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -166,13 +166,14 @@ migration-safe chunked interface: with `section: "storage_runner"` and a `storage_id`, using the same StorageRunner `exportStorage({ pageSize, startAfter })` RPC as the dedicated storage export capability. User meter counters use `section: "user_meter"` and - the `UserMeter.exportCounters` RPC. R2 raw MIME, attachment, avatar, and icon - objects use `section: "r2_object"`; each response contains at most one 256 KiB - base64 chunk and an opaque cursor. Each request uses bounded `LIMIT 1` - ownership queries rather than reconstructing inventory. Continuation cursors - bind the source row, object key, size, and ETag; ownership/key mutations and - object overwrites are reported instead of mixing generations. Missing objects - are represented explicitly. + the `UserMeter.exportCounters` RPC (daily counters plus additive + `storageBytesShadow` when present; explicitly non-authoritative). R2 raw MIME, + attachment, avatar, and icon objects use `section: "r2_object"`; each response + contains at most one 256 KiB base64 chunk and an opaque cursor. Each request + uses bounded `LIMIT 1` ownership queries rather than reconstructing inventory. + Continuation cursors bind the source row, object key, size, and ETag; + ownership/key mutations and object overwrites are reported instead of mixing + generations. Missing objects are represented explicitly. D1 manifest counts use bounded SQL `COUNT(*)` queries. D1 section rows are read with SQL-level keyset pagination: every query orders by the table's `rowid`, @@ -196,8 +197,12 @@ Durable Object export behavior: never pruned by retention. See [Run records](./run-records.md). - `UserMeter` exports daily entitlement counter rows through the `user_meter` section (`exportCounters` RPC; keyset pagination by UTC `day` and `resource`). - Retention is self-enforced inside the DO (seven UTC days of counter and - inbound-delivery-claim rows); see + The same RPC may return additive `storageBytesShadow` when the schema-v4 + shadow row exists (`null` when never shadowed); section totals count it as one + row when present, but it is explicitly **non-authoritative** — usage and + enforcement read D1 `users.d1_storage_bytes`. Retention is self-enforced + inside the DO (seven UTC days of counter and inbound-delivery-claim rows); + shadow storage-byte state is not time-pruned. See [Entitlements](./entitlements.md#usermeter-expand-phase). - `RemoteConnectorSession` exposes persisted connector metadata and tool descriptors through an export RPC. @@ -238,11 +243,15 @@ The schema is defined by migrations in `packages/worker/migrations/`: default `public`) come from migration 0068. `account_type` (`'person'` default or `'platform'`, migration 0072) distinguishes normal signups from operator-provisioned platform accounts that own official package scopes (see - [Platform accounts](./platform-accounts.md)). Inbound email routing does not - reverse-resolve stable ids at all — it uses the indexed username lookup - (`findPublicUserIdentityByUsername`). Contextless paths resolve stable ids - with one indexed point read on `users.stable_user_id` (for example - `findUserAccountByStableUserId`). + [Platform accounts](./platform-accounts.md)). `d1_storage_bytes` and + `d1_storage_bytes_updated_at` (migration 0122) are the **sole authority** for + D1 payload storage-byte read, enforcement, and reconciliation. UserMeter + `storage_bytes_state` (schema v4) is an optional expand-phase shadow only — + see [Entitlements](./entitlements.md#usermeter-expand-phase). Inbound email + routing does not reverse-resolve stable ids at all — it uses the indexed + username lookup (`findPublicUserIdentityByUsername`). Contextless paths + resolve stable ids with one indexed point read on `users.stable_user_id` (for + example `findUserAccountByStableUserId`). - `platform_feedback`: attributed, user-approved Kody feedback and admin triage state. Submitter identity remains on the row; optional reviewer attribution is cleared if that admin account is deleted. Open and triaged rows remain until @@ -489,16 +498,20 @@ Storage split: Daily rate-style entitlement counters and inbound email delivery-id idempotency live in a per-user `UserMeter` Durable Object with SQLite -(`packages/worker/src/entitlements/user-meter-do.ts`). The Worker binding is -`USER_METER` (class `UserMeter`; Wrangler SQLite migration tag `v21` via -`new_sqlite_classes` in `packages/worker/wrangler.jsonc`). +(`packages/worker/src/entitlements/user-meter-do.ts`). Schema v4 adds an +optional `storage_bytes_state` singleton as a **best-effort shadow** of D1 +`users.d1_storage_bytes` for future cutover — D1 remains sole authority for +reads, reserves, and reconciliation in the current additive slice. The Worker +binding is `USER_METER` (class `UserMeter`; Wrangler SQLite migration tag `v21` +via `new_sqlite_classes` in `packages/worker/wrangler.jsonc`). Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP `userId` via `userMeterDurableObjectName(userId)` → `idFromName(userId)` in `packages/worker/src/user-scoped-durable-object-name.ts`. There is no `user_id` column inside the DO because the object identity is the user. -SQLite ownership (schema version tracked in `user_meter_meta`): +SQLite ownership (schema version tracked in `user_meter_meta`; current version +**4**): - `daily_counters` — authoritative UTC-day counters for `email_sends_per_day`, `email_receives_per_day`, `execute_calls_per_day`, and @@ -509,30 +522,40 @@ SQLite ownership (schema version tracked in `user_meter_meta`): claim's resource/day, post-charge counter, revision, and `claimed_at` so Cloudflare Email Routing retries inside the 48-hour inbound dedupe window cannot double-charge `email_receives_per_day`. +- `storage_bytes_state` — singleton **shadow** of D1 payload bytes (`id = 1` + CHECK constraint, `bytes`, monotonic `revision`, `updated_at`). Added in + schema v4. Populated by optional non-awaited shadow writes after D1 reserves + and optional reconcile shadows; never read for enforcement or usage display. + StorageRunner bucket estimates stay outside this row (see + [Entitlements](./entitlements.md#usermeter-expand-phase)). Retention is self-enforced inside the DO: every read/write path opportunistically deletes counter and claim rows older than seven UTC days (`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; the window covers timezone edge cases, recent account exports, and inbound -retries. - -**Expand-phase D1 mirror:** enforcement and point reads are authoritative in -UserMeter. D1 `entitlement_daily_counters` is **not** dropped — it remains a -best-effort mirror for existing readers and reporting. After each DO -consume/refund/inbound claim, the entitlements service schedules a non-awaited -absolute mirror write keyed by `(user_id, resource, day)` with a -revision-ordered `updated_at` token (`r/` + zero-padded revision from -`userMeterMirrorUpdatedAtToken`) so late writes cannot overwrite newer state. -See [Entitlements](./entitlements.md#usermeter-expand-phase). - -**Cold bootstrap:** a missing `(resource, day)` row returns `needs_bootstrap`. -The service performs one legacy D1 point read on `entitlement_daily_counters`, -then `initialize()` seeds the DO row with `INSERT OR IGNORE` (concurrent callers -cannot double-apply the baseline). Warm paths never read D1 for enforcement. - -Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan). -Account export pages `UserMeter.exportCounters` through the `user_meter` -manifest section / `account_export_section`. +retries. Shadow storage-byte state is not time-pruned. + +**Expand-phase D1 mirrors (daily counters only):** enforcement and point reads +are authoritative in UserMeter for daily counters. D1 +`entitlement_daily_counters` is **not** dropped — it remains a best-effort +mirror for existing readers and reporting. After each DO consume/refund/inbound +claim, the entitlements service schedules a non-awaited absolute mirror write +keyed by `(user_id, resource, day)` with a revision-ordered `updated_at` token +(`r/` + zero-padded revision from `userMeterMirrorUpdatedAtToken`) so late +writes cannot overwrite newer state. See +[Entitlements](./entitlements.md#usermeter-expand-phase). + +**Daily cold bootstrap:** a missing `(resource, day)` row returns +`needs_bootstrap`. The service performs one legacy D1 point read on +`entitlement_daily_counters`, then `initialize()` seeds the DO row with +`INSERT OR IGNORE` (concurrent callers cannot double-apply the baseline). Warm +daily paths never read D1 for enforcement. + +Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan; +`deleteAll` clears counters, claims, and any shadow storage state). Account +export pages `UserMeter.exportCounters` through the `user_meter` manifest +section / `account_export_section` (daily counters plus additive +`storageBytesShadow` when present; shadow field is non-authoritative). ### Package state model @@ -570,8 +593,9 @@ via `durableObjectNameFromParts`); domain helpers such as and dedicated unpruned state (workflow projections, job-run observability, package activation counters/milestones). See [Run records](./run-records.md). - `UserMeter` — `userMeterDurableObjectName(userId)` → `idFromName(userId)`. One - daily-entitlement meter DO per user (untrimmed stable id, same as `RunLog`). - See [Entitlements](./entitlements.md#usermeter-expand-phase). + daily-entitlement meter DO per user (untrimmed stable id, same as `RunLog`), + plus optional schema-v4 D1 storage-byte shadow. See + [Entitlements](./entitlements.md#usermeter-expand-phase). - `McpClientHub` — `mcpClientHubDurableObjectName(userId)` → `idFromName(userId.trim())`. - `StorageRunner` — `storageRunnerDurableObjectName(userId, storageId)` → @@ -886,7 +910,7 @@ to `durableObjectNameFromParts`). - `JobManager`: `idFromName(userId)` (no trim). - `RunLog`: `idFromName(userId)` (no trim); one execution-history DO per user. - `UserMeter`: `idFromName(userId)` (no trim); one daily-entitlement meter DO - per user. + per user, plus optional schema-v4 D1 storage-byte shadow. - `McpClientHub`: `idFromName(userId.trim())`. - `StorageRunner`: `idFromName(JSON.stringify([userId, storageId]))`. - `RepoSession`: `idFromName(repo_sessions.id)`; the key is not user-prefixed, diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index c4c983ba17..2879002bd5 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -119,6 +119,17 @@ per-user `UserMeter` Durable Object** (`USER_METER` binding). Code lives in `packages/worker/src/entitlements/user-meter-do.ts` and `user-meter-client.ts`; storage layout and naming are documented in [Data storage](./data-storage.md). +**D1 payload storage bytes** (`storage_bytes`) stay authoritative on D1 +`users.d1_storage_bytes` (migration 0122). UserMeter schema v4 adds an optional +`storage_bytes_state` singleton as a **best-effort shadow** for future cutover — +it does not drive reads, reserves, or reconciliation in this additive slice. + +StorageRunner bucket `estimatedBytes` and the per-bucket inventory in +`user_storage_buckets` stay a **separate** quota component. StorageRunner write +chokepoints pass `getCurrent` as a check-only composed total (D1 payload bytes +from `readUserD1StorageBytes` plus bucket estimates); that path does **not** +reserve bytes in D1 or UserMeter. + **Strong enforcement:** `consumeDailyEntitlement` and inbound `consumeInboundDelivery` RPCs check the plan limit and increment inside the DO. The Durable Object request model serializes mutations per user; counter updates @@ -167,15 +178,71 @@ mirror is scheduled from the email path via `scheduleAbsoluteDailyEntitlementMirror` so the email subsystem does not duplicate mirror SQL. -**Remaining migration dependency:** retiring the D1 mirror and -`entitlement_daily_counters` table waits until reporting-off-D1 work merges and -mirror parity is verified in production. This expand phase intentionally does -**not** move `storage_bytes` enforcement, `package_service_states` running -counts, or account-deletion write fencing to UserMeter. +### D1 payload storage bytes — UserMeter shadow (expand phase slice 3) + +D1 `users.d1_storage_bytes` and `users.d1_storage_bytes_updated_at` remain the +**sole read, enforcement, and reconciliation authority** for the D1 payload +component of `storage_bytes`. Every caller — including email — reserves through +the same atomic D1 `UPDATE` in `assertWithinStorageBytesEntitlement`; denial +semantics and `EntitlementLimitError` always come from D1. + +UserMeter schema **v4** adds a singleton `storage_bytes_state` row (`id = 1`) as +**best-effort shadow / future-cutover support only**. The DO row is never read +for usage display or enforcement in this slice. + +**Optional non-email shadow writes:** DO-backed write chokepoints (values, +secrets, memories, and saved-package projections) may pass optional `env` into +`assertWithinStorageBytesEntitlement`. After a **successful D1 reserve**, the +service schedules a non-awaited best-effort shadow via `waitUntil` when +available: re-read the latest D1 value, then absolute-set UserMeter through +`setStorageBytes`. Shadow tasks re-read D1 at execution time so a delayed shadow +cannot apply a stale absolute and leave the DO behind D1. Missing `USER_METER` +or shadow failures log `entitlement-storage-bytes-shadow-failed` and never +affect D1 enforcement. + +**Optional reconcile shadow:** `reconcileUserD1StorageBytes` absolute-sets D1 +from `calculateUserD1StorageBytes` and awaits that write. When optional `env` is +present, it best-effort shadows the post-update D1 value into UserMeter; shadow +failures are logged and never fail or poison the lane. Current +`d1_storage_reconciliation` wiring passes `{ db }` only. + +**Usage reads stay on D1:** account usage, admin drill-down, and +`readCurrentEntitlementResourceUsage` for `storage_bytes` call +`readUserD1StorageBytes` only. + +**Account export and purge:** `UserMeter.exportCounters` may return additive +non-authoritative `storageBytesShadow` when the shadow row exists (counted in +the `user_meter` section total). `UserMeter.purge()` clears counters, inbound +delivery claims, and any shadow storage state via `deleteAll`. + +### Future storage authority flip (contract follow-up) + +A separate contract PR — not a merge blocker for this additive slice — will flip +D1 payload storage-byte authority into UserMeter once: + +1. **mailbox-do** passes `env` on email inbound/outbound storage reservations, + and +2. **cron-restructure** passes `env` into `d1_storage_reconciliation` lane + wiring in `index.ts`. + +Only then do reads, reserves, and reconciliation switch to UserMeter-first with +D1 as mirror/cursor. Until that flip, shadow divergence is acceptable; D1 +remains the contract. + +**Remaining UserMeter expand roadmap:** slice 4 — `package_service_states` +running counts (service liveness); slice 5 — account-deletion write fencing. +Storage authority flip is tracked separately as the contract follow-up above. + +**Daily-counter mirror retirement:** dropping D1 `entitlement_daily_counters` +waits until reporting-off-D1 work merges and mirror parity is verified in +production. Module wiring: `consumeDailyEntitlement`, `refundDailyEntitlement`, and -`readDailyEntitlementResourceUsage` in `service.ts` require `env.USER_METER` and -fail closed when the binding is missing. +`readDailyEntitlementResourceUsage` require `env.USER_METER` and fail closed +when the binding is missing. Storage-byte helpers (`readUserD1StorageBytes`, +`assertWithinStorageBytesEntitlement`, `reconcileUserD1StorageBytes`) treat +`USER_METER` as optional shadow-only; missing binding is a no-op for shadow +scheduling. ## Schema history @@ -351,60 +418,65 @@ Rules: against the limit instead of an accumulating count: the enforcement point passes the candidate size via `getCurrent` with `requested: 0`. There is no built-in counter for these. -- **Storage-byte limits** (`storage_bytes`) store the D1 payload estimate on - `users.d1_storage_bytes` (migration 0122). Entitlement checks read that - indexed user row and D1 write chokepoints atomically reserve their existing - positive byte-delta estimate before the write, so mailbox growth and other hot - writes never rescan the user's accumulated rows. The estimate covers - user-owned rows with durable payloads (`email_messages.raw_size` plus - extracted message bodies/metadata, externally stored attachments, values, - encrypted secrets, memories, saved-package projections, jobs, repo/session - metadata, package invocation results, and published artifact metadata). Run - records in the per-user `RunLog` Durable Object are intentionally **excluded** - from the quota — they are observability history, not user content. A fixed - batch of eight oldest users is reconciled every five-minute cron tick by - `d1_storage_reconciliation`; each selected counter is replaced with the - authoritative cross-surface sum. This bounded lane corrects over-reservation - after failed writes, shrinkage from deletes, and writes on surfaces without a - storage-byte chokepoint. Failed rows retain their conservative value and - rotate to the back of the queue for a later retry. StorageRunner Durable - Object buckets expose their own `estimatedBytes`, and each bucket's latest - measurement is persisted on its `user_storage_buckets.estimated_bytes` - inventory row (migration 0118). Write chokepoints pass `getCurrent` as - `D1 counter + sum of per-bucket estimates`, where only the bucket that - triggers the baseline read (plus any inventoried bucket with no stored - estimate yet) is probed live; every other bucket contributes its stored D1 - estimate, so mutating writes do not fan `getEstimatedBytes` RPCs across the - user's whole bucket inventory. Live probe results are persisted - fire-and-forget with **UPDATE-only** statements (they can never recreate an - inventory row removed by account, package, or job deletion), and mutating - StorageRunner RPCs opportunistically refresh their own bucket's stored - estimate after the write, throttled per bucket per isolate - (`storageBucketEstimateRefreshMinIntervalMs`). Stored estimates are therefore - freshness hints with bounded lag — acceptable for an order-of-magnitude cap - because the bucket paying the baseline read is measured live and the run cache - accounts for the run's own accepted writes. `requested` is the candidate - payload size when known. Pure read-only `storage.sql` / `packageStorage().sql` - statements (`SELECT` / `EXPLAIN` / schema `PRAGMA`) skip the baseline read - entirely even when the helper marks the call writable. Mutating SQL and - `storage.set` in one sandbox share a per-run baseline cache so repeated writes - do not re-read the baseline; a later write in the same run that targets a - **different** already-inventoried bucket reuses that bucket's stored estimate - rather than probing it live (bounded staleness, same trade-off as peers). Each - live estimate read waits at most ~2s via `Promise.race` and is retried with - backoff (`storageEstimateReadRetryDelaysMs`; a single 150ms retry lost to - transient per-bucket DO read failures in production) before failing closed for - the caller; the underlying DO RPC is not cancelled if the runtime keeps it - running. A cron lane (`storage_bucket_estimate_backfill`, - `packages/worker/src/storage-buckets/estimate-backfill.ts`) sweeps inventory - rows without a stored estimate in bounded batches (failed probes stay - unmeasured and are retried on later sweeps) so freshly migrated inventories - converge to stored estimates within a few ticks instead of making each user's - first mutating write pay (and possibly fail on) the whole-inventory probe. The - counter intentionally does **not** attempt to scan Cloudflare Artifacts - repository contents, KV snapshot/bundle bodies, R2 object listings beyond - `email_messages.raw_size`, or Vectorize: those stores either lack reliable - byte metadata or are derived from D1 and are documented in `data-storage.md`. +- **Storage-byte limits** (`storage_bytes`) split into two quota components: + 1. **D1 payload bytes (authoritative on D1):** user-owned D1 rows with durable + payloads (`email_messages.raw_size` plus extracted message bodies/metadata, + externally stored attachments, values, encrypted secrets, memories, + saved-package projections, jobs, repo/session metadata, package invocation + results, and published artifact metadata). Run records in the per-user + `RunLog` Durable Object are intentionally **excluded** — they are + observability history, not user content. Write chokepoints atomically + reserve positive byte deltas via `assertWithinStorageBytesEntitlement` + against `users.d1_storage_bytes` (every caller, including email). Optional + `env` on non-email paths schedules a best-effort UserMeter shadow after a + successful D1 reserve; shadow state is not read for enforcement or usage. + The bounded `d1_storage_reconciliation` lane recomputes the authoritative + cross-surface sum via `calculateUserD1StorageBytes` and absolute-sets D1; + optional `env` shadows into UserMeter after the D1 write — see + [UserMeter (expand phase)](#usermeter-expand-phase). + + 2. **StorageRunner bucket estimates (separate):** each bucket exposes + `estimatedBytes`; inventory rows persist the latest measurement on + `user_storage_buckets.estimated_bytes` (migration 0118). Write chokepoints + that pass `getCurrent` compose + `D1 payload (readUserD1StorageBytes) + sum of per-bucket estimates` for a + **check-only** entitlement comparison — `getCurrent` never reserves in D1 + or UserMeter. Only the bucket that triggers the baseline read (plus any + inventoried bucket with no stored estimate yet) is probed live; every other + bucket contributes its stored D1 estimate, so mutating writes do not fan + `getEstimatedBytes` RPCs across the whole inventory. Live probe results are + persisted fire-and-forget with **UPDATE-only** statements (they can never + recreate an inventory row removed by account, package, or job deletion), + and mutating StorageRunner RPCs opportunistically refresh their own + bucket's stored estimate after the write, throttled per bucket per isolate + (`storageBucketEstimateRefreshMinIntervalMs`). Stored estimates are + freshness hints with bounded lag — acceptable for an order-of-magnitude cap + because the bucket paying the baseline read is measured live and the run + cache accounts for the run's own accepted writes. `requested` is the + candidate payload size when known. Pure read-only `storage.sql` / + `packageStorage().sql` statements (`SELECT` / `EXPLAIN` / schema `PRAGMA`) + skip the baseline read entirely even when the helper marks the call + writable. Mutating SQL and `storage.set` in one sandbox share a per-run + baseline cache so repeated writes do not re-read the baseline; a later + write in the same run that targets a **different** already-inventoried + bucket reuses that bucket's stored estimate rather than probing it live + (bounded staleness, same trade-off as peers). Each live estimate read waits + at most ~2s via `Promise.race` and is retried with backoff + (`storageEstimateReadRetryDelaysMs`; a single 150ms retry lost to transient + per-bucket DO read failures in production) before failing closed for the + caller; the underlying DO RPC is not cancelled if the runtime keeps it + running. A cron lane (`storage_bucket_estimate_backfill`, + `packages/worker/src/storage-buckets/estimate-backfill.ts`) sweeps + inventory rows without a stored estimate in bounded batches (failed probes + stay unmeasured and are retried on later sweeps) so freshly migrated + inventories converge to stored estimates within a few ticks instead of + making each user's first mutating write pay (and possibly fail on) the + whole-inventory probe. The D1 payload counter intentionally does **not** + attempt to scan Cloudflare Artifacts repository contents, KV + snapshot/bundle bodies, R2 object listings beyond + `email_messages.raw_size`, or Vectorize: those stores either lack reliable + byte metadata or are derived from D1 and are documented in + `data-storage.md`. ### Concurrency @@ -458,20 +530,20 @@ Use `getCurrent` only when the built-in D1 counter cannot express the resource. ## Enforcement points -| Resource | Enforcement point | -| ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------ | -| `scheduled_jobs` | `createJob` in `packages/worker/src/jobs/service.ts` (exemplar) | -| `saved_packages` | new-package branch of `package_save` and projection insert | -| `package_services` | `service_start` capability path (count from `package_service_states`) | -| `persistent_package_services` | `service_start` for services declared `mode: 'persistent'` | -| `repo_sessions` | `repo_open_session` before creating a new session | -| `email_sends_per_day` | `sendOutboundEmail` (`consumeDailyEntitlement`; plan limit from `resolvePlanLimit`) | -| `email_receives_per_day` | `handleInboundEmail` (`consumeDailyEntitlement`; same plan limits; refund only on `RetryableInboundStorageError`) | -| `stored_email_messages` | `handleInboundEmail` before storage (`assertWithinEntitlement`; `max` caps from `planLimits.max`) | -| `email_message_bytes` | `handleInboundEmail` before quota/parse (per-message raw size via `resolveEmailResourceLimit`) | -| `secrets` | new-entry branch of `saveSecret` in `packages/worker/src/mcp/secrets/service.ts` | -| `concurrent_workflows` | `createDynamicCallableWorkflow` (`reserveWorkflowProjectionSlot` + `assertWithinEntitlement` getCurrent; `max` = 5,000) | -| `storage_bytes` | D1 payload writes (values, secrets, memories, saved-package projections, email storage) and StorageRunner write tools/app RPCs | +| Resource | Enforcement point | +| ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `scheduled_jobs` | `createJob` in `packages/worker/src/jobs/service.ts` (exemplar) | +| `saved_packages` | new-package branch of `package_save` and projection insert | +| `package_services` | `service_start` capability path (count from `package_service_states`) | +| `persistent_package_services` | `service_start` for services declared `mode: 'persistent'` | +| `repo_sessions` | `repo_open_session` before creating a new session | +| `email_sends_per_day` | `sendOutboundEmail` (`consumeDailyEntitlement`; plan limit from `resolvePlanLimit`) | +| `email_receives_per_day` | `handleInboundEmail` (`consumeDailyEntitlement`; same plan limits; refund only on `RetryableInboundStorageError`) | +| `stored_email_messages` | `handleInboundEmail` before storage (`assertWithinEntitlement`; `max` caps from `planLimits.max`) | +| `email_message_bytes` | `handleInboundEmail` before quota/parse (per-message raw size via `resolveEmailResourceLimit`) | +| `secrets` | new-entry branch of `saveSecret` in `packages/worker/src/mcp/secrets/service.ts` | +| `concurrent_workflows` | `createDynamicCallableWorkflow` (`reserveWorkflowProjectionSlot` + `assertWithinEntitlement` getCurrent; `max` = 5,000) | +| `storage_bytes` | D1 payload writes via `assertWithinStorageBytesEntitlement` (D1 reserve for all callers; optional `env` shadow on non-email paths) and StorageRunner write tools/app RPCs (`getCurrent` check-only for bucket component) | ## Billing diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index a86960099b..5302603cd6 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -513,8 +513,8 @@ primitives: group: storage name: User meter summary: - Per-user UserMeter DO SQLite for daily entitlement counters and inbound - delivery idempotency. + Per-user UserMeter DO SQLite for daily entitlement counters, inbound + delivery idempotency, and expand-phase D1 storage-byte shadow. code: - packages/worker/src/entitlements/user-meter-do.ts - packages/worker/src/entitlements/user-meter-client.ts diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 08e9964d21..97a962c81c 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1572,6 +1572,12 @@ test('account export includes user_meter counters, pages them, and warns on trun mirrorUpdatedAt: 'r/00000000000000000001', }, ] + const storageBytesShadow = { + bytes: 4_096, + revision: 3, + updatedAt: '2026-07-31T03:00:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000003', + } const exportCounters = vi.fn( async (input: { pageSize?: number; startAfter?: string | null }) => { const pageSize = input.pageSize ?? 100 @@ -1584,6 +1590,7 @@ test('account export includes user_meter counters, pages them, and warns on trun const truncated = startIndex + pageSize < counters.length return { counters: page, + storageBytesShadow, nextStartAfter: truncated ? `${page.at(-1)!.day}:${page.at(-1)!.resource}` : null, @@ -1606,9 +1613,10 @@ test('account export includes user_meter counters, pages them, and warns on trun mcpUserId: 'user-aaa', }) expect(idFromName).toHaveBeenCalledWith('user-aaa') - expect(accountExport.manifest.sections.user_meter?.count).toBe(3) + expect(accountExport.manifest.sections.user_meter?.count).toBe(4) expect(accountExport.durableObjects.userMeter).toEqual({ counters, + storageBytesShadow, nextStartAfter: null, truncated: false, }) @@ -1626,6 +1634,7 @@ test('account export includes user_meter counters, pages them, and warns on trun pageSize: 2, }) expect(first.items).toEqual(counters.slice(0, 2)) + expect(first.storageBytesShadow).toEqual(storageBytesShadow) expect(first.truncated).toBe(true) expect(first.nextStartAfter).toBe('2026-07-30:execute_calls_per_day') @@ -1638,12 +1647,14 @@ test('account export includes user_meter counters, pages them, and warns on trun startAfter: first.nextStartAfter ?? undefined, }) expect(second.items).toEqual(counters.slice(2)) + expect(second.storageBytesShadow).toEqual(storageBytesShadow) expect(second.truncated).toBe(false) expect(second.nextStartAfter).toBeNull() expect(exportCounters).toHaveBeenCalled() exportCounters.mockImplementation(async () => ({ counters: [counters[0]!], + storageBytesShadow: null, nextStartAfter: 'cursor-more', truncated: true, })) @@ -1653,6 +1664,7 @@ test('account export includes user_meter counters, pages them, and warns on trun mcpUserId: 'user-aaa', }) expect(truncatedExport.durableObjects.userMeter?.truncated).toBe(true) + expect(truncatedExport.manifest.sections.user_meter?.count).toBe(1) expect(truncatedExport.manifest.warnings).toContain( 'User meter counters were truncated in the full export; use account_export_section with section "user_meter" to retrieve additional pages.', ) diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index c3bf88a67e..297abc2c38 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -204,6 +204,10 @@ export type AccountExportArtifactRepo = { type RunRecordsExportPayload = Awaited> +function countUserMeterExportEntries(result: UserMeterExportResult): number { + return result.counters.length + (result.storageBytesShadow == null ? 0 : 1) +} + function countRunRecordsExportEntries( runRecords: RunRecordsExportPayload | null | undefined, ) { @@ -279,6 +283,8 @@ export type AccountExportSectionResult = { nextStartAfter: string | null pageSize: number warnings: Array + /** Non-authoritative UserMeter storage-byte shadow; not part of paging. */ + storageBytesShadow?: UserMeterExportResult['storageBytesShadow'] } function normalizePageSize(pageSize: number | undefined) { @@ -1050,7 +1056,7 @@ async function collectManifestInventoryCounts(input: { }).exportCounters({ pageSize: maxExportPageSize, }) - return page.counters.length + return countUserMeterExportEntries(page) }), safeCount('remote connectors', async () => countScalar( @@ -1571,7 +1577,7 @@ function buildManifest(input: { count: input.durableObjects?.userMeter == null ? (input.inventoryCounts?.userMeterCounters ?? 0) - : input.durableObjects.userMeter.counters.length, + : countUserMeterExportEntries(input.durableObjects.userMeter), warnings: input.warnings.filter( (warning) => warning.startsWith('User meter ') || warning.startsWith('USER_METER '), @@ -1911,6 +1917,7 @@ export async function readAccountExportSection(input: { return { section: input.section, items: page.counters, + storageBytesShadow: page.storageBytesShadow, truncated: page.truncated, nextStartAfter: page.nextStartAfter, pageSize, diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 34d779fea8..92e3645f7b 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -130,7 +130,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray + d1StorageBytes?: number }) { const stableUserId = testStableUserIdFromEmail(input.email) const dailyCounters = input.dailyCounters?.map((row) => ({ ...row })) ?? [] + const d1StorageBytes = input.d1StorageBytes ?? 0 return { stableUserId, db: { @@ -60,10 +62,12 @@ function createUsageTestDb(input: { ) return (row ? { count: row.count } : { count: 0 }) as T } + if (normalized.includes('d1_storage_bytes')) { + return { bytes: d1StorageBytes } as T + } if ( normalized.includes('count(*)') || - normalized.includes('sum(') || - normalized.includes('d1_storage_bytes') + normalized.includes('sum(') ) { return { count: 0, total: 0, bytes: 0 } as T } @@ -216,3 +220,24 @@ test('loadAccountUsageData returns plan rows and authoritative UserMeter daily c expect(currentFor(authoritative, 'concurrent_workflows')).toBe(3) expect(currentFor(authoritative, 'saved_packages')).toBe(4) }) + +test('loadAccountUsageData prefers D1 storage_bytes over a divergent UserMeter shadow', async () => { + const now = new Date('2026-07-25T12:00:00.000Z') + const email = 'usage-storage@example.com' + const userId = testStableUserIdFromEmail(email) + const { db } = createUsageTestDb({ + userId: 11, + email, + plan: 'pro', + packageCount: 0, + d1StorageBytes: 4_321, + }) + const usageEnv = withUsageEnv({ APP_DB: db }) + await usageEnv.meter.seedStorageBytes({ userId, bytes: 321 }) + const data = await loadAccountUsageData({ + env: usageEnv as Env, + userId: 11, + now, + }) + expect(currentFor(data, 'storage_bytes')).toBe(4_321) +}) From 3c13b53db7b5f79574cc3c9cc8cbdf1556e517fc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 07:16:36 +0000 Subject: [PATCH 15/39] fix(entitlements): make storage reserves fail closed Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 31 ++--- .../contributing/architecture/entitlements.md | 15 ++- .../worker/src/account/export.node.test.ts | 13 ++- packages/worker/src/account/export.ts | 5 +- .../worker/src/account/user-owned-surfaces.ts | 2 +- .../entitlements/entitlements.node.test.ts | 104 +++++++++++++++++ packages/worker/src/entitlements/service.ts | 110 ++++++++++++------ .../worker/src/entitlements/user-meter-do.ts | 12 +- .../entitlements/user-meter.workers.test.ts | 32 ++++- .../worker/src/test-support/user-meter.ts | 13 ++- 10 files changed, 273 insertions(+), 64 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 8e74c6138e..7bd1099f08 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -167,13 +167,14 @@ migration-safe chunked interface: StorageRunner `exportStorage({ pageSize, startAfter })` RPC as the dedicated storage export capability. User meter counters use `section: "user_meter"` and the `UserMeter.exportCounters` RPC (daily counters plus additive - `storageBytesShadow` when present; explicitly non-authoritative). R2 raw MIME, - attachment, avatar, and icon objects use `section: "r2_object"`; each response - contains at most one 256 KiB base64 chunk and an opaque cursor. Each request - uses bounded `LIMIT 1` ownership queries rather than reconstructing inventory. - Continuation cursors bind the source row, object key, size, and ETag; - ownership/key mutations and object overwrites are reported instead of mixing - generations. Missing objects are represented explicitly. + `storageBytesShadow` on the first page only when present; explicitly + non-authoritative). R2 raw MIME, attachment, avatar, and icon objects use + `section: "r2_object"`; each response contains at most one 256 KiB base64 + chunk and an opaque cursor. Each request uses bounded `LIMIT 1` ownership + queries rather than reconstructing inventory. Continuation cursors bind the + source row, object key, size, and ETag; ownership/key mutations and object + overwrites are reported instead of mixing generations. Missing objects are + represented explicitly. D1 manifest counts use bounded SQL `COUNT(*)` queries. D1 section rows are read with SQL-level keyset pagination: every query orders by the table's `rowid`, @@ -197,12 +198,13 @@ Durable Object export behavior: never pruned by retention. See [Run records](./run-records.md). - `UserMeter` exports daily entitlement counter rows through the `user_meter` section (`exportCounters` RPC; keyset pagination by UTC `day` and `resource`). - The same RPC may return additive `storageBytesShadow` when the schema-v4 - shadow row exists (`null` when never shadowed); section totals count it as one - row when present, but it is explicitly **non-authoritative** — usage and - enforcement read D1 `users.d1_storage_bytes`. Retention is self-enforced - inside the DO (seven UTC days of counter and inbound-delivery-claim rows); - shadow storage-byte state is not time-pruned. See + The same RPC may return additive `storageBytesShadow` on the first page only + (`startAfter` absent) when the schema-v4 shadow row exists (`null` on later + pages and when never shadowed); section totals count it as one row when + present, but it is explicitly **non-authoritative** — usage and enforcement + read D1 `users.d1_storage_bytes`. Retention is self-enforced inside the DO + (seven UTC days of counter and inbound-delivery-claim rows); shadow + storage-byte state is not time-pruned. See [Entitlements](./entitlements.md#usermeter-expand-phase). - `RemoteConnectorSession` exposes persisted connector metadata and tool descriptors through an export RPC. @@ -555,7 +557,8 @@ Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan; `deleteAll` clears counters, claims, and any shadow storage state). Account export pages `UserMeter.exportCounters` through the `user_meter` manifest section / `account_export_section` (daily counters plus additive -`storageBytesShadow` when present; shadow field is non-authoritative). +`storageBytesShadow` on the first page only when present; shadow field is +non-authoritative). ### Package state model diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 2879002bd5..7171beca42 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -183,8 +183,11 @@ duplicate mirror SQL. D1 `users.d1_storage_bytes` and `users.d1_storage_bytes_updated_at` remain the **sole read, enforcement, and reconciliation authority** for the D1 payload component of `storage_bytes`. Every caller — including email — reserves through -the same atomic D1 `UPDATE` in `assertWithinStorageBytesEntitlement`; denial -semantics and `EntitlementLimitError` always come from D1. +the same conditional D1 `UPDATE` in `assertWithinStorageBytesEntitlement` +(bounded retry while a real users row remains under limit; fail closed if +contention never resolves; missing-user synthetic contexts keep prior +under-limit allow / over-limit deny). Denial semantics and +`EntitlementLimitError` always come from D1. UserMeter schema **v4** adds a singleton `storage_bytes_state` row (`id = 1`) as **best-effort shadow / future-cutover support only**. The DO row is never read @@ -211,9 +214,11 @@ failures are logged and never fail or poison the lane. Current `readUserD1StorageBytes` only. **Account export and purge:** `UserMeter.exportCounters` may return additive -non-authoritative `storageBytesShadow` when the shadow row exists (counted in -the `user_meter` section total). `UserMeter.purge()` clears counters, inbound -delivery claims, and any shadow storage state via `deleteAll`. +non-authoritative `storageBytesShadow` on the first page only (`startAfter` +absent) when the shadow row exists; subsequent pages return `null` so paged +consumers never double-count it (still counted once in the `user_meter` section +total). `UserMeter.purge()` clears counters, inbound delivery claims, and any +shadow storage state via `deleteAll`. ### Future storage authority flip (contract follow-up) diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 97a962c81c..815d16f8cb 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1588,9 +1588,11 @@ test('account export includes user_meter counters, pages them, and warns on trun : 0 const page = counters.slice(startIndex, startIndex + pageSize) const truncated = startIndex + pageSize < counters.length + const isFirstPage = + typeof input.startAfter !== 'string' || input.startAfter.length === 0 return { counters: page, - storageBytesShadow, + storageBytesShadow: isFirstPage ? storageBytesShadow : null, nextStartAfter: truncated ? `${page.at(-1)!.day}:${page.at(-1)!.resource}` : null, @@ -1647,10 +1649,15 @@ test('account export includes user_meter counters, pages them, and warns on trun startAfter: first.nextStartAfter ?? undefined, }) expect(second.items).toEqual(counters.slice(2)) - expect(second.storageBytesShadow).toEqual(storageBytesShadow) + expect(second.storageBytesShadow).toBeNull() expect(second.truncated).toBe(false) expect(second.nextStartAfter).toBeNull() - expect(exportCounters).toHaveBeenCalled() + expect(exportCounters).toHaveBeenCalledWith( + expect.objectContaining({ + pageSize: 2, + startAfter: first.nextStartAfter, + }), + ) exportCounters.mockImplementation(async () => ({ counters: [counters[0]!], diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index 297abc2c38..a2ac581228 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -283,7 +283,10 @@ export type AccountExportSectionResult = { nextStartAfter: string | null pageSize: number warnings: Array - /** Non-authoritative UserMeter storage-byte shadow; not part of paging. */ + /** + * Non-authoritative UserMeter storage-byte shadow. Present only on the + * first `user_meter` page (`startAfter` absent); later pages omit it. + */ storageBytesShadow?: UserMeterExportResult['storageBytesShadow'] } diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 92e3645f7b..47a46d6fc3 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -130,7 +130,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray { + const userId = await createStableUserIdFromEmail(plannedEmail) + const proLimit = planLimits.pro.maxStorageBytes + let reserveAttempts = 0 + let bytes = proLimit - 10 + const db = { + prepare(query: string) { + return { + bind(...params: Array) { + return { + async first() { + if ( + query.includes('SELECT plan, stripe_plan FROM users') || + query.includes('SELECT plan FROM users') + ) { + return { plan: 'pro', stripe_plan: null } as T + } + if (query.includes('SELECT d1_storage_bytes AS bytes')) { + return { bytes } as T + } + throw new Error(`Unsupported first query: ${query}`) + }, + async run() { + if ( + query.includes('SET d1_storage_bytes = d1_storage_bytes + ?') + ) { + reserveAttempts += 1 + if (reserveAttempts === 1) { + return { meta: { changes: 0 } } + } + const requested = Number(params[0]) + if (bytes + Number(params[3]) > Number(params[4])) { + return { meta: { changes: 0 } } + } + bytes += requested + return { meta: { changes: 1 } } + } + throw new Error(`Unsupported run query: ${query}`) + }, + } + }, + } + }, + } as unknown as D1Database + + await assertWithinStorageBytesEntitlement({ + db, + userId, + email: plannedEmail, + requested: 5, + }) + expect(reserveAttempts).toBe(2) + expect(bytes).toBe(proLimit - 5) +}) + +test('storage byte reserve fails closed when contention never resolves', async () => { + const userId = await createStableUserIdFromEmail(plannedEmail) + const proLimit = planLimits.pro.maxStorageBytes + let reserveAttempts = 0 + const db = { + prepare(query: string) { + return { + bind() { + return { + async first() { + if ( + query.includes('SELECT plan, stripe_plan FROM users') || + query.includes('SELECT plan FROM users') + ) { + return { plan: 'pro', stripe_plan: null } as T + } + if (query.includes('SELECT d1_storage_bytes AS bytes')) { + return { bytes: proLimit - 10 } as T + } + throw new Error(`Unsupported first query: ${query}`) + }, + async run() { + if ( + query.includes('SET d1_storage_bytes = d1_storage_bytes + ?') + ) { + reserveAttempts += 1 + return { meta: { changes: 0 } } + } + throw new Error(`Unsupported run query: ${query}`) + }, + } + }, + } + }, + } as unknown as D1Database + + await expect( + assertWithinStorageBytesEntitlement({ + db, + userId, + email: plannedEmail, + requested: 5, + }), + ).rejects.toThrow( + 'Storage byte reservation could not complete under concurrent updates.', + ) + expect(reserveAttempts).toBe(5) +}) + test('getPlanRank orders free < pro < partner < max', () => { expect(getPlanRank('free')).toBeLessThan(getPlanRank('pro')) expect(getPlanRank('pro')).toBeLessThan(getPlanRank('partner')) diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index ac3d13183b..5d3a5e6ae5 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -725,10 +725,10 @@ function scheduleUserMeterStorageBytesShadow(input: { void tracked } -export async function readUserD1StorageBytes(input: { +async function readUserD1StorageBytesRow(input: { db: D1Database userId: string -}): Promise { +}): Promise<{ bytes: number } | null> { const row = await input.db .prepare( `SELECT d1_storage_bytes AS bytes @@ -737,7 +737,16 @@ export async function readUserD1StorageBytes(input: { ) .bind(input.userId) .first<{ bytes: number }>() - return Math.max(0, Number(row?.bytes ?? 0)) + if (row == null) return null + return { bytes: Math.max(0, Number(row.bytes ?? 0)) } +} + +export async function readUserD1StorageBytes(input: { + db: D1Database + userId: string +}): Promise { + const row = await readUserD1StorageBytesRow(input) + return row?.bytes ?? 0 } /** @@ -985,11 +994,20 @@ export async function readCurrentEntitlementResourceUsage(input: { }) } +/** Bounded retries when a conditional D1 reserve loses to concurrent writers. */ +const storageBytesReserveMaxAttempts = 5 + /** * Atomically reserve D1 payload storage bytes against the plan limit. * Optional `env` schedules a best-effort UserMeter shadow after a successful * reserve (never awaited; failures logged only). `getCurrent` is check-only * and does not reserve in D1 or UserMeter. + * + * Real users never succeed unreserved: a failed conditional update re-reads + * row existence/current, throws {@link EntitlementLimitError} at the limit, + * retries while a real row remains under limit, and fails closed with an + * operational error if contention never resolves. Missing-user synthetic + * contexts keep their prior under-limit allow / over-limit deny semantics. */ export async function assertWithinStorageBytesEntitlement(input: { db: D1Database @@ -1020,42 +1038,66 @@ export async function assertWithinStorageBytesEntitlement(input: { }) const limit = resolvePlanLimit(plan, 'storage_bytes') const requested = Math.max(0, input.requested ?? 1) - const now = new Date().toISOString() - const result = await input.db - .prepare( - `UPDATE users - SET d1_storage_bytes = d1_storage_bytes + ?, - d1_storage_bytes_updated_at = ? - WHERE stable_user_id = ? - AND d1_storage_bytes + ? <= ?`, - ) - .bind(requested, now, input.userId, requested, limit) - .run() - if ((result.meta.changes ?? 0) > 0) { - scheduleUserMeterStorageBytesShadow({ + + for (let attempt = 0; attempt < storageBytesReserveMaxAttempts; attempt++) { + const now = new Date().toISOString() + const result = await input.db + .prepare( + `UPDATE users + SET d1_storage_bytes = d1_storage_bytes + ?, + d1_storage_bytes_updated_at = ? + WHERE stable_user_id = ? + AND d1_storage_bytes + ? <= ?`, + ) + .bind(requested, now, input.userId, requested, limit) + .run() + if ((result.meta.changes ?? 0) > 0) { + scheduleUserMeterStorageBytesShadow({ + db: input.db, + env: input.env, + userId: input.userId, + waitUntil: input.waitUntil, + }) + return + } + + const row = await readUserD1StorageBytesRow({ db: input.db, - env: input.env, userId: input.userId, - waitUntil: input.waitUntil, }) - return + // Synthetic contexts have no users row to reserve against. Their plan + // has already failed closed to `free`; preserve that finite limit + // without inventing a user-owned counter row that reconciliation + // cannot resolve. + if (row == null) { + const current = 0 + if (current + requested <= limit) return + throw new EntitlementLimitError({ + resource: 'storage_bytes', + plan, + limit, + current, + upgradeHint: buildEntitlementUpgradeHint('storage_bytes'), + }) + } + + const current = row.bytes + if (current + requested > limit) { + throw new EntitlementLimitError({ + resource: 'storage_bytes', + plan, + limit, + current, + upgradeHint: buildEntitlementUpgradeHint('storage_bytes'), + }) + } + // Real row still under limit: another writer raced the conditional + // UPDATE. Retry the reserve rather than succeeding unreserved. } - const current = await readUserD1StorageBytes({ - db: input.db, - userId: input.userId, - }) - // Synthetic contexts have no users row to reserve against. Their plan has - // already failed closed to `free`; preserve that finite limit without - // inventing a user-owned counter row that reconciliation cannot resolve. - if (current + requested <= limit) return - throw new EntitlementLimitError({ - resource: 'storage_bytes', - plan, - limit, - current, - upgradeHint: buildEntitlementUpgradeHint('storage_bytes'), - }) + throw new Error( + 'Storage byte reservation could not complete under concurrent updates.', + ) } export type AssertWithinEntitlementInput = { diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index adc12de48d..8633bf0b91 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -112,7 +112,11 @@ export type UserMeterStorageBytesSetResult = UserMeterStorageBytesReadyState & { export type UserMeterExportResult = { counters: Array - /** Non-authoritative D1 storage-byte shadow; `null` when never synced. */ + /** + * Non-authoritative D1 storage-byte shadow. Emitted only on the first + * export page (`startAfter` absent); subsequent pages return `null` so + * paged consumers never double-count the singleton shadow. + */ storageBytesShadow: UserMeterStorageBytesState | null nextStartAfter: string | null truncated: boolean @@ -838,7 +842,11 @@ class UserMeterBase extends DurableObject { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(revision), }) } - const storageRow = this.readStorageRow() + // Shadow is a singleton outside keyset paging — emit it once on the + // first page only so section totals and multi-page consumers do not + // double-count. + const includeStorageShadow = cursor == null + const storageRow = includeStorageShadow ? this.readStorageRow() : null const last = pageRows[pageRows.length - 1] return { counters, diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index 3722d4b156..5c1440355a 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -938,13 +938,41 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a updatedAt: '2026-07-31T17:02:00.000Z', }) - const exported = await meter.exportCounters({}) - expect(exported.storageBytesShadow).toEqual({ + const day = '2026-07-31' + for (const resource of [ + 'email_receives_per_day', + 'email_sends_per_day', + 'execute_calls_per_day', + 'outbound_fetches_per_day', + ] as const) { + await meter.initialize({ + resource, + day, + count: 1, + updatedAt: '2026-07-31T17:03:00.000Z', + }) + } + + const firstPage = await meter.exportCounters({ pageSize: 2 }) + expect(firstPage.counters).toHaveLength(2) + expect(firstPage.truncated).toBe(true) + expect(firstPage.nextStartAfter).toEqual(expect.any(String)) + expect(firstPage.storageBytesShadow).toEqual({ bytes: 11, revision: 3, updatedAt: '2026-07-31T17:02:00.000Z', mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(3), }) + + const secondPage = await meter.exportCounters({ + pageSize: 2, + startAfter: firstPage.nextStartAfter, + }) + expect(secondPage.counters).toHaveLength(2) + expect(secondPage.truncated).toBe(false) + expect(secondPage.nextStartAfter).toBeNull() + expect(secondPage.storageBytesShadow).toBeNull() + await expect( readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), ).resolves.toBe(0) diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index c02fe7f3b5..8452d2d284 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -182,7 +182,12 @@ export function createInMemoryUserMeterEnv() { storageByUser.delete(userId) return { ok: true as const } }, - async exportCounters() { + async exportCounters( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ) { const counters = [...rows.entries()].map(([entryKey, row]) => { const [resource, day] = entryKey.split('\0') return { @@ -194,7 +199,11 @@ export function createInMemoryUserMeterEnv() { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), } }) - const storage = storageByUser.get(userId) + const includeStorageShadow = + typeof input.startAfter !== 'string' || input.startAfter.length === 0 + const storage = includeStorageShadow + ? storageByUser.get(userId) + : undefined const storageBytesShadow: UserMeterStorageBytesState | null = storage ? { bytes: storage.bytes, From e454e7666d31713345a15bce1d89cb62b44424dd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 07:48:58 +0000 Subject: [PATCH 16/39] feat(services): shadow liveness in UserMeter Co-authored-by: Kent C. Dodds --- .../worker/src/entitlements/user-meter-do.ts | 469 +++++++++++++++++- .../entitlements/user-meter.workers.test.ts | 191 +++++++ .../package-service.node.test.ts | 339 +++++++++++++ .../src/package-runtime/package-service.ts | 74 ++- .../worker/src/test-support/user-meter.ts | 229 ++++++++- 5 files changed, 1285 insertions(+), 17 deletions(-) diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 8633bf0b91..0103decadf 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -26,16 +26,36 @@ export const userMeterDailyCounterRetentionDays = 7 const metaSchemaVersionKey = 'schema_version' /** Bump when initializeSchema DDL changes; warm objects skip DDL. */ -const userMeterSchemaVersion = 4 +const userMeterSchemaVersion = 5 /** Singleton row id for expand-phase D1 storage-byte shadow (schema v4). */ const storageBytesStateRowId = 1 +/** Matches D1 `packageServiceStateStaleMs` (24h); cutover-support RPC only. */ +export const userMeterPackageServiceStateStaleMs = 24 * 60 * 60 * 1000 const defaultExportPageSize = 100 const maxExportPageSize = 500 const maxInboundDeliveryIdLength = 256 +const maxPackageServiceIdLength = 256 const inboundReceiveResource = 'email_receives_per_day' satisfies DailyEntitlementResource const utcDayKeyPattern = /^\d{4}-\d{2}-\d{2}$/ +const packageServiceShadowStatuses = [ + 'running', + 'idle', + 'stopped', + 'error', +] as const + +export type UserMeterPackageServiceStatus = + (typeof packageServiceShadowStatuses)[number] + +export function isUserMeterPackageServiceStatus( + status: string, +): status is UserMeterPackageServiceStatus { + return (packageServiceShadowStatuses as ReadonlyArray).includes( + status, + ) +} /** * Legacy D1 mirror `updated_at` token. Lexicographic order matches revision @@ -110,6 +130,42 @@ export type UserMeterStorageBytesSetResult = UserMeterStorageBytesReadyState & { created: boolean } +export type UserMeterPackageServiceState = { + packageId: string + serviceName: string + status: UserMeterPackageServiceStatus + startedAt: string | null + sourceUpdatedAt: string + revision: number + updatedAt: string + mirrorUpdatedAt: string +} + +export type UserMeterPackageServiceUpsertResult = { + applied: boolean + created: boolean + state: UserMeterPackageServiceState +} + +export type UserMeterPackageServiceDeleteResult = { + deleted: boolean +} + +export type UserMeterPackageServiceListResult = { + states: Array + nextStartAfter: string | null + truncated: boolean +} + +export type UserMeterPackageServiceCountResult = { + count: number +} + +export type UserMeterPackageServiceBootstrapResult = { + applied: number + skipped: number +} + export type UserMeterExportResult = { counters: Array /** @@ -118,6 +174,12 @@ export type UserMeterExportResult = { * paged consumers never double-count the singleton shadow. */ storageBytesShadow: UserMeterStorageBytesState | null + /** + * Non-authoritative D1 `package_service_states` shadow rows. Emitted only + * on the first export page (`startAfter` absent); subsequent pages return + * `null` so paged consumers never double-count the inventory. + */ + packageServiceStatesShadow: Array | null nextStartAfter: string | null truncated: boolean } @@ -141,6 +203,21 @@ type ExportCursor = { resource: string } +type PackageServiceExportCursor = { + packageId: string + serviceName: string +} + +type PackageServiceSqlRow = { + package_id: string + service_name: string + status: string + started_at: string | null + source_updated_at: string + revision: number + updated_at: string +} + function assertDailyResource(resource: string): DailyEntitlementResource { if (!isDailyEntitlementResource(resource)) { throw new Error( @@ -184,6 +261,39 @@ function assertInboundDeliveryId(deliveryId: string): string { return deliveryId } +function assertPackageServiceId(label: string, value: string): string { + if ( + typeof value !== 'string' || + value.length === 0 || + value.length > maxPackageServiceIdLength + ) { + throw new Error( + `UserMeter ${label} must be a non-empty string up to ${maxPackageServiceIdLength} characters.`, + ) + } + return value +} + +function assertPackageServiceStatus( + status: string, +): UserMeterPackageServiceStatus { + if (!isUserMeterPackageServiceStatus(status)) { + throw new Error( + `UserMeter package service status must be one of ${packageServiceShadowStatuses.join(', ')}; got ${JSON.stringify(status)}.`, + ) + } + return status +} + +function assertSourceUpdatedAt(sourceUpdatedAt: string): string { + if (typeof sourceUpdatedAt !== 'string' || sourceUpdatedAt.length === 0) { + throw new Error( + 'UserMeter package service sourceUpdatedAt must be a non-empty string.', + ) + } + return sourceUpdatedAt +} + function retentionCutoffDay(now: Date): string { const cutoff = new Date(now) cutoff.setUTCDate( @@ -208,6 +318,26 @@ function decodeExportCursor(startAfter: string): ExportCursor | null { } } +function encodePackageServiceCursor(cursor: PackageServiceExportCursor) { + return JSON.stringify([cursor.packageId, cursor.serviceName]) +} + +function decodePackageServiceCursor( + startAfter: string, +): PackageServiceExportCursor | null { + try { + const parsed = JSON.parse(startAfter) as unknown + if (!Array.isArray(parsed) || parsed.length !== 2) return null + const [packageId, serviceName] = parsed + if (typeof packageId !== 'string' || typeof serviceName !== 'string') { + return null + } + return { packageId, serviceName } + } catch { + return null + } +} + function normalizePageSize(pageSize: number | undefined) { const requested = typeof pageSize === 'number' && Number.isFinite(pageSize) @@ -301,6 +431,23 @@ class UserMeterBase extends DurableObject { updated_at TEXT NOT NULL ) `) + // Expand-phase shadow of D1 `package_service_states` (schema v5). + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS package_service_states ( + package_id TEXT NOT NULL, + service_name TEXT NOT NULL, + status TEXT NOT NULL, + started_at TEXT, + source_updated_at TEXT NOT NULL, + revision INTEGER NOT NULL, + updated_at TEXT NOT NULL, + PRIMARY KEY (package_id, service_name) + ) + `) + this.ctx.storage.sql.exec( + `CREATE INDEX IF NOT EXISTS idx_package_service_states_status_source + ON package_service_states (status, source_updated_at)`, + ) this.ctx.storage.sql.exec( `INSERT INTO user_meter_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, @@ -772,6 +919,276 @@ class UserMeterBase extends DurableObject { } } + private packageServiceStateFromRow( + row: PackageServiceSqlRow, + ): UserMeterPackageServiceState { + const revision = Math.max(0, Number(row.revision ?? 0)) + const status = assertPackageServiceStatus(String(row.status)) + return { + packageId: String(row.package_id), + serviceName: String(row.service_name), + status, + startedAt: + status === 'running' && row.started_at != null + ? String(row.started_at) + : null, + sourceUpdatedAt: String(row.source_updated_at), + revision, + updatedAt: String(row.updated_at), + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(revision), + } + } + + private readPackageServiceRow( + packageId: string, + serviceName: string, + ): UserMeterPackageServiceState | null { + const row = this.ctx.storage.sql + .exec( + `SELECT package_id, service_name, status, started_at, + source_updated_at, revision, updated_at + FROM package_service_states + WHERE package_id = ? AND service_name = ?`, + packageId, + serviceName, + ) + .toArray()[0] + if (!row) return null + return this.packageServiceStateFromRow(row) + } + + private listAllPackageServiceRows(): Array { + const rows = this.ctx.storage.sql + .exec( + `SELECT package_id, service_name, status, started_at, + source_updated_at, revision, updated_at + FROM package_service_states + ORDER BY package_id ASC, service_name ASC`, + ) + .toArray() + return rows.map((row) => this.packageServiceStateFromRow(row)) + } + + /** + * Expand-phase shadow upsert (monotonic on `sourceUpdatedAt`). Stale writes + * are rejected; expand-phase enforcement still reads D1. + */ + async upsertPackageServiceState(input: { + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }): Promise { + const packageId = assertPackageServiceId('packageId', input.packageId) + const serviceName = assertPackageServiceId('serviceName', input.serviceName) + const status = assertPackageServiceStatus(input.status) + const sourceUpdatedAt = assertSourceUpdatedAt(input.sourceUpdatedAt) + const updatedAt = + typeof input.updatedAt === 'string' && input.updatedAt.length > 0 + ? input.updatedAt + : sourceUpdatedAt + const startedAt = + status === 'running' && + typeof input.startedAt === 'string' && + input.startedAt.length > 0 + ? input.startedAt + : null + + const existing = this.readPackageServiceRow(packageId, serviceName) + if (existing && existing.sourceUpdatedAt > sourceUpdatedAt) { + return { + applied: false, + created: false, + state: existing, + } + } + if (!existing) { + this.ctx.storage.sql.exec( + `INSERT INTO package_service_states ( + package_id, service_name, status, started_at, + source_updated_at, revision, updated_at + ) VALUES (?, ?, ?, ?, ?, 1, ?)`, + packageId, + serviceName, + status, + startedAt, + sourceUpdatedAt, + updatedAt, + ) + const row = this.readPackageServiceRow(packageId, serviceName) + if (!row) { + throw new Error( + 'UserMeter upsertPackageServiceState failed to materialize shadow row.', + ) + } + return { applied: true, created: true, state: row } + } + + const nextRevision = existing.revision + 1 + this.ctx.storage.sql.exec( + `UPDATE package_service_states + SET status = ?, + started_at = ?, + source_updated_at = ?, + revision = ?, + updated_at = ? + WHERE package_id = ? AND service_name = ? AND revision = ?`, + status, + startedAt, + sourceUpdatedAt, + nextRevision, + updatedAt, + packageId, + serviceName, + existing.revision, + ) + const row = this.readPackageServiceRow(packageId, serviceName) + if (!row) { + throw new Error( + 'UserMeter upsertPackageServiceState lost the shadow row.', + ) + } + return { applied: true, created: false, state: row } + } + + /** Expand-phase shadow delete; discovery still uses D1. */ + async deletePackageServiceState(input: { + packageId: string + serviceName: string + }): Promise { + const packageId = assertPackageServiceId('packageId', input.packageId) + const serviceName = assertPackageServiceId('serviceName', input.serviceName) + const cursor = this.ctx.storage.sql.exec( + `DELETE FROM package_service_states + WHERE package_id = ? AND service_name = ?`, + packageId, + serviceName, + ) + return { deleted: cursor.rowsWritten > 0 } + } + + /** Cutover-support paged shadow list; expand-phase discovery uses D1. */ + async listPackageServiceStates(input: { + pageSize?: number + startAfter?: string | null + }): Promise { + const pageSize = normalizePageSize(input.pageSize) + const cursor = + typeof input.startAfter === 'string' && input.startAfter.length > 0 + ? decodePackageServiceCursor(input.startAfter) + : null + const rows = ( + cursor + ? this.ctx.storage.sql.exec( + `SELECT package_id, service_name, status, started_at, + source_updated_at, revision, updated_at + FROM package_service_states + WHERE package_id > ? + OR (package_id = ? AND service_name > ?) + ORDER BY package_id ASC, service_name ASC + LIMIT ?`, + cursor.packageId, + cursor.packageId, + cursor.serviceName, + pageSize + 1, + ) + : this.ctx.storage.sql.exec( + `SELECT package_id, service_name, status, started_at, + source_updated_at, revision, updated_at + FROM package_service_states + ORDER BY package_id ASC, service_name ASC + LIMIT ?`, + pageSize + 1, + ) + ).toArray() + const truncated = rows.length > pageSize + const pageRows = truncated ? rows.slice(0, pageSize) : rows + const states = pageRows.map((row) => this.packageServiceStateFromRow(row)) + const last = pageRows[pageRows.length - 1] + return { + states, + nextStartAfter: + truncated && last + ? encodePackageServiceCursor({ + packageId: String(last.package_id), + serviceName: String(last.service_name), + }) + : null, + truncated, + } + } + + /** + * Cutover-support running count from the shadow table. Expand-phase + * enforcement still uses D1 `countRunningPackageServices`. + */ + async countRunningPackageServices(input: { + staleAfterMs?: number + excludeService?: { packageId: string; serviceName: string } + now?: string + }): Promise { + const now = input.now ? new Date(input.now) : new Date() + const safeNow = Number.isNaN(now.valueOf()) ? new Date() : now + const staleAfterMs = + typeof input.staleAfterMs === 'number' && + Number.isFinite(input.staleAfterMs) && + input.staleAfterMs >= 0 + ? Math.trunc(input.staleAfterMs) + : userMeterPackageServiceStateStaleMs + const freshAfter = new Date(safeNow.valueOf() - staleAfterMs).toISOString() + const exclusion = input.excludeService + const packageId = exclusion + ? assertPackageServiceId('packageId', exclusion.packageId) + : null + const serviceName = exclusion + ? assertPackageServiceId('serviceName', exclusion.serviceName) + : null + const row = ( + packageId && serviceName + ? this.ctx.storage.sql.exec<{ count: number }>( + `SELECT COUNT(*) AS count + FROM package_service_states + WHERE status = 'running' + AND source_updated_at >= ? + AND NOT (package_id = ? AND service_name = ?)`, + freshAfter, + packageId, + serviceName, + ) + : this.ctx.storage.sql.exec<{ count: number }>( + `SELECT COUNT(*) AS count + FROM package_service_states + WHERE status = 'running' + AND source_updated_at >= ?`, + freshAfter, + ) + ).toArray()[0] + return { count: Math.max(0, Number(row?.count ?? 0)) } + } + + /** Cutover-support bulk seed; same monotonic guard as upsert. */ + async bootstrapPackageServiceStates(input: { + states: ReadonlyArray<{ + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }> + }): Promise { + let applied = 0 + let skipped = 0 + for (const state of input.states) { + const result = await this.upsertPackageServiceState(state) + if (result.applied) applied += 1 + else skipped += 1 + } + return { applied, skipped } + } + async purge(): Promise<{ ok: true }> { await this.ctx.blockConcurrencyWhile(async () => { await this.ctx.storage.deleteAll() @@ -842,11 +1259,14 @@ class UserMeterBase extends DurableObject { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(revision), }) } - // Shadow is a singleton outside keyset paging — emit it once on the - // first page only so section totals and multi-page consumers do not - // double-count. - const includeStorageShadow = cursor == null - const storageRow = includeStorageShadow ? this.readStorageRow() : null + // Shadow inventories sit outside counter keyset paging — emit them once + // on the first page only so section totals and multi-page consumers do + // not double-count. + const includeShadows = cursor == null + const storageRow = includeShadows ? this.readStorageRow() : null + const packageServiceStatesShadow = includeShadows + ? this.listAllPackageServiceRows() + : null const last = pageRows[pageRows.length - 1] return { counters, @@ -858,6 +1278,7 @@ class UserMeterBase extends DurableObject { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(storageRow.revision), } : null, + packageServiceStatesShadow, nextStartAfter: truncated && last ? encodeExportCursor({ @@ -923,6 +1344,42 @@ export type UserMeterRpc = { bytes: number updatedAt: string }) => Promise + /** Expand-phase shadow upsert; not used for enforcement. */ + upsertPackageServiceState: (input: { + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }) => Promise + /** Expand-phase shadow delete. */ + deletePackageServiceState: (input: { + packageId: string + serviceName: string + }) => Promise + /** Cutover-support paged shadow list; discovery uses D1. */ + listPackageServiceStates: (input: { + pageSize?: number + startAfter?: string | null + }) => Promise + /** Cutover-support running count; service-start still uses D1. */ + countRunningPackageServices: (input: { + staleAfterMs?: number + excludeService?: { packageId: string; serviceName: string } + now?: string + }) => Promise + /** Cutover-support bulk seed with sourceUpdatedAt monotonic guards. */ + bootstrapPackageServiceStates: (input: { + states: ReadonlyArray<{ + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }> + }) => Promise purge: () => Promise<{ ok: true }> exportCounters: (input: { pageSize?: number diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index 5c1440355a..ccafdb5e59 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -561,6 +561,7 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p expect(await meterA.exportCounters({})).toEqual({ counters: [], storageBytesShadow: null, + packageServiceStatesShadow: [], nextStartAfter: null, truncated: false, }) @@ -690,6 +691,7 @@ test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore expect(exportDuringPurge).toEqual({ counters: [], storageBytesShadow: null, + packageServiceStatesShadow: [], nextStartAfter: null, truncated: false, }) @@ -963,6 +965,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a updatedAt: '2026-07-31T17:02:00.000Z', mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(3), }) + expect(firstPage.packageServiceStatesShadow).toEqual([]) const secondPage = await meter.exportCounters({ pageSize: 2, @@ -972,6 +975,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a expect(secondPage.truncated).toBe(false) expect(secondPage.nextStartAfter).toBeNull() expect(secondPage.storageBytesShadow).toBeNull() + expect(secondPage.packageServiceStatesShadow).toBeNull() await expect( readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), @@ -984,6 +988,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a expect(await meter.exportCounters({})).toEqual({ counters: [], storageBytesShadow: null, + packageServiceStatesShadow: [], nextStartAfter: null, truncated: false, }) @@ -1024,3 +1029,189 @@ test('missing-user storage reserve preserves EntitlementLimitError semantics on }, }) }, 30_000) + +test('UserMeter package-service shadow upserts are monotonic, isolated, exportable, and purgeable', async () => { + const userA = await seedFreeUser('meter-pkg-svc-a') + const userB = await seedFreeUser('meter-pkg-svc-b') + const meterA = userMeterRpc({ env, userId: userA.userId }) + const meterB = userMeterRpc({ env, userId: userB.userId }) + + const created = await meterA.upsertPackageServiceState({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + startedAt: '2026-08-01T10:00:00.000Z', + sourceUpdatedAt: '2026-08-01T10:00:00.000Z', + }) + expect(created).toMatchObject({ + applied: true, + created: true, + state: { + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + startedAt: '2026-08-01T10:00:00.000Z', + sourceUpdatedAt: '2026-08-01T10:00:00.000Z', + revision: 1, + }, + }) + + const heartbeat = await meterA.upsertPackageServiceState({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + startedAt: '2026-08-01T10:00:00.000Z', + sourceUpdatedAt: '2026-08-01T11:00:00.000Z', + }) + expect(heartbeat).toMatchObject({ + applied: true, + created: false, + state: { + status: 'running', + sourceUpdatedAt: '2026-08-01T11:00:00.000Z', + revision: 2, + }, + }) + + const stale = await meterA.upsertPackageServiceState({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'stopped', + startedAt: null, + sourceUpdatedAt: '2026-08-01T10:30:00.000Z', + }) + expect(stale).toMatchObject({ + applied: false, + created: false, + state: { + status: 'running', + sourceUpdatedAt: '2026-08-01T11:00:00.000Z', + revision: 2, + }, + }) + + await meterA.upsertPackageServiceState({ + packageId: 'pkg-2', + serviceName: 'idle-svc', + status: 'idle', + sourceUpdatedAt: '2026-08-01T11:05:00.000Z', + }) + await meterB.upsertPackageServiceState({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + startedAt: '2026-08-01T11:00:00.000Z', + sourceUpdatedAt: '2026-08-01T11:00:00.000Z', + }) + + expect( + await meterA.countRunningPackageServices({ + now: '2026-08-01T11:30:00.000Z', + }), + ).toEqual({ count: 1 }) + expect( + await meterA.countRunningPackageServices({ + now: '2026-08-01T11:30:00.000Z', + excludeService: { packageId: 'pkg-1', serviceName: 'worker' }, + }), + ).toEqual({ count: 0 }) + + const listed = await meterA.listPackageServiceStates({ pageSize: 1 }) + expect(listed.states).toHaveLength(1) + expect(listed.truncated).toBe(true) + expect(listed.nextStartAfter).toEqual(expect.any(String)) + const listedRest = await meterA.listPackageServiceStates({ + pageSize: 10, + startAfter: listed.nextStartAfter, + }) + expect(listedRest.states).toHaveLength(1) + expect(listedRest.truncated).toBe(false) + + const bootstrap = await meterA.bootstrapPackageServiceStates({ + states: [ + { + packageId: 'pkg-1', + serviceName: 'worker', + status: 'error', + sourceUpdatedAt: '2026-08-01T10:00:00.000Z', + }, + { + packageId: 'pkg-3', + serviceName: 'new', + status: 'stopped', + sourceUpdatedAt: '2026-08-01T12:00:00.000Z', + }, + ], + }) + expect(bootstrap).toEqual({ applied: 1, skipped: 1 }) + + const day = '2026-08-01' + for (const resource of [ + 'email_receives_per_day', + 'email_sends_per_day', + ] as const) { + await meterA.initialize({ + resource, + day, + count: 1, + updatedAt: '2026-08-01T12:00:00.000Z', + }) + } + + const firstPage = await meterA.exportCounters({ pageSize: 1 }) + expect(firstPage.truncated).toBe(true) + expect(firstPage.nextStartAfter).toEqual(expect.any(String)) + expect(firstPage.packageServiceStatesShadow).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + revision: 2, + }), + expect.objectContaining({ + packageId: 'pkg-2', + serviceName: 'idle-svc', + status: 'idle', + }), + expect.objectContaining({ + packageId: 'pkg-3', + serviceName: 'new', + status: 'stopped', + }), + ]), + ) + const secondPage = await meterA.exportCounters({ + pageSize: 1, + startAfter: firstPage.nextStartAfter, + }) + expect(secondPage.packageServiceStatesShadow).toBeNull() + + await expect( + meterA.deletePackageServiceState({ + packageId: 'pkg-1', + serviceName: 'worker', + }), + ).resolves.toEqual({ deleted: true }) + expect( + await meterA.countRunningPackageServices({ + now: '2026-08-01T11:30:00.000Z', + }), + ).toEqual({ count: 0 }) + + await expect(meterA.purge()).resolves.toEqual({ ok: true }) + expect(await meterA.listPackageServiceStates({})).toEqual({ + states: [], + nextStartAfter: null, + truncated: false, + }) + expect(await meterB.listPackageServiceStates({})).toMatchObject({ + states: [ + expect.objectContaining({ + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running', + }), + ], + }) +}, 30_000) diff --git a/packages/worker/src/package-runtime/package-service.node.test.ts b/packages/worker/src/package-runtime/package-service.node.test.ts index 5ca9d226d6..7aa921d924 100644 --- a/packages/worker/src/package-runtime/package-service.node.test.ts +++ b/packages/worker/src/package-runtime/package-service.node.test.ts @@ -1,4 +1,9 @@ import { expect, test, vi } from 'vitest' +import { + consoleWarn, + silenceExpectedConsoleWarns, +} from '#worker/test-support/console-spies.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' const mockModule = vi.hoisted(() => ({ getSavedPackageById: vi.fn(), @@ -799,3 +804,337 @@ test('package service restore projects current state into package_service_states vi.useRealTimers() } }) + +function createRecordingPackageServiceDb() { + const upserts: Array> = [] + const deletes: Array> = [] + const appDb = { + prepare(query: string) { + return { + bind(...params: Array) { + return { + async run() { + if (query.includes('INSERT INTO package_service_states')) { + upserts.push(params) + } + if (query.includes('DELETE FROM package_service_states')) { + deletes.push(params) + } + return { meta: { changes: 1 } } + }, + } + }, + } + }, + } as unknown as D1Database + return { appDb, upserts, deletes } +} + +function createNoopPackageServiceDb() { + return { + prepare() { + return { + bind() { + return { + async run() { + return { meta: { changes: 1 } } + }, + } + }, + } + }, + } as unknown as D1Database +} + +function packageServiceShadowRow( + meter: ReturnType, + userId: string, + packageId: string, + serviceName: string, +) { + return meter.packageServicesByUser + .get(userId) + ?.get(`${packageId}\0${serviceName}`) +} + +async function postPackageService( + instance: InstanceType, + action: 'start' | 'stop' | 'purge', + binding: typeof serviceBinding = serviceBinding, +) { + return instance.fetch( + new Request(`https://package-service.invalid/service/${action}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ binding }), + }), + ) +} + +test('package service start/heartbeat/stop/purge shadow UserMeter transitions', async () => { + resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + const { appDb, upserts, deletes } = createRecordingPackageServiceDb() + const meter = createInMemoryUserMeterEnv() + + try { + setupSavedPackage('persistent') + mockModule.runBundledModuleWithRegistry.mockImplementation( + () => new Promise(() => {}), + ) + + const created = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: meter.env.USER_METER, + } as Env) + const start = await postPackageService(created.instance, 'start') + expect(start.status).toBe(200) + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ + status: 'running', + startedAt: '2026-07-05T12:00:00.000Z', + sourceUpdatedAt: '2026-07-05T12:00:00.000Z', + revision: 1, + }) + expect(created.getAlarmAt()).toBe(Date.parse('2026-07-05T13:00:00.000Z')) + + vi.setSystemTime(new Date('2026-07-05T13:00:00.000Z')) + await created.instance.alarm() + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ + status: 'running', + sourceUpdatedAt: '2026-07-05T13:00:00.000Z', + revision: 2, + }) + expect(created.getAlarmAt()).toBe(Date.parse('2026-07-05T14:00:00.000Z')) + + await postPackageService(created.instance, 'stop') + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ + status: 'running', + sourceUpdatedAt: '2026-07-05T13:00:00.000Z', + revision: 3, + }) + + await postPackageService(created.instance, 'purge') + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toBeUndefined() + expect(deletes).toContainEqual([ + 'user-123', + 'package-1', + 'realtime-supervisor', + ]) + expect(upserts.length).toBeGreaterThanOrEqual(3) + } finally { + vi.useRealTimers() + } +}) + +test('package service UserMeter shadow failure cannot break D1 projection or lifecycle', async () => { + resetMocks() + silenceExpectedConsoleWarns(['package-service-user-meter-shadow-failed']) + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + const { appDb, upserts } = createRecordingPackageServiceDb() + const failingMeter = { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: () => ({ + async upsertPackageServiceState() { + throw new Error('shadow write failed') + }, + async deletePackageServiceState() { + throw new Error('shadow delete failed') + }, + }), + }, + } + + try { + setupSavedPackage('bounded') + mockModule.runBundledModuleWithRegistry.mockImplementation(async () => { + vi.setSystemTime(new Date('2026-07-05T12:00:05.000Z')) + return { result: { ok: true }, error: null } + }) + + const created = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: failingMeter.USER_METER, + } as Env) + const start = await postPackageService(created.instance, 'start') + expect(start.status).toBe(200) + expect(upserts[0]).toEqual([ + 'user-123', + 'package-1', + 'realtime-supervisor', + 'running', + '2026-07-05T12:00:00.000Z', + '2026-07-05T12:00:00.000Z', + ]) + await flushWaitUntilTasks(created.waitUntilTasks) + expect(upserts.at(-1)).toEqual([ + 'user-123', + 'package-1', + 'realtime-supervisor', + 'stopped', + null, + '2026-07-05T12:00:05.000Z', + ]) + expect(consoleWarn).toHaveBeenCalledWith( + 'package-service-user-meter-shadow-failed', + expect.any(Error), + ) + + const purge = await postPackageService(created.instance, 'purge') + expect(purge.status).toBe(200) + } finally { + vi.useRealTimers() + } +}) + +test('package service UserMeter shadows are isolated per owning user', async () => { + resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + const appDb = createNoopPackageServiceDb() + const meter = createInMemoryUserMeterEnv() + const otherBinding = { + ...serviceBinding, + userId: 'user-other', + packageId: 'package-2', + kodyId: '@scope/package-2', + } + + try { + setupSavedPackage('bounded') + mockModule.runBundledModuleWithRegistry.mockImplementation( + () => new Promise(() => {}), + ) + mockModule.getSavedPackageById.mockImplementation( + async (_db: unknown, input: { packageId: string }) => + input.packageId === 'package-2' + ? { + id: 'package-2', + sourceId: 'source-1', + kodyId: '@scope/package-2', + } + : savedPackage, + ) + mockModule.loadPackageSourceBySourceId.mockImplementation(async () => + createPackageSource('bounded'), + ) + + const first = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: meter.env.USER_METER, + } as Env) + const second = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: meter.env.USER_METER, + } as Env) + + await postPackageService(first.instance, 'start') + await postPackageService(second.instance, 'start', otherBinding) + + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ status: 'running' }) + expect( + packageServiceShadowRow( + meter, + 'user-other', + 'package-2', + 'realtime-supervisor', + ), + ).toMatchObject({ status: 'running' }) + expect(meter.packageServicesByUser.get('user-123')?.size).toBe(1) + expect(meter.packageServicesByUser.get('user-other')?.size).toBe(1) + expect( + meter.packageServicesByUser + .get('user-123') + ?.has('package-2\0realtime-supervisor'), + ).toBe(false) + } finally { + vi.useRealTimers() + } +}) + +test('package service restore shadows current D1 projection into UserMeter', async () => { + resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T14:00:00.000Z')) + const appDb = createNoopPackageServiceDb() + const meter = createInMemoryUserMeterEnv() + + try { + const restored = createPackageServiceState() + restored.persistedEntries.set('package-service-state', { + binding: serviceBinding, + autoStart: false, + mode: 'bounded', + timeoutMs: 30_000, + stopRequested: false, + currentRunId: null, + nextAlarmAt: null, + nextAlarmSource: null, + lastStartedAt: '2026-07-05T13:00:00.000Z', + lastStoppedAt: '2026-07-05T13:05:00.000Z', + status: 'error', + lastError: 'boom', + lastResult: null, + lastRunFinishedAt: '2026-07-05T13:05:00.000Z', + consecutiveFailureCount: 1, + }) + new PackageServiceInstance(restored.state, { + APP_DB: appDb, + USER_METER: meter.env.USER_METER, + } as Env) + await waitForRestoreState(restored.state) + + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ + status: 'error', + startedAt: null, + sourceUpdatedAt: '2026-07-05T14:00:00.000Z', + revision: 1, + }) + } finally { + vi.useRealTimers() + } +}) diff --git a/packages/worker/src/package-runtime/package-service.ts b/packages/worker/src/package-runtime/package-service.ts index 56e8bc9f67..a6da78bf01 100644 --- a/packages/worker/src/package-runtime/package-service.ts +++ b/packages/worker/src/package-runtime/package-service.ts @@ -3,6 +3,10 @@ import * as Sentry from '@sentry/cloudflare' import { DurableObject } from 'cloudflare:workers' import { z } from 'zod' import { createMcpCallerContext } from '#mcp/context.ts' +import { + userMeterNamespace, + userMeterRpc, +} from '#worker/entitlements/user-meter-client.ts' import { getPackageServiceEntryPath, listPackageServices, @@ -415,26 +419,80 @@ class PackageServiceInstanceBase extends DurableObject { } /** - * Best-effort D1 projection of service liveness for entitlement counting. - * Failures must never break the service path; stop/error/idle writes still - * attempt to clear `running` so quota is released when D1 is healthy. + * Best-effort D1 projection of service liveness for entitlement counting, + * plus an expand-phase UserMeter shadow. D1 remains sole count/discovery + * authority; shadow failures never break the service path. Stop/error/idle + * writes still attempt to clear `running` so quota is released when D1 is + * healthy. */ private async projectServiceStateToD1() { const binding = this.stateSnapshot.binding if (!binding) return + const status = projectPackageServiceStatus(this.stateSnapshot.status) + const startedAt = this.stateSnapshot.lastStartedAt + const updatedAt = new Date().toISOString() try { await upsertPackageServiceState({ db: this.env.APP_DB, userId: binding.userId, packageId: binding.packageId, serviceName: binding.serviceName, - status: projectPackageServiceStatus(this.stateSnapshot.status), - startedAt: this.stateSnapshot.lastStartedAt, - updatedAt: new Date().toISOString(), + status, + startedAt, + updatedAt, }) } catch { // Best-effort: D1 outages must not take down package services. } + await this.shadowPackageServiceStateToUserMeter({ + binding, + status, + startedAt, + sourceUpdatedAt: updatedAt, + }) + } + + /** Best-effort UserMeter shadow; optional when `USER_METER` is unbound. */ + private async shadowPackageServiceStateToUserMeter(input: { + binding: PackageServiceBindingState + status: PackageServiceProjectedStatus + startedAt: string | null + sourceUpdatedAt: string + }) { + if (!userMeterNamespace(this.env)) return + try { + const meter = userMeterRpc({ + env: this.env, + userId: input.binding.userId, + }) + await meter.upsertPackageServiceState({ + packageId: input.binding.packageId, + serviceName: input.binding.serviceName, + status: input.status, + startedAt: input.startedAt, + sourceUpdatedAt: input.sourceUpdatedAt, + }) + } catch (error) { + console.warn('package-service-user-meter-shadow-failed', error) + } + } + + private async deletePackageServiceShadow( + binding: PackageServiceBindingState, + ) { + if (!userMeterNamespace(this.env)) return + try { + const meter = userMeterRpc({ + env: this.env, + userId: binding.userId, + }) + await meter.deletePackageServiceState({ + packageId: binding.packageId, + serviceName: binding.serviceName, + }) + } catch (error) { + console.warn('package-service-user-meter-shadow-failed', error) + } } private async deleteProjectedServiceState( @@ -450,6 +508,7 @@ class PackageServiceInstanceBase extends DurableObject { } catch { // Best-effort cleanup on purge. } + await this.deletePackageServiceShadow(binding) } private async ensureRunningHeartbeat() { @@ -956,8 +1015,9 @@ class PackageServiceInstanceBase extends DurableObject { await this.ctx.storage.deleteAlarm().catch(() => { // Best effort cleanup before deleteAll. }) - await this.ctx.storage.deleteAll() + // Clear D1 + UserMeter shadow before deleteAll. await this.deleteProjectedServiceState(binding) + await this.ctx.storage.deleteAll() return Response.json({ ok: true, }) diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index 8452d2d284..2eac2d4ec2 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -1,13 +1,24 @@ import { isDailyEntitlementResource, + isUserMeterPackageServiceStatus, userMeterMirrorUpdatedAtToken, + userMeterPackageServiceStateStaleMs, type DailyEntitlementResource, + type UserMeterPackageServiceState, + type UserMeterPackageServiceStatus, type UserMeterStorageBytesState, } from '#worker/entitlements/user-meter-do.ts' import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' type MeterRow = { count: number; revision: number } type StorageRow = { bytes: number; revision: number; updatedAt: string } +type PackageServiceRow = { + status: UserMeterPackageServiceStatus + startedAt: string | null + sourceUpdatedAt: string + revision: number + updatedAt: string +} /** * In-memory UserMeter stub keyed by `idFromName` (stable userId) for node-unit @@ -17,16 +28,31 @@ type StorageRow = { bytes: number; revision: number; updatedAt: string } export function createInMemoryUserMeterEnv() { const metersByUser = new Map>() const storageByUser = new Map() + const packageServicesByUser = new Map< + string, + Map + >() function counterKey(resource: string, day: string) { return `${resource}\0${day}` } + function packageServiceKey(packageId: string, serviceName: string) { + return `${packageId}\0${serviceName}` + } + function meterFor(userId: string) { const existingRows = metersByUser.get(userId) const rows = existingRows ?? new Map() if (!existingRows) metersByUser.set(userId, rows) + const existingPackageServices = packageServicesByUser.get(userId) + const packageServices = + existingPackageServices ?? new Map() + if (!existingPackageServices) { + packageServicesByUser.set(userId, packageServices) + } + function readRow(resource: string, day: string) { return rows.get(counterKey(resource, day)) ?? null } @@ -49,6 +75,36 @@ export function createInMemoryUserMeterEnv() { } } + function packageServiceState( + packageId: string, + serviceName: string, + row: PackageServiceRow, + ): UserMeterPackageServiceState { + return { + packageId, + serviceName, + status: row.status, + startedAt: row.status === 'running' ? row.startedAt : null, + sourceUpdatedAt: row.sourceUpdatedAt, + revision: row.revision, + updatedAt: row.updatedAt, + mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), + } + } + + function listPackageServiceStatesSorted() { + return [...packageServices.entries()] + .map(([entryKey, row]) => { + const [packageId, serviceName] = entryKey.split('\0') + return packageServiceState(packageId!, serviceName!, row) + }) + .sort((left, right) => { + const byPackage = left.packageId.localeCompare(right.packageId) + if (byPackage !== 0) return byPackage + return left.serviceName.localeCompare(right.serviceName) + }) + } + return { async initialize(input: { resource: string @@ -177,9 +233,171 @@ export function createInMemoryUserMeterEnv() { storageByUser.set(userId, next) return { ...storageReady(next), created: false } }, + async upsertPackageServiceState(input: { + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }) { + if (!isUserMeterPackageServiceStatus(input.status)) { + throw new Error( + `UserMeter package service status must be running, idle, stopped, or error; got ${JSON.stringify(input.status)}.`, + ) + } + const status = input.status + const key = packageServiceKey(input.packageId, input.serviceName) + const existing = packageServices.get(key) + if (existing && existing.sourceUpdatedAt > input.sourceUpdatedAt) { + return { + applied: false, + created: false, + state: packageServiceState( + input.packageId, + input.serviceName, + existing, + ), + } + } + const startedAt = + status === 'running' && + typeof input.startedAt === 'string' && + input.startedAt.length > 0 + ? input.startedAt + : null + const updatedAt = + typeof input.updatedAt === 'string' && input.updatedAt.length > 0 + ? input.updatedAt + : input.sourceUpdatedAt + const row: PackageServiceRow = { + status, + startedAt, + sourceUpdatedAt: input.sourceUpdatedAt, + revision: existing ? existing.revision + 1 : 1, + updatedAt, + } + packageServices.set(key, row) + return { + applied: true, + created: !existing, + state: packageServiceState(input.packageId, input.serviceName, row), + } + }, + async deletePackageServiceState(input: { + packageId: string + serviceName: string + }) { + const deleted = packageServices.delete( + packageServiceKey(input.packageId, input.serviceName), + ) + return { deleted } + }, + async listPackageServiceStates( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ) { + const pageSize = + typeof input.pageSize === 'number' && Number.isFinite(input.pageSize) + ? Math.min(Math.max(Math.trunc(input.pageSize), 1), 500) + : 100 + const all = listPackageServiceStatesSorted() + let startIndex = 0 + if ( + typeof input.startAfter === 'string' && + input.startAfter.length > 0 + ) { + try { + const parsed = JSON.parse(input.startAfter) as unknown + if ( + Array.isArray(parsed) && + parsed.length === 2 && + typeof parsed[0] === 'string' && + typeof parsed[1] === 'string' + ) { + const packageId = parsed[0] + const serviceName = parsed[1] + startIndex = + all.findIndex( + (row) => + row.packageId === packageId && + row.serviceName === serviceName, + ) + 1 + } + } catch { + startIndex = 0 + } + } + const page = all.slice(startIndex, startIndex + pageSize) + const truncated = startIndex + pageSize < all.length + const last = page[page.length - 1] + return { + states: page, + nextStartAfter: + truncated && last + ? JSON.stringify([last.packageId, last.serviceName]) + : null, + truncated, + } + }, + async countRunningPackageServices( + input: { + staleAfterMs?: number + excludeService?: { packageId: string; serviceName: string } + now?: string + } = {}, + ) { + const now = input.now ? new Date(input.now) : new Date() + const staleAfterMs = + typeof input.staleAfterMs === 'number' && + Number.isFinite(input.staleAfterMs) && + input.staleAfterMs >= 0 + ? Math.trunc(input.staleAfterMs) + : userMeterPackageServiceStateStaleMs + const freshAfter = new Date(now.valueOf() - staleAfterMs).toISOString() + let count = 0 + for (const [entryKey, row] of packageServices) { + if (row.status !== 'running') continue + if (row.sourceUpdatedAt < freshAfter) continue + if (input.excludeService) { + const [packageId, serviceName] = entryKey.split('\0') + if ( + packageId === input.excludeService.packageId && + serviceName === input.excludeService.serviceName + ) { + continue + } + } + count += 1 + } + return { count } + }, + async bootstrapPackageServiceStates(input: { + states: ReadonlyArray<{ + packageId: string + serviceName: string + status: string + startedAt?: string | null + sourceUpdatedAt: string + updatedAt?: string + }> + }) { + let applied = 0 + let skipped = 0 + for (const state of input.states) { + const result = await this.upsertPackageServiceState(state) + if (result.applied) applied += 1 + else skipped += 1 + } + return { applied, skipped } + }, async purge() { rows.clear() storageByUser.delete(userId) + packageServices.clear() + packageServicesByUser.delete(userId) return { ok: true as const } }, async exportCounters( @@ -199,11 +417,9 @@ export function createInMemoryUserMeterEnv() { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(row.revision), } }) - const includeStorageShadow = + const includeShadows = typeof input.startAfter !== 'string' || input.startAfter.length === 0 - const storage = includeStorageShadow - ? storageByUser.get(userId) - : undefined + const storage = includeShadows ? storageByUser.get(userId) : undefined const storageBytesShadow: UserMeterStorageBytesState | null = storage ? { bytes: storage.bytes, @@ -212,9 +428,13 @@ export function createInMemoryUserMeterEnv() { mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(storage.revision), } : null + const packageServiceStatesShadow = includeShadows + ? listPackageServiceStatesSorted() + : null return { counters, storageBytesShadow, + packageServiceStatesShadow, nextStartAfter: null, truncated: false, } @@ -233,6 +453,7 @@ export function createInMemoryUserMeterEnv() { env, metersByUser, storageByUser, + packageServicesByUser, async seed(input: { userId: string resource: DailyEntitlementResource From 9f7be032638d4f49cd1e606ea8dba379bf584a91 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 07:49:02 +0000 Subject: [PATCH 17/39] docs(services): inventory liveness shadow Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 94 ++++++++++------ .../contributing/architecture/entitlements.md | 104 +++++++++++++++--- .../contributing/architecture/primitives.yaml | 4 +- .../worker/src/account/export.node.test.ts | 21 +++- packages/worker/src/account/export.ts | 14 ++- .../worker/src/account/user-owned-surfaces.ts | 2 +- 6 files changed, 186 insertions(+), 53 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 7bd1099f08..26b49846c8 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -166,15 +166,15 @@ migration-safe chunked interface: with `section: "storage_runner"` and a `storage_id`, using the same StorageRunner `exportStorage({ pageSize, startAfter })` RPC as the dedicated storage export capability. User meter counters use `section: "user_meter"` and - the `UserMeter.exportCounters` RPC (daily counters plus additive - `storageBytesShadow` on the first page only when present; explicitly - non-authoritative). R2 raw MIME, attachment, avatar, and icon objects use - `section: "r2_object"`; each response contains at most one 256 KiB base64 - chunk and an opaque cursor. Each request uses bounded `LIMIT 1` ownership - queries rather than reconstructing inventory. Continuation cursors bind the - source row, object key, size, and ETag; ownership/key mutations and object - overwrites are reported instead of mixing generations. Missing objects are - represented explicitly. + the `UserMeter.exportCounters` RPC (daily counters plus additive shadow fields + on the first page only when present: `storageBytesShadow` and + `packageServiceStatesShadow`; explicitly non-authoritative). R2 raw MIME, + attachment, avatar, and icon objects use `section: "r2_object"`; each response + contains at most one 256 KiB base64 chunk and an opaque cursor. Each request + uses bounded `LIMIT 1` ownership queries rather than reconstructing inventory. + Continuation cursors bind the source row, object key, size, and ETag; + ownership/key mutations and object overwrites are reported instead of mixing + generations. Missing objects are represented explicitly. D1 manifest counts use bounded SQL `COUNT(*)` queries. D1 section rows are read with SQL-level keyset pagination: every query orders by the table's `rowid`, @@ -198,13 +198,15 @@ Durable Object export behavior: never pruned by retention. See [Run records](./run-records.md). - `UserMeter` exports daily entitlement counter rows through the `user_meter` section (`exportCounters` RPC; keyset pagination by UTC `day` and `resource`). - The same RPC may return additive `storageBytesShadow` on the first page only - (`startAfter` absent) when the schema-v4 shadow row exists (`null` on later - pages and when never shadowed); section totals count it as one row when - present, but it is explicitly **non-authoritative** — usage and enforcement - read D1 `users.d1_storage_bytes`. Retention is self-enforced inside the DO - (seven UTC days of counter and inbound-delivery-claim rows); shadow - storage-byte state is not time-pruned. See + The same RPC may return additive shadow fields on the first page only + (`startAfter` absent): `storageBytesShadow` when the schema-v4 row exists, and + `packageServiceStatesShadow` when schema-v5 service rows exist (`null` on + later pages and when never shadowed). Section totals count each shadow + inventory once when present, but both are explicitly **non-authoritative** — + usage and enforcement read D1 (`users.d1_storage_bytes` and + `package_service_states` respectively). Retention is self-enforced inside the + DO (seven UTC days of counter and inbound-delivery-claim rows); shadow + storage-byte and package-service liveness rows are not time-pruned. See [Entitlements](./entitlements.md#usermeter-expand-phase). - `RemoteConnectorSession` exposes persisted connector metadata and tool descriptors through an export RPC. @@ -248,8 +250,9 @@ The schema is defined by migrations in `packages/worker/migrations/`: [Platform accounts](./platform-accounts.md)). `d1_storage_bytes` and `d1_storage_bytes_updated_at` (migration 0122) are the **sole authority** for D1 payload storage-byte read, enforcement, and reconciliation. UserMeter - `storage_bytes_state` (schema v4) is an optional expand-phase shadow only — - see [Entitlements](./entitlements.md#usermeter-expand-phase). Inbound email + `storage_bytes_state` (schema v4) and `package_service_states` (schema v5) are + optional expand-phase shadows only — see + [Entitlements](./entitlements.md#usermeter-expand-phase). Inbound email routing does not reverse-resolve stable ids at all — it uses the indexed username lookup (`findPublicUserIdentityByUsername`). Contextless paths resolve stable ids with one indexed point read on `users.stable_user_id` (for @@ -279,8 +282,12 @@ The schema is defined by migrations in `packages/worker/migrations/`: [Run records](./run-records.md)). Execution history rows live in the same DO. - `package_service_states` (`0095-package-service-states.sql`): authoritative per-service liveness projection (`running` / `idle` / `stopped` / `error`) for - entitlement concurrency. Upserted and heartbeaten by the package-service - Durable Object; not derived from run history. + entitlement concurrency, discovery, and export/deletion inventory. Upserted + and heartbeaten (1h) by the `PackageServiceInstance` Durable Object; running + counts treat rows stale after 24h without a fresh heartbeat. Not derived from + run history. Expand-phase slice 4 Phase A also best-effort shadows each row + into the per-user `UserMeter` DO (schema v5); D1 remains sole authority in + that slice — see [Entitlements](./entitlements.md#usermeter-expand-phase). - `entity_sources`: durable mapping from user-facing entities to Artifacts repos and their latest published commit - `saved_packages`: package metadata/search projection derived from published @@ -502,10 +509,12 @@ Daily rate-style entitlement counters and inbound email delivery-id idempotency live in a per-user `UserMeter` Durable Object with SQLite (`packages/worker/src/entitlements/user-meter-do.ts`). Schema v4 adds an optional `storage_bytes_state` singleton as a **best-effort shadow** of D1 -`users.d1_storage_bytes` for future cutover — D1 remains sole authority for -reads, reserves, and reconciliation in the current additive slice. The Worker -binding is `USER_METER` (class `UserMeter`; Wrangler SQLite migration tag `v21` -via `new_sqlite_classes` in `packages/worker/wrangler.jsonc`). +`users.d1_storage_bytes`. Schema v5 adds an optional `package_service_states` +table as a **best-effort shadow** of D1 `package_service_states` for future +cutover — D1 remains sole authority for reads, running counts, discovery, and +`service_start` enforcement in expand-phase slice 4 Phase A. The Worker binding +is `USER_METER` (class `UserMeter`; Wrangler SQLite migration tag `v21` via +`new_sqlite_classes` in `packages/worker/wrangler.jsonc`). Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP `userId` via `userMeterDurableObjectName(userId)` → `idFromName(userId)` in @@ -513,7 +522,7 @@ Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP column inside the DO because the object identity is the user. SQLite ownership (schema version tracked in `user_meter_meta`; current version -**4**): +**5**): - `daily_counters` — authoritative UTC-day counters for `email_sends_per_day`, `email_receives_per_day`, `execute_calls_per_day`, and @@ -530,12 +539,23 @@ SQLite ownership (schema version tracked in `user_meter_meta`; current version and optional reconcile shadows; never read for enforcement or usage display. StorageRunner bucket estimates stay outside this row (see [Entitlements](./entitlements.md#usermeter-expand-phase)). +- `package_service_states` — per-service **shadow** of D1 liveness rows + (`package_id`, `service_name`, `status`, `started_at`, `source_updated_at`, + monotonic `revision`, `updated_at`; primary key `(package_id, service_name)`). + Added in schema v5. Populated by best-effort dual-writes from + `PackageServiceInstance` on every D1 projection/delete; monotonic on + `source_updated_at`. Never read for enforcement, running counts, discovery, or + usage display in expand-phase slice 4 Phase A. Cutover-support RPCs + (`listPackageServiceStates`, `countRunningPackageServices`, + `bootstrapPackageServiceStates`) mirror D1 semantics for future parity review + only. Retention is self-enforced inside the DO: every read/write path opportunistically deletes counter and claim rows older than seven UTC days (`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; the window covers timezone edge cases, recent account exports, and inbound -retries. Shadow storage-byte state is not time-pruned. +retries. Shadow storage-byte and package-service liveness rows are not +time-pruned. **Expand-phase D1 mirrors (daily counters only):** enforcement and point reads are authoritative in UserMeter for daily counters. D1 @@ -554,11 +574,12 @@ writes cannot overwrite newer state. See daily paths never read D1 for enforcement. Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan; -`deleteAll` clears counters, claims, and any shadow storage state). Account -export pages `UserMeter.exportCounters` through the `user_meter` manifest -section / `account_export_section` (daily counters plus additive -`storageBytesShadow` on the first page only when present; shadow field is -non-authoritative). +`deleteAll` clears counters, claims, and all shadow state including storage +bytes and package-service liveness). Account export pages +`UserMeter.exportCounters` through the `user_meter` manifest section / +`account_export_section` (daily counters plus additive `storageBytesShadow` and +`packageServiceStatesShadow` on the first page only when present; shadow fields +are non-authoritative). ### Package state model @@ -574,7 +595,9 @@ storage homes as follows: `package:{encodeURIComponent(packageId)}` via `buildPackageStorageId` / `packageStorage()`. Shared durable data for every package surface. - **Package coordination** — `PackageServiceInstance` DO holds lifecycle and - alarms only; durable data stays in package storage. App facets and + alarms only; durable data stays in package storage. Each lifecycle projection + dual-writes D1 `package_service_states` (authority) and a best-effort + UserMeter shadow (expand-phase slice 4 Phase A). App facets and package-internal DO namespaces are extra StorageRunner buckets under the package id, not a general actor model. - **Package jobs** — schedule metadata in D1 `jobs`; run-local scratch in @@ -597,8 +620,8 @@ via `durableObjectNameFromParts`); domain helpers such as package activation counters/milestones). See [Run records](./run-records.md). - `UserMeter` — `userMeterDurableObjectName(userId)` → `idFromName(userId)`. One daily-entitlement meter DO per user (untrimmed stable id, same as `RunLog`), - plus optional schema-v4 D1 storage-byte shadow. See - [Entitlements](./entitlements.md#usermeter-expand-phase). + plus optional schema-v4 D1 storage-byte shadow and schema-v5 package-service + liveness shadow. See [Entitlements](./entitlements.md#usermeter-expand-phase). - `McpClientHub` — `mcpClientHubDurableObjectName(userId)` → `idFromName(userId.trim())`. - `StorageRunner` — `storageRunnerDurableObjectName(userId, storageId)` → @@ -913,7 +936,8 @@ to `durableObjectNameFromParts`). - `JobManager`: `idFromName(userId)` (no trim). - `RunLog`: `idFromName(userId)` (no trim); one execution-history DO per user. - `UserMeter`: `idFromName(userId)` (no trim); one daily-entitlement meter DO - per user, plus optional schema-v4 D1 storage-byte shadow. + per user, plus optional schema-v4 D1 storage-byte shadow and schema-v5 + package-service liveness shadow. - `McpClientHub`: `idFromName(userId.trim())`. - `StorageRunner`: `idFromName(JSON.stringify([userId, storageId]))`. - `RepoSession`: `idFromName(repo_sessions.id)`; the key is not user-prefixed, diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 7171beca42..58fa5e601a 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -124,6 +124,12 @@ storage layout and naming are documented in [Data storage](./data-storage.md). `storage_bytes_state` singleton as a **best-effort shadow** for future cutover — it does not drive reads, reserves, or reconciliation in this additive slice. +**D1 package service liveness** (`package_service_states`) stays authoritative +for running counts, discovery, and `service_start` enforcement. UserMeter schema +v5 adds an optional per-service shadow table as **best-effort future-cutover +support only** — see +[Package service liveness — UserMeter shadow](#package-service-liveness--usermeter-shadow-expand-phase-slice-4-phase-a). + StorageRunner bucket `estimatedBytes` and the per-bucket inventory in `user_storage_buckets` stay a **separate** quota component. StorageRunner write chokepoints pass `getCurrent` as a check-only composed total (D1 payload bytes @@ -214,11 +220,78 @@ failures are logged and never fail or poison the lane. Current `readUserD1StorageBytes` only. **Account export and purge:** `UserMeter.exportCounters` may return additive -non-authoritative `storageBytesShadow` on the first page only (`startAfter` -absent) when the shadow row exists; subsequent pages return `null` so paged -consumers never double-count it (still counted once in the `user_meter` section -total). `UserMeter.purge()` clears counters, inbound delivery claims, and any -shadow storage state via `deleteAll`. +non-authoritative shadow fields on the first page only (`startAfter` absent): +`storageBytesShadow` when the schema-v4 row exists, and +`packageServiceStatesShadow` when schema-v5 service rows exist. Subsequent pages +return `null` for each shadow so paged consumers never double-count them +(section totals still count each shadow inventory once when present). +`UserMeter.purge()` clears counters, inbound delivery claims, and all shadow +state (storage bytes and package-service liveness) via `deleteAll`. + +### Package service liveness — UserMeter shadow (expand phase slice 4, Phase A) + +D1 `package_service_states` remains the **sole authority** for running-service +**count**, **discovery**, and **`service_start` enforcement** in this PR. +Nothing in Phase A switches those reads or the `assertWithinEntitlement` path +for `package_services` / `persistent_package_services`. + +UserMeter schema **v5** adds a per-service `package_service_states` table inside +the DO as **best-effort shadow / future-cutover support only** (`status`, +`started_at`, monotonic `source_updated_at` from the D1 projection timestamp, +`revision`, `updated_at`). User scope is the DO identity — there is no `user_id` +column. Shadow rows are never read for usage display, entitlement enforcement, +or account-deletion inventory in this slice. + +**Dual-write from `PackageServiceInstance`:** every D1 projection also attempts +a best-effort UserMeter shadow on the same lifecycle surface: + +- lifecycle transitions and warm-start restore after upgrades + (`projectServiceStateToD1`) +- running-service heartbeat alarms (1h `packageServiceStateHeartbeatMs`, + unchanged) +- stop, error, and idle projections that clear `running` +- purge (`deleteProjectedServiceState` deletes D1 then shadow before + `deleteAll`) + +D1 upsert/delete runs first; shadow RPCs are optional when `USER_METER` is +unbound and failures log `package-service-user-meter-shadow-failed` without +affecting the service path. Shadow upserts reject stale/out-of-order writes when +`sourceUpdatedAt` is older than the existing shadow row so cold bootstrap cannot +clobber fresher state. + +**Timing unchanged:** live services heartbeat D1 `updated_at` every **1 hour** +(`packageServiceStateHeartbeatMs`). Running counts still treat rows as stale +after **24 hours** without a fresh heartbeat (`packageServiceStateStaleMs` in +`entitlements/service.ts`). The UserMeter cutover-support RPC +`countRunningPackageServices` uses the same 24h window on shadow +`source_updated_at` but is **not** wired to enforcement in Phase A. + +**Account export:** `UserMeter.exportCounters` returns additive +`packageServiceStatesShadow` on the first page only (`startAfter` absent); later +pages return `null`. Section totals count the shadow inventory once when +present; the field is explicitly non-authoritative — authoritative liveness +remains on D1. + +**Account purge:** `UserMeter.purge()` clears package-service shadow rows with +the rest of DO state via `deleteAll`. + +### Future package-service authority flip (contract follow-up) + +A separate **high-risk contract PR** — not a merge blocker for Phase A — will +flip running-service count/discovery/enforcement into UserMeter only after: + +1. at least one full **24h stale-window soak** with shadow/D1 parity review, and +2. a **cold-bootstrap design** for accounts whose DO shadow is empty while D1 + still holds rows (`bootstrapPackageServiceStates` / equivalent). + +Until that flip, shadow divergence is acceptable; D1 remains the contract. **D1 +likely stays the enumeration index** for account export, deletion, and admin +discovery until an alternate inventory exists — UserMeter would become the +running-count authority first, not a wholesale replacement for every D1 reader. + +**Remaining expand roadmap:** slice 4 Phase A (this shadow slice) is additive +only; slice 5 — account-deletion write fencing — follows independently. Storage +authority flip remains the separate contract follow-up above. ### Future storage authority flip (contract follow-up) @@ -234,9 +307,10 @@ Only then do reads, reserves, and reconciliation switch to UserMeter-first with D1 as mirror/cursor. Until that flip, shadow divergence is acceptable; D1 remains the contract. -**Remaining UserMeter expand roadmap:** slice 4 — `package_service_states` -running counts (service liveness); slice 5 — account-deletion write fencing. -Storage authority flip is tracked separately as the contract follow-up above. +**Remaining UserMeter expand roadmap:** slice 4 Phase A — +`package_service_states` UserMeter shadow (this PR; D1 authority unchanged); +slice 5 — account-deletion write fencing. Package-service and storage authority +flips are separate high-risk contract follow-ups after soak/parity review. **Daily-counter mirror retirement:** dropping D1 `entitlement_daily_counters` waits until reporting-off-D1 work merges and mirror parity is verified in @@ -382,12 +456,14 @@ Rules: running package services) are counted **directly from their source D1 tables at the enforcement point** via built-in counters in `service.ts`. They do not depend on any metering or rollup tables. Running package services are counted - from `package_service_states` (status `running` and freshly heartbeaten), not - from run-history rows — see [Run records](./run-records.md) - (`state-vs-history`). **Concurrent workflows** are authoritative in per-user - RunLog `workflow_projections`: create reserves atomically via - `reserveWorkflowProjectionSlot`, and usage readers call - `countActiveWorkflowProjections` through + from D1 `package_service_states` (status `running` and freshly heartbeaten; 1h + heartbeat, 24h staleness), not from run-history rows — see + [Run records](./run-records.md) (`state-vs-history`). Expand-phase slice 4 + Phase A dual-writes the same projection into UserMeter as a non-authoritative + shadow; enforcement and `service_start` still read D1 only. **Concurrent + workflows** are authoritative in per-user RunLog `workflow_projections`: + create reserves atomically via `reserveWorkflowProjectionSlot`, and usage + readers call `countActiveWorkflowProjections` through `readCurrentEntitlementResourceUsage`. Expand-phase D1 `workflow_runs` is a compatibility mirror only. - **Rate-style limits** (email sends/receives per day, execute calls per day, diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index 5302603cd6..e66f4a2402 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -514,7 +514,9 @@ primitives: name: User meter summary: Per-user UserMeter DO SQLite for daily entitlement counters, inbound - delivery idempotency, and expand-phase D1 storage-byte shadow. + delivery idempotency, expand-phase D1 storage-byte shadow, and + expand-phase package-service liveness shadow (D1 authority unchanged in + slice 4 Phase A). code: - packages/worker/src/entitlements/user-meter-do.ts - packages/worker/src/entitlements/user-meter-client.ts diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 815d16f8cb..d1850823d2 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1578,6 +1578,18 @@ test('account export includes user_meter counters, pages them, and warns on trun updatedAt: '2026-07-31T03:00:00.000Z', mirrorUpdatedAt: 'r/00000000000000000003', } + const packageServiceStatesShadow = [ + { + packageId: 'pkg-1', + serviceName: 'worker', + status: 'running' as const, + startedAt: '2026-07-31T03:00:00.000Z', + sourceUpdatedAt: '2026-07-31T03:05:00.000Z', + revision: 2, + updatedAt: '2026-07-31T03:05:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000002', + }, + ] const exportCounters = vi.fn( async (input: { pageSize?: number; startAfter?: string | null }) => { const pageSize = input.pageSize ?? 100 @@ -1593,6 +1605,9 @@ test('account export includes user_meter counters, pages them, and warns on trun return { counters: page, storageBytesShadow: isFirstPage ? storageBytesShadow : null, + packageServiceStatesShadow: isFirstPage + ? packageServiceStatesShadow + : null, nextStartAfter: truncated ? `${page.at(-1)!.day}:${page.at(-1)!.resource}` : null, @@ -1615,10 +1630,11 @@ test('account export includes user_meter counters, pages them, and warns on trun mcpUserId: 'user-aaa', }) expect(idFromName).toHaveBeenCalledWith('user-aaa') - expect(accountExport.manifest.sections.user_meter?.count).toBe(4) + expect(accountExport.manifest.sections.user_meter?.count).toBe(5) expect(accountExport.durableObjects.userMeter).toEqual({ counters, storageBytesShadow, + packageServiceStatesShadow, nextStartAfter: null, truncated: false, }) @@ -1637,6 +1653,7 @@ test('account export includes user_meter counters, pages them, and warns on trun }) expect(first.items).toEqual(counters.slice(0, 2)) expect(first.storageBytesShadow).toEqual(storageBytesShadow) + expect(first.packageServiceStatesShadow).toEqual(packageServiceStatesShadow) expect(first.truncated).toBe(true) expect(first.nextStartAfter).toBe('2026-07-30:execute_calls_per_day') @@ -1650,6 +1667,7 @@ test('account export includes user_meter counters, pages them, and warns on trun }) expect(second.items).toEqual(counters.slice(2)) expect(second.storageBytesShadow).toBeNull() + expect(second.packageServiceStatesShadow).toBeNull() expect(second.truncated).toBe(false) expect(second.nextStartAfter).toBeNull() expect(exportCounters).toHaveBeenCalledWith( @@ -1662,6 +1680,7 @@ test('account export includes user_meter counters, pages them, and warns on trun exportCounters.mockImplementation(async () => ({ counters: [counters[0]!], storageBytesShadow: null, + packageServiceStatesShadow: null, nextStartAfter: 'cursor-more', truncated: true, })) diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index a2ac581228..3ceb559e36 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -205,7 +205,13 @@ export type AccountExportArtifactRepo = { type RunRecordsExportPayload = Awaited> function countUserMeterExportEntries(result: UserMeterExportResult): number { - return result.counters.length + (result.storageBytesShadow == null ? 0 : 1) + return ( + result.counters.length + + (result.storageBytesShadow == null ? 0 : 1) + + (result.packageServiceStatesShadow == null + ? 0 + : result.packageServiceStatesShadow.length) + ) } function countRunRecordsExportEntries( @@ -288,6 +294,11 @@ export type AccountExportSectionResult = { * first `user_meter` page (`startAfter` absent); later pages omit it. */ storageBytesShadow?: UserMeterExportResult['storageBytesShadow'] + /** + * Non-authoritative UserMeter package-service shadow rows. Present only on + * the first `user_meter` page (`startAfter` absent); later pages omit it. + */ + packageServiceStatesShadow?: UserMeterExportResult['packageServiceStatesShadow'] } function normalizePageSize(pageSize: number | undefined) { @@ -1921,6 +1932,7 @@ export async function readAccountExportSection(input: { section: input.section, items: page.counters, storageBytesShadow: page.storageBytesShadow, + packageServiceStatesShadow: page.packageServiceStatesShadow, truncated: page.truncated, nextStartAfter: page.nextStartAfter, pageSize, diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 47a46d6fc3..d0d751617f 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -130,7 +130,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray Date: Sat, 1 Aug 2026 07:58:52 +0000 Subject: [PATCH 18/39] perf(services): defer liveness shadow writes Co-authored-by: Kent C. Dodds --- .../package-service.node.test.ts | 132 ++++++++++++++++++ .../src/package-runtime/package-service.ts | 44 ++++-- 2 files changed, 167 insertions(+), 9 deletions(-) diff --git a/packages/worker/src/package-runtime/package-service.node.test.ts b/packages/worker/src/package-runtime/package-service.node.test.ts index 7aa921d924..66fc1f98a5 100644 --- a/packages/worker/src/package-runtime/package-service.node.test.ts +++ b/packages/worker/src/package-runtime/package-service.node.test.ts @@ -179,6 +179,37 @@ async function flushWaitUntilTasks(waitUntilTasks: Array>) { await Promise.all(waitUntilTasks) } +/** + * Drain UserMeter shadow (and other promptly settling) waitUntil tasks without + * hanging on long-lived background service runs also scheduled via waitUntil. + */ +async function drainSettledWaitUntilTasks( + waitUntilTasks: Array>, + fromIndex = 0, +) { + const pending = waitUntilTasks.slice(fromIndex) + await Promise.all( + pending.map(async (task) => { + const state = { settled: false } + void Promise.resolve(task).then( + () => { + state.settled = true + }, + () => { + state.settled = true + }, + ) + for (let i = 0; i < 40; i++) { + await Promise.resolve() + if (state.settled) { + await task + return + } + } + }), + ) +} + function resetMocks() { mockModule.getSavedPackageById.mockReset() mockModule.loadPackageSourceBySourceId.mockReset() @@ -890,6 +921,9 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', } as Env) const start = await postPackageService(created.instance, 'start') expect(start.status).toBe(200) + // D1 projection stays on the awaited lifecycle path. + expect(upserts[0]?.[3]).toBe('running') + await drainSettledWaitUntilTasks(created.waitUntilTasks) expect( packageServiceShadowRow( meter, @@ -906,7 +940,12 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', expect(created.getAlarmAt()).toBe(Date.parse('2026-07-05T13:00:00.000Z')) vi.setSystemTime(new Date('2026-07-05T13:00:00.000Z')) + const afterStartWaitUntilCount = created.waitUntilTasks.length await created.instance.alarm() + await drainSettledWaitUntilTasks( + created.waitUntilTasks, + afterStartWaitUntilCount, + ) expect( packageServiceShadowRow( meter, @@ -921,7 +960,12 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', }) expect(created.getAlarmAt()).toBe(Date.parse('2026-07-05T14:00:00.000Z')) + const afterHeartbeatWaitUntilCount = created.waitUntilTasks.length await postPackageService(created.instance, 'stop') + await drainSettledWaitUntilTasks( + created.waitUntilTasks, + afterHeartbeatWaitUntilCount, + ) expect( packageServiceShadowRow( meter, @@ -935,7 +979,12 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', revision: 3, }) + const afterStopWaitUntilCount = created.waitUntilTasks.length await postPackageService(created.instance, 'purge') + await drainSettledWaitUntilTasks( + created.waitUntilTasks, + afterStopWaitUntilCount, + ) expect( packageServiceShadowRow( meter, @@ -955,6 +1004,77 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', } }) +test('package service lifecycle does not await a gated UserMeter shadow', async () => { + resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + const { appDb, upserts } = createRecordingPackageServiceDb() + let releaseShadow!: () => void + const shadowGate = new Promise((resolve) => { + releaseShadow = resolve + }) + let shadowFinished = false + const gatedMeter = { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: () => ({ + async upsertPackageServiceState() { + await shadowGate + shadowFinished = true + return { + applied: true, + created: true, + state: { + packageId: 'package-1', + serviceName: 'realtime-supervisor', + status: 'running' as const, + startedAt: '2026-07-05T12:00:00.000Z', + sourceUpdatedAt: '2026-07-05T12:00:00.000Z', + revision: 1, + updatedAt: '2026-07-05T12:00:00.000Z', + mirrorUpdatedAt: 'r/00000000000000000001', + }, + } + }, + async deletePackageServiceState() { + return { deleted: true } + }, + }), + }, + } + + try { + setupSavedPackage('persistent') + mockModule.runBundledModuleWithRegistry.mockImplementation( + () => new Promise(() => {}), + ) + + const created = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: gatedMeter.USER_METER, + } as Env) + const start = await postPackageService(created.instance, 'start') + expect(start.status).toBe(200) + expect(upserts[0]).toEqual([ + 'user-123', + 'package-1', + 'realtime-supervisor', + 'running', + '2026-07-05T12:00:00.000Z', + '2026-07-05T12:00:00.000Z', + ]) + // Lifecycle returned while the shadow DO hop is still gated. + expect(shadowFinished).toBe(false) + expect(created.waitUntilTasks.length).toBeGreaterThan(0) + + releaseShadow() + await drainSettledWaitUntilTasks(created.waitUntilTasks) + expect(shadowFinished).toBe(true) + } finally { + vi.useRealTimers() + } +}) + test('package service UserMeter shadow failure cannot break D1 projection or lifecycle', async () => { resetMocks() silenceExpectedConsoleWarns(['package-service-user-meter-shadow-failed']) @@ -1010,8 +1130,17 @@ test('package service UserMeter shadow failure cannot break D1 projection or lif expect.any(Error), ) + const beforePurgeWaitUntilCount = created.waitUntilTasks.length const purge = await postPackageService(created.instance, 'purge') expect(purge.status).toBe(200) + await drainSettledWaitUntilTasks( + created.waitUntilTasks, + beforePurgeWaitUntilCount, + ) + expect(consoleWarn).toHaveBeenCalledWith( + 'package-service-user-meter-shadow-failed', + expect.any(Error), + ) } finally { vi.useRealTimers() } @@ -1060,6 +1189,8 @@ test('package service UserMeter shadows are isolated per owning user', async () await postPackageService(first.instance, 'start') await postPackageService(second.instance, 'start', otherBinding) + await drainSettledWaitUntilTasks(first.waitUntilTasks) + await drainSettledWaitUntilTasks(second.waitUntilTasks) expect( packageServiceShadowRow( @@ -1120,6 +1251,7 @@ test('package service restore shadows current D1 projection into UserMeter', asy USER_METER: meter.env.USER_METER, } as Env) await waitForRestoreState(restored.state) + await drainSettledWaitUntilTasks(restored.waitUntilTasks) expect( packageServiceShadowRow( diff --git a/packages/worker/src/package-runtime/package-service.ts b/packages/worker/src/package-runtime/package-service.ts index a6da78bf01..51a97ea2b0 100644 --- a/packages/worker/src/package-runtime/package-service.ts +++ b/packages/worker/src/package-runtime/package-service.ts @@ -420,10 +420,11 @@ class PackageServiceInstanceBase extends DurableObject { /** * Best-effort D1 projection of service liveness for entitlement counting, - * plus an expand-phase UserMeter shadow. D1 remains sole count/discovery - * authority; shadow failures never break the service path. Stop/error/idle - * writes still attempt to clear `running` so quota is released when D1 is - * healthy. + * plus a non-awaited expand-phase UserMeter shadow via `waitUntil`. D1 + * remains sole count/discovery authority and stays on the awaited path; + * shadow DO hops must not gate lifecycle/heartbeat responses. + * Stop/error/idle writes still attempt to clear `running` so quota is + * released when D1 is healthy. */ private async projectServiceStateToD1() { const binding = this.stateSnapshot.binding @@ -444,7 +445,7 @@ class PackageServiceInstanceBase extends DurableObject { } catch { // Best-effort: D1 outages must not take down package services. } - await this.shadowPackageServiceStateToUserMeter({ + this.schedulePackageServiceStateShadow({ binding, status, startedAt, @@ -452,14 +453,28 @@ class PackageServiceInstanceBase extends DurableObject { }) } - /** Best-effort UserMeter shadow; optional when `USER_METER` is unbound. */ - private async shadowPackageServiceStateToUserMeter(input: { + /** + * Schedule a best-effort UserMeter shadow upsert. Never awaited by + * lifecycle/heartbeat paths; the helper catches so `waitUntil` cannot + * reject. + */ + private schedulePackageServiceStateShadow(input: { binding: PackageServiceBindingState status: PackageServiceProjectedStatus startedAt: string | null sourceUpdatedAt: string }) { if (!userMeterNamespace(this.env)) return + this.ctx.waitUntil(this.shadowPackageServiceStateToUserMeter(input)) + } + + /** Best-effort UserMeter shadow upsert; catches/logs so waitUntil settles. */ + private async shadowPackageServiceStateToUserMeter(input: { + binding: PackageServiceBindingState + status: PackageServiceProjectedStatus + startedAt: string | null + sourceUpdatedAt: string + }) { try { const meter = userMeterRpc({ env: this.env, @@ -477,10 +492,21 @@ class PackageServiceInstanceBase extends DurableObject { } } - private async deletePackageServiceShadow( + /** + * Schedule a best-effort UserMeter shadow delete. Never awaited by purge; + * the helper catches so `waitUntil` cannot reject. + */ + private schedulePackageServiceShadowDelete( binding: PackageServiceBindingState, ) { if (!userMeterNamespace(this.env)) return + this.ctx.waitUntil(this.deletePackageServiceShadow(binding)) + } + + /** Best-effort UserMeter shadow delete; catches/logs so waitUntil settles. */ + private async deletePackageServiceShadow( + binding: PackageServiceBindingState, + ) { try { const meter = userMeterRpc({ env: this.env, @@ -508,7 +534,7 @@ class PackageServiceInstanceBase extends DurableObject { } catch { // Best-effort cleanup on purge. } - await this.deletePackageServiceShadow(binding) + this.schedulePackageServiceShadowDelete(binding) } private async ensureRunningHeartbeat() { From 12e7c0280d414b776b8be4c2446861646363181c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 08:26:49 +0000 Subject: [PATCH 19/39] fix(services): serialize liveness shadows Co-authored-by: Kent C. Dodds --- packages/worker/src/account/export.ts | 3 +- .../worker/src/entitlements/user-meter-do.ts | 20 +- .../entitlements/user-meter.workers.test.ts | 57 ++++++ .../package-service.node.test.ts | 181 ++++++++++++++++++ .../src/package-runtime/package-service.ts | 23 ++- 5 files changed, 277 insertions(+), 7 deletions(-) diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index 6614fd085d..b3a31eb93e 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -301,7 +301,8 @@ export type AccountExportSectionResult = { storageBytesShadow?: UserMeterExportResult['storageBytesShadow'] /** * Non-authoritative UserMeter package-service shadow rows. Present only on - * the first `user_meter` page (`startAfter` absent); later pages omit it. + * the first `user_meter` page (`startAfter` absent); later pages set it to + * `null`. */ packageServiceStatesShadow?: UserMeterExportResult['packageServiceStatesShadow'] } diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 0103decadf..12a7f82461 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -39,6 +39,9 @@ const maxPackageServiceIdLength = 256 const inboundReceiveResource = 'email_receives_per_day' satisfies DailyEntitlementResource const utcDayKeyPattern = /^\d{4}-\d{2}-\d{2}$/ +/** Matches `new Date().toISOString()` — required for lexicographic monotonicity. */ +const packageServiceSourceUpdatedAtPattern = + /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/ const packageServiceShadowStatuses = [ 'running', 'idle', @@ -286,9 +289,22 @@ function assertPackageServiceStatus( } function assertSourceUpdatedAt(sourceUpdatedAt: string): string { - if (typeof sourceUpdatedAt !== 'string' || sourceUpdatedAt.length === 0) { + if ( + typeof sourceUpdatedAt !== 'string' || + !packageServiceSourceUpdatedAtPattern.test(sourceUpdatedAt) + ) { + throw new Error( + 'UserMeter package service sourceUpdatedAt must be an ISO-8601 UTC timestamp.', + ) + } + // Pattern-first: avoid Invalid Date → RangeError from toISOString(). + const parsedMs = Date.parse(sourceUpdatedAt) + if ( + !Number.isFinite(parsedMs) || + new Date(parsedMs).toISOString() !== sourceUpdatedAt + ) { throw new Error( - 'UserMeter package service sourceUpdatedAt must be a non-empty string.', + 'UserMeter package service sourceUpdatedAt must be an ISO-8601 UTC timestamp.', ) } return sourceUpdatedAt diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index ccafdb5e59..a330f75f52 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -1215,3 +1215,60 @@ test('UserMeter package-service shadow upserts are monotonic, isolated, exportab ], }) }, 30_000) + +test('UserMeter package-service sourceUpdatedAt accepts only canonical UTC ISO timestamps', async () => { + const user = await seedFreeUser('meter-pkg-svc-iso') + const meter = userMeterRpc({ env, userId: user.userId }) + const stub = env.USER_METER.get( + env.USER_METER.idFromName(userMeterDurableObjectName(user.userId)), + ) + const valid = '2026-08-01T10:00:00.000Z' + + await expect( + meter.upsertPackageServiceState({ + packageId: 'pkg-iso', + serviceName: 'worker', + status: 'running', + startedAt: valid, + sourceUpdatedAt: valid, + }), + ).resolves.toMatchObject({ + applied: true, + state: { sourceUpdatedAt: valid }, + }) + + const rejected = [ + '', + 'not-a-timestamp', + '2026-08-01T10:00:00Z', + '2026-08-01 10:00:00.000Z', + '2026-08-01T10:00:00.000+02:00', + '2026-08-01T08:00:00.000+00:00', + '2026-13-01T00:00:00.000Z', + '2026-02-30T00:00:00.000Z', + ] as const + await runInDurableObject(stub, async (instance: UserMeter) => { + for (const sourceUpdatedAt of rejected) { + await expect( + instance.upsertPackageServiceState({ + packageId: 'pkg-iso', + serviceName: 'worker', + status: 'stopped', + startedAt: null, + sourceUpdatedAt, + }), + ).rejects.toThrow(/ISO-8601 UTC timestamp/) + } + }) + + // Prior valid row must remain; rejected writes must not throw RangeError past RPC. + expect(await meter.listPackageServiceStates({})).toMatchObject({ + states: [ + expect.objectContaining({ + packageId: 'pkg-iso', + status: 'running', + sourceUpdatedAt: valid, + }), + ], + }) +}, 30_000) diff --git a/packages/worker/src/package-runtime/package-service.node.test.ts b/packages/worker/src/package-runtime/package-service.node.test.ts index 66fc1f98a5..99093306bb 100644 --- a/packages/worker/src/package-runtime/package-service.node.test.ts +++ b/packages/worker/src/package-runtime/package-service.node.test.ts @@ -1004,6 +1004,187 @@ test('package service start/heartbeat/stop/purge shadow UserMeter transitions', } }) +test('package service serializes gated UserMeter shadows so purge stays deleted and running→stopped cannot regress', async () => { + resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + const { appDb } = createRecordingPackageServiceDb() + const meter = createInMemoryUserMeterEnv() + + type Gate = { + release: () => void + promise: Promise + label: string + } + const gates: Array = [] + const callOrder: Array = [] + const baseStub = meter.env.USER_METER.get( + meter.env.USER_METER.idFromName('user-123'), + ) + + function pushGate(label: string): Gate { + const gate: Gate = { + label, + release: () => {}, + promise: Promise.resolve(), + } + gate.promise = new Promise((resolve) => { + gate.release = resolve + }) + gates.push(gate) + return gate + } + + const gatedMeter = { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: () => ({ + async upsertPackageServiceState( + input: Parameters[0], + ) { + const gate = pushGate(`upsert:${input.status}`) + await gate.promise + callOrder.push(`upsert:${input.status}`) + return baseStub.upsertPackageServiceState(input) + }, + async deletePackageServiceState( + input: Parameters[0], + ) { + const gate = pushGate('delete') + await gate.promise + callOrder.push('delete') + return baseStub.deletePackageServiceState(input) + }, + }), + }, + } + + async function settleMicrotasks() { + await Promise.resolve() + await Promise.resolve() + await Promise.resolve() + } + + async function waitForGate(label: string) { + for (let i = 0; i < 100; i++) { + if (gates.some((gate) => gate.label === label)) return + await Promise.resolve() + } + throw new Error(`Timed out waiting for gated shadow op ${label}.`) + } + + try { + setupSavedPackage('bounded') + mockModule.runBundledModuleWithRegistry.mockImplementation(async () => { + vi.setSystemTime(new Date('2026-07-05T12:00:00.000Z')) + return { result: { ok: true }, error: null } + }) + + const created = await createPackageServiceInstance({ + APP_DB: appDb, + USER_METER: gatedMeter.USER_METER, + } as Env) + + const start = await postPackageService(created.instance, 'start') + expect(start.status).toBe(200) + // Lifecycle returned while the first shadow hop is still gated. + expect(created.waitUntilTasks.length).toBeGreaterThanOrEqual(2) + await waitForGate('upsert:running') + expect(gates.map((gate) => gate.label)).toEqual(['upsert:running']) + expect(callOrder).toEqual([]) + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toBeUndefined() + + // Finish the bounded background run while the running shadow stays gated so + // the same-ms stopped upsert is chained behind it (not started yet). + const backgroundRunTask = created.waitUntilTasks[1] + expect(backgroundRunTask).toBeDefined() + await backgroundRunTask + await settleMicrotasks() + expect(gates.map((gate) => gate.label)).toEqual(['upsert:running']) + expect(callOrder).toEqual([]) + + // Same-ms running→stopped: stopped cannot begin until running is released, + // so a reordered drain cannot let stopped win first and then regress. + gates[0]?.release() + await waitForGate('upsert:stopped') + expect(callOrder).toEqual(['upsert:running']) + expect(gates.map((gate) => gate.label)).toEqual([ + 'upsert:running', + 'upsert:stopped', + ]) + gates[1]?.release() + await drainSettledWaitUntilTasks(created.waitUntilTasks) + expect(callOrder).toEqual(['upsert:running', 'upsert:stopped']) + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toMatchObject({ + status: 'stopped', + sourceUpdatedAt: '2026-07-05T12:00:00.000Z', + }) + + callOrder.length = 0 + gates.length = 0 + const beforePurgeWaitUntilCount = created.waitUntilTasks.length + const purge = await postPackageService(created.instance, 'purge') + expect(purge.status).toBe(200) + await waitForGate('upsert:stopped') + // Purge schedules stopped upsert then delete; delete cannot start first. + expect(gates.map((gate) => gate.label)).toEqual(['upsert:stopped']) + expect(callOrder).toEqual([]) + + // Prefer settling the trailing waitUntil chain entry first while the head + // upsert is still gated. Serialization keeps delete behind the upsert, so + // the purge final state stays deleted instead of resurrecting. + const trailingPurgeChain = created.waitUntilTasks.at(-1) + expect(trailingPurgeChain).toBeDefined() + let trailingSettled = false + void Promise.resolve(trailingPurgeChain).then(() => { + trailingSettled = true + }) + await settleMicrotasks() + expect(trailingSettled).toBe(false) + expect(callOrder).toEqual([]) + + gates[0]?.release() + await waitForGate('delete') + expect(callOrder).toEqual(['upsert:stopped']) + expect(gates.map((gate) => gate.label)).toEqual([ + 'upsert:stopped', + 'delete', + ]) + expect(trailingSettled).toBe(false) + gates[1]?.release() + await drainSettledWaitUntilTasks( + created.waitUntilTasks, + beforePurgeWaitUntilCount, + ) + expect(trailingSettled).toBe(true) + expect(callOrder).toEqual(['upsert:stopped', 'delete']) + expect( + packageServiceShadowRow( + meter, + 'user-123', + 'package-1', + 'realtime-supervisor', + ), + ).toBeUndefined() + } finally { + vi.useRealTimers() + } +}) + test('package service lifecycle does not await a gated UserMeter shadow', async () => { resetMocks() vi.useFakeTimers() diff --git a/packages/worker/src/package-runtime/package-service.ts b/packages/worker/src/package-runtime/package-service.ts index 51a97ea2b0..af98468616 100644 --- a/packages/worker/src/package-runtime/package-service.ts +++ b/packages/worker/src/package-runtime/package-service.ts @@ -344,6 +344,8 @@ class PackageServiceInstanceBase extends DurableObject { private stateSnapshot: PackageServiceState = createInitialPackageServiceState() private activeRunPromise: Promise | null = null + /** Per-instance shadow write chain so waitUntil hops cannot reorder. */ + private shadowQueue: Promise = Promise.resolve() constructor(state: DurableObjectState, env: Env) { super(state, env) @@ -456,7 +458,8 @@ class PackageServiceInstanceBase extends DurableObject { /** * Schedule a best-effort UserMeter shadow upsert. Never awaited by * lifecycle/heartbeat paths; the helper catches so `waitUntil` cannot - * reject. + * reject. Appended to the per-instance shadow queue so rapid transitions + * and purge (stopped upsert then delete) cannot reorder. */ private schedulePackageServiceStateShadow(input: { binding: PackageServiceBindingState @@ -465,7 +468,9 @@ class PackageServiceInstanceBase extends DurableObject { sourceUpdatedAt: string }) { if (!userMeterNamespace(this.env)) return - this.ctx.waitUntil(this.shadowPackageServiceStateToUserMeter(input)) + this.enqueueShadowTask(() => + this.shadowPackageServiceStateToUserMeter(input), + ) } /** Best-effort UserMeter shadow upsert; catches/logs so waitUntil settles. */ @@ -494,13 +499,23 @@ class PackageServiceInstanceBase extends DurableObject { /** * Schedule a best-effort UserMeter shadow delete. Never awaited by purge; - * the helper catches so `waitUntil` cannot reject. + * the helper catches so `waitUntil` cannot reject. Queued after any prior + * shadow upserts for this instance. */ private schedulePackageServiceShadowDelete( binding: PackageServiceBindingState, ) { if (!userMeterNamespace(this.env)) return - this.ctx.waitUntil(this.deletePackageServiceShadow(binding)) + this.enqueueShadowTask(() => this.deletePackageServiceShadow(binding)) + } + + /** + * Run shadow writes in scheduling order and register the chain with + * `waitUntil`. Helpers already catch, so the queue itself does not reject. + */ + private enqueueShadowTask(task: () => Promise) { + this.shadowQueue = this.shadowQueue.then(task, task) + this.ctx.waitUntil(this.shadowQueue) } /** Best-effort UserMeter shadow delete; catches/logs so waitUntil settles. */ From 5e8d1b9fb3a7818db49bd72d85f32f0bd06d5253 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 09:25:09 +0000 Subject: [PATCH 20/39] feat(account): shadow deletion leases in UserMeter Co-authored-by: Kent C. Dodds --- .../src/account/deletion-state.node.test.ts | 433 +++++++++++++++++- packages/worker/src/account/deletion-state.ts | 96 ++++ .../src/app/account-deletion.node.test.ts | 15 +- packages/worker/src/app/account-deletion.ts | 1 + .../src/app/account-write-lease-middleware.ts | 1 + .../src/billing/stripe-plan-refresh-client.ts | 1 + .../src/billing/stripe-plan-refresh-do.ts | 2 + packages/worker/src/community/avatar.ts | 7 +- .../worker/src/community/community-icon.ts | 2 + .../worker/src/entitlements/user-meter-do.ts | 395 +++++++++++++++- .../entitlements/user-meter.workers.test.ts | 165 +++++++ .../worker/src/jobs/job-schedule-watchdog.ts | 2 + packages/worker/src/jobs/service.ts | 11 + packages/worker/src/mcp-auth.ts | 1 + .../admin/admin-account-write-lease-repair.ts | 1 + packages/worker/src/mcp/memory/service.ts | 4 + .../worker/src/package-registry/service.ts | 3 + .../worker/src/test-support/user-meter.ts | 165 +++++++ 18 files changed, 1295 insertions(+), 10 deletions(-) diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index a44e9fc8ce..0a9fc0f3f7 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -1,7 +1,17 @@ import { DatabaseSync } from 'node:sqlite' -import { expect, test } from 'vitest' +import { expect, test, vi } from 'vitest' import { createD1FromSqlite } from '#worker/test-support/create-d1-from-sqlite.ts' import { + userMeterRpc, + type UserMeterEnv, +} from '#worker/entitlements/user-meter-client.ts' +import { silenceExpectedConsoleWarns } from '#worker/test-support/console-spies.ts' +import { + createInMemoryUserMeterEnv, + createWaitUntilDrain, +} from '#worker/test-support/user-meter.ts' +import { + AccountDeletionInProgressError, AccountWriteLeaseLostError, listActiveAccountWriteLeases, markAccountDeleting, @@ -110,6 +120,21 @@ function createLeaseTestDb() { return { sqlite, db: createD1FromSqlite(sqlite) } } +function addLeaseRepairsTable(sqlite: DatabaseSync) { + sqlite.exec(` + CREATE TABLE account_write_lease_repairs ( + id TEXT PRIMARY KEY, + target_user_id TEXT NOT NULL, + lease_token TEXT NOT NULL, + lease_holder TEXT NOT NULL, + lease_acquired_at TEXT NOT NULL, + repaired_by_user_id TEXT NOT NULL, + reason TEXT NOT NULL, + created_at TEXT NOT NULL + ); + `) +} + function countLeaseRows(sqlite: DatabaseSync) { const row = sqlite .prepare(`SELECT COUNT(*) AS count FROM account_write_leases`) @@ -280,3 +305,409 @@ test('waitUntil moves lease release off the response path', async () => { finishRelease() await releasePromise }) + +test('optional env shadows mark/acquire/release/repair into UserMeter without changing D1 authority', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const meter = createInMemoryUserMeterEnv() + const drain = createWaitUntilDrain() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const meterB = userMeterRpc({ env: meter.env, userId: 'user-b' }) + + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:shadowed', + env: meter.env, + waitUntil: drain.waitUntil, + async write() { + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) + return 'ok' + }, + }) + await drain.drain() + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + + let heldToken = '' + let heldAcquiredAt = '' + let startWrite: () => void = () => undefined + let finishWrite: () => void = () => undefined + const started = new Promise((resolve) => { + startWrite = resolve + }) + const finish = new Promise((resolve) => { + finishWrite = resolve + }) + const operation = withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:repair-shadow', + env: meter.env, + waitUntil: drain.waitUntil, + async write() { + const [lease] = await listActiveAccountWriteLeases(db, 'user-a') + heldToken = lease!.token + heldAcquiredAt = lease!.acquired_at + startWrite() + await finish + return 'lost' + }, + }) + await started + await drain.drain() + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: meter.env, + waitUntil: drain.waitUntil, + }), + ).resolves.toBe(1) + await drain.drain() + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2099-01-01 00:00:00', + }) + + await repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: heldToken, + expectedAcquiredAt: heldAcquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + waitUntil: drain.waitUntil, + }) + await drain.drain() + finishWrite() + await expect(operation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + env: meter.env, + async write() { + return 'blocked' + }, + }), + ).rejects.toBeInstanceOf(AccountDeletionInProgressError) + await withAccountWriteLease({ + db, + stableUserId: 'user-b', + holder: 'test:other-user', + env: meter.env, + waitUntil: drain.waitUntil, + async write() { + return 'ok' + }, + }) + await drain.drain() + expect(await meterB.readDeletionState()).toEqual({ deletingAt: null }) +}) + +function createFailingShadowEnv(): UserMeterEnv { + return { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: () => ({ + async shadowMarkDeleting() { + throw new Error('shadow mark failed') + }, + async shadowAcquireWriteLease() { + throw new Error('shadow acquire failed') + }, + async shadowReleaseWriteLease() { + throw new Error('shadow release failed') + }, + }), + }, + } as unknown as UserMeterEnv +} + +test('UserMeter shadow failures never alter D1 lease or deletion results', async () => { + silenceExpectedConsoleWarns(['account-deletion-user-meter-shadow-failed']) + const { db } = createLeaseTestDb() + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const failingEnv = createFailingShadowEnv() + + let finishRelease: () => void = () => undefined + const releaseGate = new Promise((resolve) => { + finishRelease = resolve + }) + const originalBatch = db.batch.bind(db) + let releaseBatchSeen = false + db.batch = (async (statements: Parameters[0]) => { + const sql = statements + .map((statement) => { + const withQuery = statement as { query?: string } + return typeof withQuery.query === 'string' ? withQuery.query : '' + }) + .join('\n') + const isRelease = + sql.includes('MAX(active_write_count - 1') || + sql.includes('DELETE FROM account_write_leases') + if (isRelease && !releaseBatchSeen) { + releaseBatchSeen = true + await releaseGate + } + return await originalBatch(statements) + }) as D1Database['batch'] + + const pending: Array> = [] + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:shadow-fail', + env: failingEnv, + waitUntil(promise) { + pending.push(promise) + }, + async write() { + return 'mutated' + }, + }), + ).resolves.toBe('mutated') + expect(pending.length).toBeGreaterThan(0) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) + finishRelease() + await Promise.all(pending) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: failingEnv, + }), + ).resolves.toBe(0) + expect(warn).toHaveBeenCalledWith( + 'account-deletion-user-meter-shadow-failed', + expect.any(Error), + ) + warn.mockRestore() +}) + +function createDeferred() { + let resolve: () => void = () => undefined + const promise = new Promise((next) => { + resolve = next + }) + return { promise, resolve: () => resolve() } +} + +function createGatedShadowEnv(input: { + releaseGate: Promise + markGate?: Promise +}) { + const base = createInMemoryUserMeterEnv() + const namespace = base.env.USER_METER! + const gates = { + markEntered: false, + releaseEntered: false, + } + const env = { + USER_METER: { + idFromName: namespace.idFromName.bind(namespace), + get(id: DurableObjectId) { + const stub = namespace.get(id) + return { + shadowMarkDeleting: async (args: { deletingAt: string }) => { + gates.markEntered = true + if (input.markGate) await input.markGate + return stub.shadowMarkDeleting(args) + }, + shadowAcquireWriteLease: (args: { + token: string + holder: string + acquiredAt: string + }) => stub.shadowAcquireWriteLease(args), + shadowReleaseWriteLease: async (args: { token: string }) => { + gates.releaseEntered = true + await input.releaseGate + return stub.shadowReleaseWriteLease(args) + }, + countActiveWriteLeases: () => stub.countActiveWriteLeases(), + readDeletionState: () => stub.readDeletionState(), + } + }, + }, + } as unknown as UserMeterEnv + return { + env, + gates, + meterFor(userId: string) { + return userMeterRpc({ env, userId }) + }, + } +} + +async function waitFor(predicate: () => Promise, label: string) { + const deadline = Date.now() + 2_000 + while (Date.now() < deadline) { + if (await predicate()) return + await new Promise((resolve) => setTimeout(resolve, 5)) + } + throw new Error(`Timed out waiting for ${label}`) +} + +test('without waitUntil, lease return awaits release shadow', async () => { + const { db } = createLeaseTestDb() + const release = createDeferred() + const gated = createGatedShadowEnv({ releaseGate: release.promise }) + const meterA = gated.meterFor('user-a') + + let settled = false + const operation = withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:await-release-shadow', + env: gated.env, + async write() { + return 'ok' + }, + }).then((value) => { + settled = true + return value + }) + await waitFor( + async () => + gated.gates.releaseEntered && + (await listActiveAccountWriteLeases(db, 'user-a')).length === 0 && + (await meterA.countActiveWriteLeases()).count === 1, + 'D1 release with gated shadow lease still held', + ) + expect(settled).toBe(false) + release.resolve() + await expect(operation).resolves.toBe('ok') + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) +}) + +test('without waitUntil, mark and repair await their shadows', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const release = createDeferred() + const mark = createDeferred() + const gated = createGatedShadowEnv({ + releaseGate: release.promise, + markGate: mark.promise, + }) + const meterA = gated.meterFor('user-a') + const drain = createWaitUntilDrain() + + let heldToken = '' + let heldAcquiredAt = '' + let startWrite: () => void = () => undefined + let finishWrite: () => void = () => undefined + const started = new Promise((resolve) => { + startWrite = resolve + }) + const finish = new Promise((resolve) => { + finishWrite = resolve + }) + const held = withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:await-mark-repair', + env: gated.env, + waitUntil: drain.waitUntil, + async write() { + const [lease] = await listActiveAccountWriteLeases(db, 'user-a') + heldToken = lease!.token + heldAcquiredAt = lease!.acquired_at + startWrite() + await finish + return 'lost' + }, + }) + await started + await drain.drain() + + let markSettled = false + const markPromise = markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: gated.env, + }).then((count) => { + markSettled = true + return count + }) + await waitFor( + async () => gated.gates.markEntered && !markSettled, + 'mark shadow gate', + ) + expect(await meterA.readDeletionState()).toEqual({ deletingAt: null }) + mark.resolve() + await expect(markPromise).resolves.toBe(1) + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2099-01-01 00:00:00', + }) + + let repairSettled = false + const repairPromise = repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: heldToken, + expectedAcquiredAt: heldAcquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: gated.env, + }).then((result) => { + repairSettled = true + return result + }) + await waitFor( + async () => + gated.gates.releaseEntered && + (await meterA.countActiveWriteLeases()).count === 1 && + !repairSettled, + 'repair shadow gate', + ) + release.resolve() + await expect(repairPromise).resolves.toEqual( + expect.objectContaining({ repaired: true }), + ) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + finishWrite() + await expect(held).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + await drain.drain() +}) + +test('with waitUntil, lease return completes before release shadow settles', async () => { + const { db } = createLeaseTestDb() + const release = createDeferred() + const gated = createGatedShadowEnv({ releaseGate: release.promise }) + const meterA = gated.meterFor('user-a') + const drain = createWaitUntilDrain() + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:detach-release-shadow', + env: gated.env, + waitUntil: drain.waitUntil, + async write() { + return 'ok' + }, + }), + ).resolves.toBe('ok') + await waitFor( + async () => + gated.gates.releaseEntered && + (await listActiveAccountWriteLeases(db, 'user-a')).length === 0 && + (await meterA.countActiveWriteLeases()).count === 1, + 'detached release shadow still held', + ) + + release.resolve() + await drain.drain() + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) +}) diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index 41286fe6d7..495da004e0 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -1,5 +1,10 @@ import { AsyncLocalStorage } from 'node:async_hooks' import { utcSqliteTimestamp } from '@kody-internal/shared/date-keys.ts' +import { + userMeterNamespace, + userMeterRpc, + type UserMeterEnv, +} from '#worker/entitlements/user-meter-client.ts' export class AccountDeletionInProgressError extends Error { constructor() { @@ -34,10 +39,31 @@ export type AccountWriteLease = { acquiredAt: string } +const accountDeletionShadowFailedLog = + 'account-deletion-user-meter-shadow-failed' + +function scheduleAccountDeletionShadow(input: { + env?: UserMeterEnv + waitUntil?: (promise: Promise) => void + task: (env: UserMeterEnv) => Promise +}): Promise { + const env = input.env + if (!env || !userMeterNamespace(env)) return Promise.resolve() + const tracked = input.task(env).catch((error: unknown) => { + console.warn(accountDeletionShadowFailedLog, error) + }) + if (input.waitUntil) input.waitUntil(tracked) + return tracked +} + export async function markAccountDeleting(input: { db: D1Database dbUserId: number now?: Date + /** Optional expand-phase UserMeter shadow after a successful D1 mark. */ + env?: UserMeterEnv + /** Prefer `ctx.waitUntil`; awaited (non-rejecting) when omitted. */ + waitUntil?: (promise: Promise) => void }) { const now = utcSqliteTimestamp(input.now ?? new Date()) const result = await input.db @@ -61,6 +87,29 @@ export async function markAccountDeleting(input: { ) .bind(input.dbUserId) .first<{ count: number }>() + if (input.env && userMeterNamespace(input.env)) { + const userRow = await input.db + .prepare( + `SELECT stable_user_id, deleting_at + FROM users + WHERE id = ?`, + ) + .bind(input.dbUserId) + .first<{ stable_user_id: string; deleting_at: string | null }>() + if (userRow?.stable_user_id && userRow.deleting_at) { + const { stable_user_id: stableUserId, deleting_at: deletingAt } = userRow + const markShadowPromise = scheduleAccountDeletionShadow({ + env: input.env, + waitUntil: input.waitUntil, + task: async (env) => { + await userMeterRpc({ env, userId: stableUserId }).shadowMarkDeleting({ + deletingAt, + }) + }, + }) + if (!input.waitUntil) await markShadowPromise + } + } return Number(row?.count ?? 0) } @@ -107,6 +156,8 @@ export async function withAccountWriteLease(input: { stableUserId: string holder?: string waitUntil?: (promise: Promise) => void + /** Optional expand-phase UserMeter shadow after D1 acquire/release. */ + env?: UserMeterEnv write: () => Promise }) { const heldLeases = heldAccountWriteLeaseStorage.getStore() @@ -130,6 +181,7 @@ async function acquireAccountWriteLeaseAndWrite(input: { stableUserId: string holder?: string waitUntil?: (promise: Promise) => void + env?: UserMeterEnv frame: AccountWriteLeaseFrame write: () => Promise }) { @@ -228,6 +280,21 @@ async function acquireAccountWriteLeaseAndWrite(input: { ) { throw new AccountDeletionInProgressError() } + // Hold acquire so a detached release shadow cannot land first. + const acquireShadowPromise = scheduleAccountDeletionShadow({ + env: input.env, + waitUntil: input.waitUntil, + task: async (env) => { + await userMeterRpc({ + env, + userId: lease.stableUserId, + }).shadowAcquireWriteLease({ + token: lease.token, + holder: lease.holder, + acquiredAt: lease.acquiredAt, + }) + }, + }) try { const result = await input.write() const held = await input.db @@ -292,6 +359,20 @@ async function acquireAccountWriteLeaseAndWrite(input: { const releasePromise = release() if (input.waitUntil) input.waitUntil(releasePromise) else await releasePromise + // Ordered after D1 release; await when waitUntil is omitted. + const releaseShadowPromise = scheduleAccountDeletionShadow({ + env: input.env, + waitUntil: input.waitUntil, + task: async (env) => { + await releasePromise + await acquireShadowPromise + await userMeterRpc({ + env, + userId: lease.stableUserId, + }).shadowReleaseWriteLease({ token: lease.token }) + }, + }) + if (!input.waitUntil) await releaseShadowPromise } } @@ -318,6 +399,10 @@ export async function repairAccountWriteLease(input: { expectedAcquiredAt: string repairedByUserId: string reason: string + /** Optional expand-phase UserMeter shadow after a successful D1 repair. */ + env?: UserMeterEnv + /** Prefer `ctx.waitUntil`; awaited (non-rejecting) when omitted. */ + waitUntil?: (promise: Promise) => void }) { if (input.reason.trim().length < 10) { throw new Error('Lease repair requires a detailed audit reason.') @@ -377,5 +462,16 @@ export async function repairAccountWriteLease(input: { ) { throw new Error('Active account write lease did not match repair request.') } + const repairShadowPromise = scheduleAccountDeletionShadow({ + env: input.env, + waitUntil: input.waitUntil, + task: async (env) => { + await userMeterRpc({ + env, + userId: input.stableUserId, + }).shadowReleaseWriteLease({ token: input.token }) + }, + }) + if (!input.waitUntil) await repairShadowPromise return { repaired: true as const, repairId } } diff --git a/packages/worker/src/app/account-deletion.node.test.ts b/packages/worker/src/app/account-deletion.node.test.ts index a50251963c..3ee082b6b0 100644 --- a/packages/worker/src/app/account-deletion.node.test.ts +++ b/packages/worker/src/app/account-deletion.node.test.ts @@ -20,6 +20,7 @@ import { consoleError, consoleWarn, } from '#worker/test-support/console-spies.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' type RowMap = Record>> @@ -513,6 +514,7 @@ function createSuccessfulDeletionEnv( ) { const durableObjectId = (name: string) => name as unknown as DurableObjectId const fetchOk = async () => Response.json({ ok: true }) + const userMeter = createInMemoryUserMeterEnv() return { APP_DB: db, CAPABILITY_VECTOR_INDEX: { @@ -529,10 +531,7 @@ function createSuccessfulDeletionEnv( listStorageIds: async () => [] as Array, }), }, - USER_METER: { - idFromName: durableObjectId, - get: () => ({ purge: async () => ({ ok: true as const }) }), - }, + USER_METER: userMeter.env.USER_METER, STRIPE_PLAN_REFRESH: { idFromName: durableObjectId, get: () => ({ purgeUser: async () => ({ ok: true as const }) }), @@ -1131,6 +1130,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async const purgeMcpAgentSessionMock = vi.fn(async () => undefined) const doFetchMock = vi.fn(async () => Response.json({ ok: true })) const deleteVectorsMock = vi.fn(async () => undefined) + const userMeter = createInMemoryUserMeterEnv() const env = createSuccessfulDeletionEnv(db, { BUNDLE_ARTIFACTS_KV: kv, COMMUNITY_ASSETS: communityAssets, @@ -1150,8 +1150,11 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async }), }, USER_METER: { - idFromName: (name: string) => name as unknown as DurableObjectId, - get: () => ({ purge: purgeUserMeterMock }), + idFromName: (name: string) => userMeter.env.USER_METER!.idFromName(name), + get: (id: DurableObjectId) => ({ + ...userMeter.env.USER_METER!.get(id), + purge: async () => purgeUserMeterMock(), + }), }, STRIPE_PLAN_REFRESH: { idFromName: stripePlanRefreshIdFromNameMock, diff --git a/packages/worker/src/app/account-deletion.ts b/packages/worker/src/app/account-deletion.ts index 4cd8b15b98..81a2ce0472 100644 --- a/packages/worker/src/app/account-deletion.ts +++ b/packages/worker/src/app/account-deletion.ts @@ -1064,6 +1064,7 @@ export async function deleteUserAccount(input: { const activeWriteCount = await markAccountDeleting({ db: input.env.APP_DB, dbUserId: input.dbUserId, + env: input.env, }) if (activeWriteCount > 0) { throw new AccountDeletionWritersActiveError(activeWriteCount) diff --git a/packages/worker/src/app/account-write-lease-middleware.ts b/packages/worker/src/app/account-write-lease-middleware.ts index ab1831db17..24cdcc01c8 100644 --- a/packages/worker/src/app/account-write-lease-middleware.ts +++ b/packages/worker/src/app/account-write-lease-middleware.ts @@ -39,6 +39,7 @@ export function createAccountWriteLeaseMiddleware(env: Env): Middleware { db: env.APP_DB, stableUserId: auth.user.mcpUser.userId, holder: `web:${request.method} ${url.pathname}`, + env, write: next, }) } catch (error) { diff --git a/packages/worker/src/billing/stripe-plan-refresh-client.ts b/packages/worker/src/billing/stripe-plan-refresh-client.ts index 3a4fd66dde..514555b9a8 100644 --- a/packages/worker/src/billing/stripe-plan-refresh-client.ts +++ b/packages/worker/src/billing/stripe-plan-refresh-client.ts @@ -25,6 +25,7 @@ export async function scheduleStripePlanRefreshBackstop(input: { db: input.env.APP_DB, stableUserId: userId, holder: 'stripe_plan_refresh_schedule', + env: input.env, write: async () => { await stub.schedule({ userId, diff --git a/packages/worker/src/billing/stripe-plan-refresh-do.ts b/packages/worker/src/billing/stripe-plan-refresh-do.ts index 5c8828dbdc..12439126c4 100644 --- a/packages/worker/src/billing/stripe-plan-refresh-do.ts +++ b/packages/worker/src/billing/stripe-plan-refresh-do.ts @@ -60,6 +60,8 @@ class StripePlanRefreshBase extends DurableObject { db: this.env.APP_DB, stableUserId: userId, holder: 'stripe_plan_refresh_alarm', + env: this.env, + waitUntil: this.ctx.waitUntil.bind(this.ctx), write: async () => { await refreshStripePlanForUser({ env: this.env, diff --git a/packages/worker/src/community/avatar.ts b/packages/worker/src/community/avatar.ts index 7c35fd08b3..2b853f98ab 100644 --- a/packages/worker/src/community/avatar.ts +++ b/packages/worker/src/community/avatar.ts @@ -4,6 +4,7 @@ import { assertAccountWritableDb, withAccountWriteLease, } from '#worker/account/deletion-state.ts' +import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' import { readJpegDimensions, readPngDimensions, @@ -146,7 +147,7 @@ async function sha256Hex(bytes: Uint8Array) { } export async function saveUserAvatar(input: { - env: Pick + env: Pick & UserMeterEnv numericUserId: number stableUserId: string bytes: Uint8Array @@ -155,6 +156,7 @@ export async function saveUserAvatar(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.stableUserId, + env: input.env, async write() { await assertAccountWritableDb(input.env.APP_DB, input.stableUserId) const contentHash = await sha256Hex(input.bytes) @@ -226,13 +228,14 @@ export async function saveUserAvatar(input: { } export async function deleteUserAvatar(input: { - env: Pick + env: Pick & UserMeterEnv numericUserId: number stableUserId: string }): Promise { await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.stableUserId, + env: input.env, async write() { await assertAccountWritableDb(input.env.APP_DB, input.stableUserId) const existing = await input.env.APP_DB.prepare( diff --git a/packages/worker/src/community/community-icon.ts b/packages/worker/src/community/community-icon.ts index a9b2ef84e0..f77b4c0c67 100644 --- a/packages/worker/src/community/community-icon.ts +++ b/packages/worker/src/community/community-icon.ts @@ -100,6 +100,7 @@ export async function getCommunityIconObject(input: { await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.listing.ownerUserId, + env: input.env, write, }) } else { @@ -361,6 +362,7 @@ async function createCommunityIconDescriptor(input: { ? await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.listing.ownerUserId, + env: input.env, write, }) : await write() diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 12a7f82461..c4eaf94849 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -26,9 +26,11 @@ export const userMeterDailyCounterRetentionDays = 7 const metaSchemaVersionKey = 'schema_version' /** Bump when initializeSchema DDL changes; warm objects skip DDL. */ -const userMeterSchemaVersion = 5 +const userMeterSchemaVersion = 6 /** Singleton row id for expand-phase D1 storage-byte shadow (schema v4). */ const storageBytesStateRowId = 1 +/** Singleton row id for expand-phase D1 deletion-fence shadow (schema v6). */ +const deletionStateRowId = 1 /** Matches D1 `packageServiceStateStaleMs` (24h); cutover-support RPC only. */ export const userMeterPackageServiceStateStaleMs = 24 * 60 * 60 * 1000 @@ -36,6 +38,9 @@ const defaultExportPageSize = 100 const maxExportPageSize = 500 const maxInboundDeliveryIdLength = 256 const maxPackageServiceIdLength = 256 +const maxWriteLeaseTokenLength = 64 +const maxWriteLeaseHolderLength = 256 +const maxDeletionTimestampLength = 64 const inboundReceiveResource = 'email_receives_per_day' satisfies DailyEntitlementResource const utcDayKeyPattern = /^\d{4}-\d{2}-\d{2}$/ @@ -169,6 +174,46 @@ export type UserMeterPackageServiceBootstrapResult = { skipped: number } +export type UserMeterWriteLeaseShadow = { + token: string + holder: string + acquiredAt: string +} + +export type UserMeterDeletionShadow = { + deletingAt: string | null + writeLeases: Array +} + +export type UserMeterShadowMarkDeletingResult = { + deletingAt: string + created: boolean +} + +export type UserMeterShadowAcquireWriteLeaseResult = { + acquired: boolean +} + +export type UserMeterShadowReleaseWriteLeaseResult = { + released: boolean +} + +export type UserMeterWriteLeaseListResult = { + leases: Array + nextStartAfter: string | null + truncated: boolean +} + +export type UserMeterWriteLeaseCountResult = { + count: number +} + +export type UserMeterDeletionBootstrapResult = { + deletingAtApplied: boolean + leasesApplied: number + leasesSkipped: number +} + export type UserMeterExportResult = { counters: Array /** @@ -183,6 +228,12 @@ export type UserMeterExportResult = { * `null` so paged consumers never double-count the inventory. */ packageServiceStatesShadow: Array | null + /** + * Non-authoritative D1 deletion-fence / write-lease shadow. Emitted only + * on the first export page (`startAfter` absent); subsequent pages return + * `null` so paged consumers never double-count the inventory. + */ + deletionShadow: UserMeterDeletionShadow | null nextStartAfter: string | null truncated: boolean } @@ -211,6 +262,11 @@ type PackageServiceExportCursor = { serviceName: string } +type WriteLeaseExportCursor = { + acquiredAt: string + token: string +} + type PackageServiceSqlRow = { package_id: string service_name: string @@ -221,6 +277,12 @@ type PackageServiceSqlRow = { updated_at: string } +type WriteLeaseSqlRow = { + token: string + holder: string + acquired_at: string +} + function assertDailyResource(resource: string): DailyEntitlementResource { if (!isDailyEntitlementResource(resource)) { throw new Error( @@ -310,6 +372,45 @@ function assertSourceUpdatedAt(sourceUpdatedAt: string): string { return sourceUpdatedAt } +function assertDeletionTimestamp(label: string, value: string): string { + if ( + typeof value !== 'string' || + value.length === 0 || + value.length > maxDeletionTimestampLength + ) { + throw new Error( + `UserMeter ${label} must be a non-empty string up to ${maxDeletionTimestampLength} characters.`, + ) + } + return value +} + +function assertWriteLeaseToken(token: string): string { + if ( + typeof token !== 'string' || + token.length === 0 || + token.length > maxWriteLeaseTokenLength + ) { + throw new Error( + `UserMeter write lease token must be a non-empty string up to ${maxWriteLeaseTokenLength} characters.`, + ) + } + return token +} + +function assertWriteLeaseHolder(holder: string): string { + if ( + typeof holder !== 'string' || + holder.length === 0 || + holder.length > maxWriteLeaseHolderLength + ) { + throw new Error( + `UserMeter write lease holder must be a non-empty string up to ${maxWriteLeaseHolderLength} characters.`, + ) + } + return holder +} + function retentionCutoffDay(now: Date): string { const cutoff = new Date(now) cutoff.setUTCDate( @@ -354,6 +455,26 @@ function decodePackageServiceCursor( } } +function encodeWriteLeaseCursor(cursor: WriteLeaseExportCursor) { + return JSON.stringify([cursor.acquiredAt, cursor.token]) +} + +function decodeWriteLeaseCursor( + startAfter: string, +): WriteLeaseExportCursor | null { + try { + const parsed = JSON.parse(startAfter) as unknown + if (!Array.isArray(parsed) || parsed.length !== 2) return null + const [acquiredAt, token] = parsed + if (typeof acquiredAt !== 'string' || typeof token !== 'string') { + return null + } + return { acquiredAt, token } + } catch { + return null + } +} + function normalizePageSize(pageSize: number | undefined) { const requested = typeof pageSize === 'number' && Number.isFinite(pageSize) @@ -464,6 +585,24 @@ class UserMeterBase extends DurableObject { `CREATE INDEX IF NOT EXISTS idx_package_service_states_status_source ON package_service_states (status, source_updated_at)`, ) + // Expand-phase shadow of D1 deletion fence / write leases (schema v6). + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS deletion_state ( + id INTEGER PRIMARY KEY NOT NULL CHECK (id = 1), + deleting_at TEXT NOT NULL + ) + `) + this.ctx.storage.sql.exec(` + CREATE TABLE IF NOT EXISTS account_write_leases ( + token TEXT PRIMARY KEY NOT NULL, + holder TEXT NOT NULL, + acquired_at TEXT NOT NULL + ) + `) + this.ctx.storage.sql.exec( + `CREATE INDEX IF NOT EXISTS idx_account_write_leases_acquired_token + ON account_write_leases (acquired_at, token)`, + ) this.ctx.storage.sql.exec( `INSERT INTO user_meter_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, @@ -1205,10 +1344,230 @@ class UserMeterBase extends DurableObject { return { applied, skipped } } + private readDeletingAt(): string | null { + const row = this.ctx.storage.sql + .exec<{ + deleting_at: string + }>( + `SELECT deleting_at FROM deletion_state WHERE id = ?`, + deletionStateRowId, + ) + .toArray()[0] + if (!row) return null + const deletingAt = String(row.deleting_at ?? '') + return deletingAt.length > 0 ? deletingAt : null + } + + private writeLeaseFromRow(row: WriteLeaseSqlRow): UserMeterWriteLeaseShadow { + return { + token: String(row.token), + holder: String(row.holder), + acquiredAt: String(row.acquired_at), + } + } + + private readWriteLease(token: string): UserMeterWriteLeaseShadow | null { + const row = this.ctx.storage.sql + .exec( + `SELECT token, holder, acquired_at + FROM account_write_leases + WHERE token = ?`, + token, + ) + .toArray()[0] + if (!row) return null + return this.writeLeaseFromRow(row) + } + + private listAllWriteLeaseRows(): Array { + const rows = this.ctx.storage.sql + .exec( + `SELECT token, holder, acquired_at + FROM account_write_leases + ORDER BY acquired_at ASC, token ASC`, + ) + .toArray() + return rows.map((row) => this.writeLeaseFromRow(row)) + } + + private readDeletionShadow(): UserMeterDeletionShadow { + return { + deletingAt: this.readDeletingAt(), + writeLeases: this.listAllWriteLeaseRows(), + } + } + + /** Expand-phase shadow of D1 `users.deleting_at` (first write wins). */ + async shadowMarkDeleting(input: { + deletingAt: string + }): Promise { + const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + const existing = this.readDeletingAt() + if (existing != null) { + return { deletingAt: existing, created: false } + } + this.ctx.storage.sql.exec( + `INSERT INTO deletion_state (id, deleting_at) + VALUES (?, ?) + ON CONFLICT(id) DO NOTHING`, + deletionStateRowId, + deletingAt, + ) + const stored = this.readDeletingAt() ?? deletingAt + return { deletingAt: stored, created: stored === deletingAt } + } + + /** Expand-phase shadow lease acquire; fencing still uses D1. */ + async shadowAcquireWriteLease(input: { + token: string + holder: string + acquiredAt: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const holder = assertWriteLeaseHolder(input.holder) + const acquiredAt = assertDeletionTimestamp('acquiredAt', input.acquiredAt) + if (this.readWriteLease(token)) return { acquired: true } + if (this.readDeletingAt() != null) return { acquired: false } + const cursor = this.ctx.storage.sql.exec( + `INSERT INTO account_write_leases (token, holder, acquired_at) + VALUES (?, ?, ?) + ON CONFLICT(token) DO NOTHING`, + token, + holder, + acquiredAt, + ) + return { + acquired: cursor.rowsWritten > 0 || this.readWriteLease(token) != null, + } + } + + /** Expand-phase shadow lease release / repair mirror; fencing uses D1. */ + async shadowReleaseWriteLease(input: { + token: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const cursor = this.ctx.storage.sql.exec( + `DELETE FROM account_write_leases WHERE token = ?`, + token, + ) + return { released: cursor.rowsWritten > 0 } + } + + /** Cutover-support deletion tombstone read; Phase A uses D1. */ + async readDeletionState(): Promise<{ deletingAt: string | null }> { + return { deletingAt: this.readDeletingAt() } + } + + /** Cutover-support paged shadow lease list; Phase A list uses D1. */ + async listWriteLeases( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ): Promise { + const pageSize = normalizePageSize(input.pageSize) + const cursor = + typeof input.startAfter === 'string' && input.startAfter.length > 0 + ? decodeWriteLeaseCursor(input.startAfter) + : null + const rows = ( + cursor + ? this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at + FROM account_write_leases + WHERE acquired_at > ? + OR (acquired_at = ? AND token > ?) + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + cursor.acquiredAt, + cursor.acquiredAt, + cursor.token, + pageSize + 1, + ) + : this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at + FROM account_write_leases + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + pageSize + 1, + ) + ).toArray() + const truncated = rows.length > pageSize + const pageRows = truncated ? rows.slice(0, pageSize) : rows + const leases = pageRows.map((row) => this.writeLeaseFromRow(row)) + const last = pageRows[pageRows.length - 1] + return { + leases, + nextStartAfter: + truncated && last + ? encodeWriteLeaseCursor({ + acquiredAt: String(last.acquired_at), + token: String(last.token), + }) + : null, + truncated, + } + } + + /** Cutover-support active lease count; Phase A drain uses D1. */ + async countActiveWriteLeases(): Promise { + const row = this.ctx.storage.sql + .exec<{ + count: number + }>(`SELECT COUNT(*) AS count FROM account_write_leases`) + .toArray()[0] + return { count: Math.max(0, Number(row?.count ?? 0)) } + } + + /** + * Cutover-support cold seed. Leases apply first (same guards as + * {@link shadowAcquireWriteLease}), then the deleting tombstone. + */ + async bootstrapDeletionState(input: { + deletingAt?: string | null + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }): Promise { + let leasesApplied = 0 + let leasesSkipped = 0 + for (const lease of input.leases ?? []) { + const token = assertWriteLeaseToken(lease.token) + if (this.readWriteLease(token)) { + leasesSkipped += 1 + continue + } + const result = await this.shadowAcquireWriteLease(lease) + if (result.acquired) leasesApplied += 1 + else leasesSkipped += 1 + } + let deletingAtApplied = false + if (typeof input.deletingAt === 'string' && input.deletingAt.length > 0) { + const marked = await this.shadowMarkDeleting({ + deletingAt: input.deletingAt, + }) + deletingAtApplied = marked.created + } + return { deletingAtApplied, leasesApplied, leasesSkipped } + } + async purge(): Promise<{ ok: true }> { await this.ctx.blockConcurrencyWhile(async () => { + // Preserve deleting tombstone across deleteAll for later cutover safety. + const deletingAt = this.readDeletingAt() await this.ctx.storage.deleteAll() this.initializeSchema() + if (deletingAt != null) { + this.ctx.storage.sql.exec( + `INSERT INTO deletion_state (id, deleting_at) + VALUES (?, ?) + ON CONFLICT(id) DO NOTHING`, + deletionStateRowId, + deletingAt, + ) + } }) return { ok: true } } @@ -1283,6 +1642,7 @@ class UserMeterBase extends DurableObject { const packageServiceStatesShadow = includeShadows ? this.listAllPackageServiceRows() : null + const deletionShadow = includeShadows ? this.readDeletionShadow() : null const last = pageRows[pageRows.length - 1] return { counters, @@ -1295,6 +1655,7 @@ class UserMeterBase extends DurableObject { } : null, packageServiceStatesShadow, + deletionShadow, nextStartAfter: truncated && last ? encodeExportCursor({ @@ -1396,6 +1757,38 @@ export type UserMeterRpc = { updatedAt?: string }> }) => Promise + /** Expand-phase shadow of D1 users.deleting_at; fence still uses D1. */ + shadowMarkDeleting: (input: { + deletingAt: string + }) => Promise + /** Expand-phase shadow lease acquire; fencing still uses D1. */ + shadowAcquireWriteLease: (input: { + token: string + holder: string + acquiredAt: string + }) => Promise + /** Expand-phase shadow lease release / repair mirror; fencing uses D1. */ + shadowReleaseWriteLease: (input: { + token: string + }) => Promise + /** Cutover-support deletion tombstone read; Phase A uses D1. */ + readDeletionState: () => Promise<{ deletingAt: string | null }> + /** Cutover-support paged shadow lease list; Phase A list uses D1. */ + listWriteLeases: (input: { + pageSize?: number + startAfter?: string | null + }) => Promise + /** Cutover-support active lease count; Phase A drain uses D1. */ + countActiveWriteLeases: () => Promise + /** Cutover-support cold seed with monotonic mark / acquire guards. */ + bootstrapDeletionState: (input: { + deletingAt?: string | null + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) => Promise purge: () => Promise<{ ok: true }> exportCounters: (input: { pageSize?: number diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index a330f75f52..6f3e09b7e4 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -562,6 +562,7 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], + deletionShadow: { deletingAt: null, writeLeases: [] }, nextStartAfter: null, truncated: false, }) @@ -692,6 +693,7 @@ test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], + deletionShadow: { deletingAt: null, writeLeases: [] }, nextStartAfter: null, truncated: false, }) @@ -966,6 +968,10 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a mirrorUpdatedAt: userMeterMirrorUpdatedAtToken(3), }) expect(firstPage.packageServiceStatesShadow).toEqual([]) + expect(firstPage.deletionShadow).toEqual({ + deletingAt: null, + writeLeases: [], + }) const secondPage = await meter.exportCounters({ pageSize: 2, @@ -976,6 +982,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a expect(secondPage.nextStartAfter).toBeNull() expect(secondPage.storageBytesShadow).toBeNull() expect(secondPage.packageServiceStatesShadow).toBeNull() + expect(secondPage.deletionShadow).toBeNull() await expect( readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }), @@ -989,6 +996,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], + deletionShadow: { deletingAt: null, writeLeases: [] }, nextStartAfter: null, truncated: false, }) @@ -1186,6 +1194,7 @@ test('UserMeter package-service shadow upserts are monotonic, isolated, exportab startAfter: firstPage.nextStartAfter, }) expect(secondPage.packageServiceStatesShadow).toBeNull() + expect(secondPage.deletionShadow).toBeNull() await expect( meterA.deletePackageServiceState({ @@ -1272,3 +1281,159 @@ test('UserMeter package-service sourceUpdatedAt accepts only canonical UTC ISO t ], }) }, 30_000) + +test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserves tombstone across purge', async () => { + const userA = await seedFreeUser('meter-deletion-a') + const userB = await seedFreeUser('meter-deletion-b') + const meterA = userMeterRpc({ env, userId: userA.userId }) + const meterB = userMeterRpc({ env, userId: userB.userId }) + + const firstMark = await meterA.shadowMarkDeleting({ + deletingAt: '2026-08-01 10:00:00', + }) + expect(firstMark).toEqual({ + deletingAt: '2026-08-01 10:00:00', + created: true, + }) + const secondMark = await meterA.shadowMarkDeleting({ + deletingAt: '2026-08-01 11:00:00', + }) + expect(secondMark).toEqual({ + deletingAt: '2026-08-01 10:00:00', + created: false, + }) + + await expect( + meterB.shadowAcquireWriteLease({ + token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + holder: 'test:writer', + acquiredAt: '2026-08-01 10:05:00', + }), + ).resolves.toEqual({ acquired: true }) + await expect( + meterB.shadowAcquireWriteLease({ + token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + holder: 'test:writer', + acquiredAt: '2026-08-01 10:05:00', + }), + ).resolves.toEqual({ acquired: true }) + await expect( + meterB.shadowAcquireWriteLease({ + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:other', + acquiredAt: '2026-08-01 10:06:00', + }), + ).resolves.toEqual({ acquired: true }) + expect(await meterB.countActiveWriteLeases()).toEqual({ count: 2 }) + + const listed = await meterB.listWriteLeases({ pageSize: 1 }) + expect(listed.leases).toHaveLength(1) + expect(listed.truncated).toBe(true) + const listedRest = await meterB.listWriteLeases({ + pageSize: 10, + startAfter: listed.nextStartAfter, + }) + expect(listedRest.leases).toHaveLength(1) + expect(listedRest.truncated).toBe(false) + + await expect( + meterB.shadowReleaseWriteLease({ + token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + }), + ).resolves.toEqual({ released: true }) + expect(await meterB.countActiveWriteLeases()).toEqual({ count: 1 }) + + await expect( + meterA.shadowAcquireWriteLease({ + token: 'cccccccc-cccc-4ccc-8ccc-cccccccccccc', + holder: 'test:blocked', + acquiredAt: '2026-08-01 10:07:00', + }), + ).resolves.toEqual({ acquired: false }) + + const bootstrap = await meterB.bootstrapDeletionState({ + deletingAt: '2026-08-01 12:00:00', + leases: [ + { + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:other', + acquiredAt: '2026-08-01 10:06:00', + }, + { + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + holder: 'test:boot', + acquiredAt: '2026-08-01 10:08:00', + }, + ], + }) + expect(bootstrap).toEqual({ + deletingAtApplied: true, + leasesApplied: 1, + leasesSkipped: 1, + }) + + const day = '2026-08-01' + for (const resource of [ + 'email_receives_per_day', + 'email_sends_per_day', + ] as const) { + await meterA.initialize({ + resource, + day, + count: 1, + updatedAt: '2026-08-01T12:00:00.000Z', + }) + } + const firstPage = await meterA.exportCounters({ pageSize: 1 }) + expect(firstPage.truncated).toBe(true) + expect(firstPage.deletionShadow).toEqual({ + deletingAt: '2026-08-01 10:00:00', + writeLeases: [], + }) + const secondPage = await meterA.exportCounters({ + pageSize: 1, + startAfter: firstPage.nextStartAfter, + }) + expect(secondPage.deletionShadow).toBeNull() + + await expect(meterA.purge()).resolves.toEqual({ ok: true }) + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2026-08-01 10:00:00', + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + expect(await meterA.read({ resource: 'email_sends_per_day', day })).toEqual({ + outcome: 'needs_bootstrap', + }) + expect(await meterA.exportCounters({})).toEqual({ + counters: [], + storageBytesShadow: null, + packageServiceStatesShadow: [], + deletionShadow: { + deletingAt: '2026-08-01 10:00:00', + writeLeases: [], + }, + nextStartAfter: null, + truncated: false, + }) + await expect( + meterA.shadowAcquireWriteLease({ + token: 'eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee', + holder: 'test:post-purge', + acquiredAt: '2026-08-01 13:00:00', + }), + ).resolves.toEqual({ acquired: false }) + + expect(await meterB.readDeletionState()).toEqual({ + deletingAt: '2026-08-01 12:00:00', + }) + expect(await meterB.listWriteLeases({})).toMatchObject({ + leases: expect.arrayContaining([ + expect.objectContaining({ + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + }), + expect.objectContaining({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + }), + ]), + }) +}, 30_000) diff --git a/packages/worker/src/jobs/job-schedule-watchdog.ts b/packages/worker/src/jobs/job-schedule-watchdog.ts index 8b529a14e6..fe95dbf6b0 100644 --- a/packages/worker/src/jobs/job-schedule-watchdog.ts +++ b/packages/worker/src/jobs/job-schedule-watchdog.ts @@ -37,6 +37,7 @@ type JobScheduleWatchdogEnv = { APP_DB: D1Database JOB_MANAGER?: Env['JOB_MANAGER'] BUNDLE_ARTIFACTS_KV?: KVNamespace + USER_METER?: DurableObjectNamespace } export function shouldRunJobScheduleWatchdogCron(now: Date) { @@ -190,6 +191,7 @@ export async function runJobScheduleWatchdogTick(input: { const repaired = await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: row.user_id, + env: input.env, async write() { return advanceStuckSkippedJobNextRunAt({ db: input.env.APP_DB, diff --git a/packages/worker/src/jobs/service.ts b/packages/worker/src/jobs/service.ts index d0a028c471..93da1ddb99 100644 --- a/packages/worker/src/jobs/service.ts +++ b/packages/worker/src/jobs/service.ts @@ -667,6 +667,7 @@ export async function syncPackageJobsForPackage(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const desiredJobs = input.manifest.kody.jobs ?? {} const existingRows = await listJobRowsByUserId( @@ -784,6 +785,7 @@ export async function deletePackageJobsForSourceId(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const existingRows = await listJobRowsByUserId( input.env.APP_DB, @@ -859,6 +861,7 @@ export async function createJob(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: callerContext.user.userId, + env: input.env, async write() { await assertWithinEntitlement({ db: input.env.APP_DB, @@ -959,6 +962,7 @@ export async function updateJob(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: callerContext.user.userId, + env: input.env, async write() { const existingRow = await getJobRowById( input.env.APP_DB, @@ -1085,6 +1089,7 @@ export async function deleteJob(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const row = await getJobRowById( input.env.APP_DB, @@ -1163,6 +1168,8 @@ export async function executeJobOnce(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.job.userId, + env: input.env, + waitUntil: input.waitUntil, async write() { const started = new Date() let execution: JobExecutionResult @@ -1337,6 +1344,8 @@ export async function runJobNow(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, + waitUntil: input.waitUntil, async write() { const row = await getJobRowById( input.env.APP_DB, @@ -1480,6 +1489,8 @@ export async function runDueJobsForUser(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, + waitUntil: input.waitUntil, async write() { const now = input.now ?? new Date() const nowIso = now.toISOString() diff --git a/packages/worker/src/mcp-auth.ts b/packages/worker/src/mcp-auth.ts index d2adb7ffe4..c3339a1c9b 100644 --- a/packages/worker/src/mcp-auth.ts +++ b/packages/worker/src/mcp-auth.ts @@ -214,6 +214,7 @@ export async function handleMcpRequest({ db: env.APP_DB, stableUserId: mcpUser.userId, holder: `mcp:${request.method} ${url.pathname}`, + env, waitUntil: ctx.waitUntil.bind(ctx), write: async () => await fetchMcp( diff --git a/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-repair.ts b/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-repair.ts index 78a7a6f1d8..0e7c494b83 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-repair.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-repair.ts @@ -40,6 +40,7 @@ export const adminAccountWriteLeaseRepairCapability = defineDomainCapability( expectedAcquiredAt: args.expected_acquired_at, repairedByUserId, reason: args.reason, + env: ctx.env, }) return { repaired: result.repaired, diff --git a/packages/worker/src/mcp/memory/service.ts b/packages/worker/src/mcp/memory/service.ts index 846afba466..313f7a1f32 100644 --- a/packages/worker/src/mcp/memory/service.ts +++ b/packages/worker/src/mcp/memory/service.ts @@ -148,6 +148,7 @@ export async function acknowledgeSurfacedMemories(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const memoryIds = Array.from( new Set(input.memoryIds.filter((id) => id.trim().length > 0)), @@ -175,6 +176,8 @@ export async function upsertMemory(input: MemoryUpsertInput): Promise<{ return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, + waitUntil: input.waitUntil, async write() { const normalized = normalizeMemoryPayload(input) const now = new Date().toISOString() @@ -329,6 +332,7 @@ export async function deleteMemory( return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const existing = await getMemoryById( input.env.APP_DB, diff --git a/packages/worker/src/package-registry/service.ts b/packages/worker/src/package-registry/service.ts index 673946208f..f3c41c6b19 100644 --- a/packages/worker/src/package-registry/service.ts +++ b/packages/worker/src/package-registry/service.ts @@ -222,6 +222,8 @@ export async function refreshSavedPackageProjection(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, + waitUntil: input.waitUntil, async write() { const shouldRebuildArtifacts = input.rebuildArtifacts ?? true const loaded = shouldRebuildArtifacts @@ -455,6 +457,7 @@ export async function deleteSavedPackageProjection(input: { return await withAccountWriteLease({ db: input.env.APP_DB, stableUserId: input.userId, + env: input.env, async write() { const savedPackage = await getSavedPackageById(input.env.APP_DB, { userId: input.userId, diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index 2eac2d4ec2..79b844de68 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -4,9 +4,11 @@ import { userMeterMirrorUpdatedAtToken, userMeterPackageServiceStateStaleMs, type DailyEntitlementResource, + type UserMeterDeletionShadow, type UserMeterPackageServiceState, type UserMeterPackageServiceStatus, type UserMeterStorageBytesState, + type UserMeterWriteLeaseShadow, } from '#worker/entitlements/user-meter-do.ts' import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' @@ -19,6 +21,14 @@ type PackageServiceRow = { revision: number updatedAt: string } +type WriteLeaseRow = { + holder: string + acquiredAt: string +} +type DeletionState = { + deletingAt: string | null + leases: Map +} /** * In-memory UserMeter stub keyed by `idFromName` (stable userId) for node-unit @@ -32,6 +42,7 @@ export function createInMemoryUserMeterEnv() { string, Map >() + const deletionByUser = new Map() function counterKey(resource: string, day: string) { return `${resource}\0${day}` @@ -41,6 +52,17 @@ export function createInMemoryUserMeterEnv() { return `${packageId}\0${serviceName}` } + function deletionStateFor(userId: string): DeletionState { + const existing = deletionByUser.get(userId) + if (existing) return existing + const created: DeletionState = { + deletingAt: null, + leases: new Map(), + } + deletionByUser.set(userId, created) + return created + } + function meterFor(userId: string) { const existingRows = metersByUser.get(userId) const rows = existingRows ?? new Map() @@ -53,6 +75,8 @@ export function createInMemoryUserMeterEnv() { packageServicesByUser.set(userId, packageServices) } + const deletion = deletionStateFor(userId) + function readRow(resource: string, day: string) { return rows.get(counterKey(resource, day)) ?? null } @@ -105,6 +129,34 @@ export function createInMemoryUserMeterEnv() { }) } + function writeLeaseState( + token: string, + row: WriteLeaseRow, + ): UserMeterWriteLeaseShadow { + return { + token, + holder: row.holder, + acquiredAt: row.acquiredAt, + } + } + + function listWriteLeasesSorted() { + return [...deletion.leases.entries()] + .map(([token, row]) => writeLeaseState(token, row)) + .sort((left, right) => { + const byAcquired = left.acquiredAt.localeCompare(right.acquiredAt) + if (byAcquired !== 0) return byAcquired + return left.token.localeCompare(right.token) + }) + } + + function readDeletionShadow(): UserMeterDeletionShadow { + return { + deletingAt: deletion.deletingAt, + writeLeases: listWriteLeasesSorted(), + } + } + return { async initialize(input: { resource: string @@ -393,11 +445,121 @@ export function createInMemoryUserMeterEnv() { } return { applied, skipped } }, + async shadowMarkDeleting(input: { deletingAt: string }) { + if (deletion.deletingAt != null) { + return { deletingAt: deletion.deletingAt, created: false } + } + deletion.deletingAt = input.deletingAt + return { deletingAt: input.deletingAt, created: true } + }, + async shadowAcquireWriteLease(input: { + token: string + holder: string + acquiredAt: string + }) { + if (deletion.leases.has(input.token)) return { acquired: true } + if (deletion.deletingAt != null) return { acquired: false } + deletion.leases.set(input.token, { + holder: input.holder, + acquiredAt: input.acquiredAt, + }) + return { acquired: true } + }, + async shadowReleaseWriteLease(input: { token: string }) { + return { released: deletion.leases.delete(input.token) } + }, + async readDeletionState() { + return { deletingAt: deletion.deletingAt } + }, + async listWriteLeases( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ) { + const pageSize = + typeof input.pageSize === 'number' && Number.isFinite(input.pageSize) + ? Math.min(Math.max(Math.trunc(input.pageSize), 1), 500) + : 100 + const all = listWriteLeasesSorted() + let startIndex = 0 + if ( + typeof input.startAfter === 'string' && + input.startAfter.length > 0 + ) { + try { + const parsed = JSON.parse(input.startAfter) as unknown + if ( + Array.isArray(parsed) && + parsed.length === 2 && + typeof parsed[0] === 'string' && + typeof parsed[1] === 'string' + ) { + const acquiredAt = parsed[0] + const token = parsed[1] + startIndex = + all.findIndex( + (row) => row.acquiredAt === acquiredAt && row.token === token, + ) + 1 + } + } catch { + startIndex = 0 + } + } + const page = all.slice(startIndex, startIndex + pageSize) + const truncated = startIndex + pageSize < all.length + const last = page[page.length - 1] + return { + leases: page, + nextStartAfter: + truncated && last + ? JSON.stringify([last.acquiredAt, last.token]) + : null, + truncated, + } + }, + async countActiveWriteLeases() { + return { count: deletion.leases.size } + }, + async bootstrapDeletionState(input: { + deletingAt?: string | null + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) { + let leasesApplied = 0 + let leasesSkipped = 0 + for (const lease of input.leases ?? []) { + if (deletion.leases.has(lease.token)) { + leasesSkipped += 1 + continue + } + const result = await this.shadowAcquireWriteLease(lease) + if (result.acquired) leasesApplied += 1 + else leasesSkipped += 1 + } + let deletingAtApplied = false + if ( + typeof input.deletingAt === 'string' && + input.deletingAt.length > 0 + ) { + const marked = await this.shadowMarkDeleting({ + deletingAt: input.deletingAt, + }) + deletingAtApplied = marked.created + } + return { deletingAtApplied, leasesApplied, leasesSkipped } + }, async purge() { + const deletingAt = deletion.deletingAt rows.clear() storageByUser.delete(userId) packageServices.clear() packageServicesByUser.delete(userId) + deletion.leases.clear() + deletion.deletingAt = deletingAt return { ok: true as const } }, async exportCounters( @@ -431,10 +593,12 @@ export function createInMemoryUserMeterEnv() { const packageServiceStatesShadow = includeShadows ? listPackageServiceStatesSorted() : null + const deletionShadow = includeShadows ? readDeletionShadow() : null return { counters, storageBytesShadow, packageServiceStatesShadow, + deletionShadow, nextStartAfter: null, truncated: false, } @@ -454,6 +618,7 @@ export function createInMemoryUserMeterEnv() { metersByUser, storageByUser, packageServicesByUser, + deletionByUser, async seed(input: { userId: string resource: DailyEntitlementResource From 2e92f03772fd0d3eb42a43a89778e18c1a7dfae8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 09:25:13 +0000 Subject: [PATCH 21/39] docs(account): inventory deletion shadow Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 29 +++++--- .../contributing/architecture/entitlements.md | 67 ++++++++++++++++--- .../contributing/architecture/primitives.yaml | 8 ++- .../worker/src/account/export.node.test.ts | 18 ++++- packages/worker/src/account/export.ts | 15 ++++- .../worker/src/account/user-owned-surfaces.ts | 2 +- 6 files changed, 116 insertions(+), 23 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 6bbeda7bb7..c55b9a5129 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -539,7 +539,7 @@ Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP column inside the DO because the object identity is the user. SQLite ownership (schema version tracked in `user_meter_meta`; current version -**5**): +**6**): - `daily_counters` — authoritative UTC-day counters for `email_sends_per_day`, `email_receives_per_day`, `execute_calls_per_day`, and @@ -566,13 +566,23 @@ SQLite ownership (schema version tracked in `user_meter_meta`; current version (`listPackageServiceStates`, `countRunningPackageServices`, `bootstrapPackageServiceStates`) mirror D1 semantics for future parity review only. +- `deletion_state` / `account_write_leases` — **shadow** of D1 + `users.deleting_at` and active `account_write_leases` (singleton + `deleting_at`; lease rows `token` / `holder` / `acquired_at`). Added in schema + v6. Populated by optional best-effort dual-writes from + `account/deletion-state.ts` after successful D1 mark/acquire/release/repair; + never read for fencing, list, repair, or drain in expand-phase slice 5 Phase + A. `purge()` preserves an existing deleting tombstone across `deleteAll` so a + later cutover cannot reopen writes. Cutover-support RPCs (`readDeletionState`, + `listWriteLeases`, `countActiveWriteLeases`, `bootstrapDeletionState`) exist + for future parity review only. Retention is self-enforced inside the DO: every read/write path opportunistically deletes counter and claim rows older than seven UTC days (`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; the window covers timezone edge cases, recent account exports, and inbound -retries. Shadow storage-byte and package-service liveness rows are not -time-pruned. +retries. Shadow storage-byte, package-service liveness, and deletion-fence rows +are not time-pruned. **Expand-phase D1 mirrors (daily counters only):** enforcement and point reads are authoritative in UserMeter for daily counters. D1 @@ -591,12 +601,13 @@ writes cannot overwrite newer state. See daily paths never read D1 for enforcement. Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan; -`deleteAll` clears counters, claims, and all shadow state including storage -bytes and package-service liveness). Account export pages -`UserMeter.exportCounters` through the `user_meter` manifest section / -`account_export_section` (daily counters plus additive `storageBytesShadow` and -`packageServiceStatesShadow` on the first page only when present; shadow fields -are non-authoritative). +`deleteAll` clears counters, claims, storage-byte shadow, package-service +shadow, and write-lease shadow, while preserving an existing deleting +tombstone). Account export pages `UserMeter.exportCounters` through the +`user_meter` manifest section / `account_export_section` (daily counters plus +additive `storageBytesShadow`, `packageServiceStatesShadow`, and +`deletionShadow` on the first page only when present; shadow fields are +non-authoritative). ## Durable Objects (`Mailbox`) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index f72c681906..20235240c1 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -130,6 +130,12 @@ v5 adds an optional per-service shadow table as **best-effort future-cutover support only** — see [Package service liveness — UserMeter shadow](#package-service-liveness--usermeter-shadow-expand-phase-slice-4-phase-a). +**D1 account-deletion write fencing** (`users.deleting_at` / +`account_write_leases`) stays authoritative for fencing, list, repair, and +drain. UserMeter schema v6 adds an optional deletion-fence / write-lease shadow +as **best-effort future-cutover support only** — see +[Account-deletion write fencing — UserMeter shadow](#account-deletion-write-fencing--usermeter-shadow-expand-phase-slice-5-phase-a). + StorageRunner bucket `estimatedBytes` and the per-bucket inventory in `user_storage_buckets` stay a **separate** quota component. StorageRunner write chokepoints pass `getCurrent` as a check-only composed total (D1 payload bytes @@ -276,6 +282,54 @@ remains on D1. **Account purge:** `UserMeter.purge()` clears package-service shadow rows with the rest of DO state via `deleteAll`. +### Account-deletion write fencing — UserMeter shadow (expand phase slice 5, Phase A) + +D1 `users.deleting_at`, `account_write_leases`, and +`account_write_lease_repairs` remain the **sole authority** for deletion +fencing, lease acquire/release, active-lease list, audited repair, and deletion +drain waits in this PR. Nothing in Phase A switches `withAccountWriteLease`, +`markAccountDeleting`, `listActiveAccountWriteLeases`, or +`repairAccountWriteLease` onto UserMeter reads — public errors, ALS nested-lease +reuse, and D1 `waitUntil` release semantics stay identical when optional `env` +is omitted. + +UserMeter schema **v6** adds a singleton `deletion_state.deleting_at` tombstone +plus per-token `account_write_leases` rows (`token`, `holder`, `acquired_at`) as +**best-effort shadow / future-cutover support only**. User scope is the DO +identity — there is no `user_id` column. Shadow rows are never read for fencing, +list, repair, or drain in this slice. + +**Optional dual-write from `account/deletion-state.ts`:** after a **successful** +D1 mark/acquire/release/repair, callers may pass optional `env?: UserMeterEnv` +(and `waitUntil` when available). The helper schedules a best-effort UserMeter +shadow (`shadowMarkDeleting`, `shadowAcquireWriteLease`, +`shadowReleaseWriteLease`). Acquire shadows start during the write; ordered +release shadows run after authoritative D1 release. When `waitUntil` is +provided, shadows detach through it; when omitted, mark/release/repair shadows +are awaited (still non-rejecting) so middleware and other non-`waitUntil` +callers cannot return with a stale shadow lease or missing tombstone. Missing +`USER_METER` is a no-op; failures log +`account-deletion-user-meter-shadow-failed` and never alter D1 results or +errors. Mark is monotonic (first `deleting_at` wins). Acquire is rejected when a +deleting tombstone is present and is idempotent for an already-held token. + +**Account export:** `UserMeter.exportCounters` returns additive `deletionShadow` +(`deletingAt` + `writeLeases`) on the first page only (`startAfter` absent); +later pages return `null`. Section totals count the tombstone (when set) and +each lease once when present; the field is explicitly non-authoritative — +authoritative fencing remains on D1. + +**Account purge:** `UserMeter.purge()` clears counters, claims, storage shadow, +package-service shadow, and write-lease shadow via `deleteAll`, then restores +any pre-existing deleting tombstone so a later authority cutover cannot reopen +writes for a purged account. + +Non-email production call sites that hold `env` pass it into deletion-state APIs +(`withAccountWriteLease`, `markAccountDeleting`, `repairAccountWriteLease`) and +forward existing `waitUntil` when available. Email paths stay untouched in this +slice (`system:email` skip unchanged). Assert-only readers +(`assertAccountWritable*`, `listActiveAccountWriteLeases`) remain D1-only. + ### Future package-service authority flip (contract follow-up) A separate **high-risk contract PR** — not a merge blocker for Phase A — will @@ -290,10 +344,6 @@ likely stays the enumeration index** for account export, deletion, and admin discovery until an alternate inventory exists — UserMeter would become the running-count authority first, not a wholesale replacement for every D1 reader. -**Remaining expand roadmap:** slice 4 Phase A (this shadow slice) is additive -only; slice 5 — account-deletion write fencing — follows independently. Storage -authority flip remains the separate contract follow-up above. - ### Future storage authority flip (contract follow-up) A separate contract PR will flip D1 payload storage-byte authority into @@ -304,10 +354,11 @@ Only then do reads, reserves, and reconciliation switch to UserMeter-first with D1 as mirror/cursor. Until that flip, shadow divergence is acceptable; D1 remains the contract. -**Remaining UserMeter expand roadmap:** slice 4 Phase A — -`package_service_states` UserMeter shadow (this PR; D1 authority unchanged); -slice 5 — account-deletion write fencing. Package-service and storage authority -flips are separate high-risk contract follow-ups after soak/parity review. +**Remaining UserMeter expand roadmap:** slice 5 Phase A — account-deletion write +fencing shadow (this PR; D1 authority unchanged; non-email dual-write wired). +Package-service and storage authority flips are separate high-risk contract +follow-ups after soak/parity review. Deletion-fence authority flip is a separate +high-risk contract follow-up after shadow/D1 parity review. **Daily-counter mirror retirement:** dropping D1 `entitlement_daily_counters` waits until reporting-off-D1 work merges and mirror parity is verified in diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index 776bcfe2f7..8946b59507 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -516,12 +516,14 @@ primitives: name: User meter summary: Per-user UserMeter DO SQLite for daily entitlement counters, inbound - delivery idempotency, expand-phase D1 storage-byte shadow, and - expand-phase package-service liveness shadow (D1 authority unchanged in - slice 4 Phase A). + delivery idempotency, expand-phase D1 storage-byte shadow, expand-phase + package-service liveness shadow, and expand-phase account-deletion + write-fence shadow (D1 remains sole fence/lease authority in slice 5 Phase + A). code: - packages/worker/src/entitlements/user-meter-do.ts - packages/worker/src/entitlements/user-meter-client.ts + - packages/worker/src/account/deletion-state.ts docs: - docs/contributing/architecture/data-storage.md - docs/contributing/architecture/entitlements.md diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 22fb0971d7..4aa9e87617 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1590,6 +1590,16 @@ test('account export includes user_meter counters, pages them, and warns on trun mirrorUpdatedAt: 'r/00000000000000000002', }, ] + const deletionShadow = { + deletingAt: '2026-07-31 03:10:00' as string | null, + writeLeases: [ + { + token: '11111111-1111-4111-8111-111111111111', + holder: 'test:export', + acquiredAt: '2026-07-31 03:00:00', + }, + ], + } const exportCounters = vi.fn( async (input: { pageSize?: number; startAfter?: string | null }) => { const pageSize = input.pageSize ?? 100 @@ -1608,6 +1618,7 @@ test('account export includes user_meter counters, pages them, and warns on trun packageServiceStatesShadow: isFirstPage ? packageServiceStatesShadow : null, + deletionShadow: isFirstPage ? deletionShadow : null, nextStartAfter: truncated ? `${page.at(-1)!.day}:${page.at(-1)!.resource}` : null, @@ -1630,11 +1641,13 @@ test('account export includes user_meter counters, pages them, and warns on trun mcpUserId: 'user-aaa', }) expect(idFromName).toHaveBeenCalledWith('user-aaa') - expect(accountExport.manifest.sections.user_meter?.count).toBe(5) + // 3 counters + storageBytesShadow + 1 package-service shadow + deletingAt + 1 lease + expect(accountExport.manifest.sections.user_meter?.count).toBe(7) expect(accountExport.durableObjects.userMeter).toEqual({ counters, storageBytesShadow, packageServiceStatesShadow, + deletionShadow, nextStartAfter: null, truncated: false, }) @@ -1654,6 +1667,7 @@ test('account export includes user_meter counters, pages them, and warns on trun expect(first.items).toEqual(counters.slice(0, 2)) expect(first.storageBytesShadow).toEqual(storageBytesShadow) expect(first.packageServiceStatesShadow).toEqual(packageServiceStatesShadow) + expect(first.deletionShadow).toEqual(deletionShadow) expect(first.truncated).toBe(true) expect(first.nextStartAfter).toBe('2026-07-30:execute_calls_per_day') @@ -1668,6 +1682,7 @@ test('account export includes user_meter counters, pages them, and warns on trun expect(second.items).toEqual(counters.slice(2)) expect(second.storageBytesShadow).toBeNull() expect(second.packageServiceStatesShadow).toBeNull() + expect(second.deletionShadow).toBeNull() expect(second.truncated).toBe(false) expect(second.nextStartAfter).toBeNull() expect(exportCounters).toHaveBeenCalledWith( @@ -1681,6 +1696,7 @@ test('account export includes user_meter counters, pages them, and warns on trun counters: [counters[0]!], storageBytesShadow: null, packageServiceStatesShadow: null, + deletionShadow: null, nextStartAfter: 'cursor-more', truncated: true, })) diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index b3a31eb93e..e3d9f65aac 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -209,12 +209,18 @@ export type AccountExportArtifactRepo = { type RunRecordsExportPayload = Awaited> function countUserMeterExportEntries(result: UserMeterExportResult): number { + const deletionShadowCount = + result.deletionShadow == null + ? 0 + : (result.deletionShadow.deletingAt == null ? 0 : 1) + + result.deletionShadow.writeLeases.length return ( result.counters.length + (result.storageBytesShadow == null ? 0 : 1) + (result.packageServiceStatesShadow == null ? 0 - : result.packageServiceStatesShadow.length) + : result.packageServiceStatesShadow.length) + + deletionShadowCount ) } @@ -305,6 +311,12 @@ export type AccountExportSectionResult = { * `null`. */ packageServiceStatesShadow?: UserMeterExportResult['packageServiceStatesShadow'] + /** + * Non-authoritative UserMeter deletion-fence / write-lease shadow. Present + * only on the first `user_meter` page (`startAfter` absent); later pages + * set it to `null`. + */ + deletionShadow?: UserMeterExportResult['deletionShadow'] } function normalizePageSize(pageSize: number | undefined) { @@ -2002,6 +2014,7 @@ export async function readAccountExportSection(input: { items: page.counters, storageBytesShadow: page.storageBytesShadow, packageServiceStatesShadow: page.packageServiceStatesShadow, + deletionShadow: page.deletionShadow, truncated: page.truncated, nextStartAfter: page.nextStartAfter, pageSize, diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 065e34438b..5afc89aa7d 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -132,7 +132,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray Date: Sat, 1 Aug 2026 09:52:14 +0000 Subject: [PATCH 22/39] fix(account): reconcile and sanitize deletion shadows Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 18 ++- .../contributing/architecture/entitlements.md | 29 ++-- .../src/account/deletion-state.node.test.ts | 137 +++++++++++++++++- packages/worker/src/account/deletion-state.ts | 69 +++++---- .../worker/src/account/export.node.test.ts | 3 +- packages/worker/src/account/export.ts | 2 +- .../worker/src/account/user-owned-surfaces.ts | 2 +- .../src/app/account-deletion.node.test.ts | 39 +++++ .../worker/src/entitlements/user-meter-do.ts | 114 ++++++++++++++- .../entitlements/user-meter.workers.test.ts | 66 ++++++++- .../worker/src/test-support/user-meter.ts | 44 +++++- 11 files changed, 456 insertions(+), 67 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index d604f4dae4..89408dcad8 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -572,17 +572,23 @@ SQLite ownership (schema version tracked in `user_meter_meta`; current version v6. Populated by optional best-effort dual-writes from `account/deletion-state.ts` after successful D1 mark/acquire/release/repair; never read for fencing, list, repair, or drain in expand-phase slice 5 Phase - A. `purge()` preserves an existing deleting tombstone across `deleteAll` so a - later cutover cannot reopen writes. Cutover-support RPCs (`readDeletionState`, - `listWriteLeases`, `countActiveWriteLeases`, `bootstrapDeletionState`) exist - for future parity review only. + A. On mark, after D1 `deleting_at` is set, the helper loads active D1 leases + and calls `shadowReplaceDeletionState` (DO-serialized tombstone set/preserve + + exact lease-set replace) so stale unreleased shadow rows are cleared when D1 + drain reaches zero and active D1 leases are preserved. `purge()` preserves an + existing deleting tombstone across `deleteAll` so a later cutover cannot + reopen writes. Cutover-support RPCs (`readDeletionState`, `listWriteLeases`, + `countActiveWriteLeases`, `bootstrapDeletionState`) exist for future parity + review only; account export emits a sanitized `deletionShadow` without raw + token/holder. Retention is self-enforced inside the DO: every read/write path opportunistically deletes counter and claim rows older than seven UTC days (`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; the window covers timezone edge cases, recent account exports, and inbound -retries. Shadow storage-byte, package-service liveness, and deletion-fence rows -are not time-pruned. +retries. Shadow storage-byte and package-service liveness rows are not +time-pruned. Deletion-fence lease shadows are bounded by the D1-backed mark +replace path above rather than time retention. **Expand-phase D1 mirrors (daily counters only):** enforcement and point reads are authoritative in UserMeter for daily counters. D1 diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index e3ef761afc..5ea3edf7eb 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -302,22 +302,29 @@ list, repair, or drain in this slice. **Optional dual-write from `account/deletion-state.ts`:** after a **successful** D1 mark/acquire/release/repair, callers may pass optional `env?: UserMeterEnv` (and `waitUntil` when available). The helper schedules a best-effort UserMeter -shadow (`shadowMarkDeleting`, `shadowAcquireWriteLease`, -`shadowReleaseWriteLease`). Acquire shadows start during the write; ordered -release shadows run after authoritative D1 release. When `waitUntil` is -provided, shadows detach through it; when omitted, mark/release/repair shadows -are awaited (still non-rejecting) so middleware and other non-`waitUntil` -callers cannot return with a stale shadow lease or missing tombstone. Missing -`USER_METER` is a no-op; failures log +shadow (`shadowReplaceDeletionState` on mark, `shadowAcquireWriteLease`, +`shadowReleaseWriteLease`). After D1 `deleting_at` is set, mark loads the +authoritative active D1 lease rows and calls `shadowReplaceDeletionState`, which +under DO serialization sets/preserves the tombstone then deletes and re-inserts +exactly that lease set — so a prior failed shadow release cannot leave stale +rows, and a deletion retry with zero D1 leases clears the shadow set. Acquire +shadows start during the write; ordered release shadows run after authoritative +D1 release. When `waitUntil` is provided, shadows detach through it; when +omitted, mark/release/repair shadows are awaited (still non-rejecting) so +middleware and other non-`waitUntil` callers cannot return with a stale shadow +lease or missing tombstone. Missing `USER_METER` is a no-op; failures log `account-deletion-user-meter-shadow-failed` and never alter D1 results or errors. Mark is monotonic (first `deleting_at` wins). Acquire is rejected when a deleting tombstone is present and is idempotent for an already-held token. +Cutover/list RPCs (`listWriteLeases`, `bootstrapDeletionState`) may still carry +token/holder for parity review. **Account export:** `UserMeter.exportCounters` returns additive `deletionShadow` -(`deletingAt` + `writeLeases`) on the first page only (`startAfter` absent); -later pages return `null`. Section totals count the tombstone (when set) and -each lease once when present; the field is explicitly non-authoritative — -authoritative fencing remains on D1. +on the first page only (`startAfter` absent); later pages return `null`. The +export shape is sanitized (`deletingAt`, `activeWriteLeaseCount`, and +`writeLeases` with `acquiredAt` only — no raw token or holder). Section totals +count the tombstone (when set) plus `activeWriteLeaseCount`; the field is +explicitly non-authoritative — authoritative fencing remains on D1. **Account purge:** `UserMeter.purge()` clears counters, claims, storage shadow, package-service shadow, and write-lease shadow via `deleteAll`, then restores diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index 0a9fc0f3f7..5838b76585 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -419,6 +419,9 @@ function createFailingShadowEnv(): UserMeterEnv { async shadowMarkDeleting() { throw new Error('shadow mark failed') }, + async shadowReplaceDeletionState() { + throw new Error('shadow replace failed') + }, async shadowAcquireWriteLease() { throw new Error('shadow acquire failed') }, @@ -518,10 +521,19 @@ function createGatedShadowEnv(input: { get(id: DurableObjectId) { const stub = namespace.get(id) return { - shadowMarkDeleting: async (args: { deletingAt: string }) => { + shadowMarkDeleting: (args: { deletingAt: string }) => + stub.shadowMarkDeleting(args), + shadowReplaceDeletionState: async (args: { + deletingAt: string + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) => { gates.markEntered = true if (input.markGate) await input.markGate - return stub.shadowMarkDeleting(args) + return stub.shadowReplaceDeletionState(args) }, shadowAcquireWriteLease: (args: { token: string @@ -535,6 +547,9 @@ function createGatedShadowEnv(input: { }, countActiveWriteLeases: () => stub.countActiveWriteLeases(), readDeletionState: () => stub.readDeletionState(), + listWriteLeases: ( + args: { pageSize?: number; startAfter?: string | null } = {}, + ) => stub.listWriteLeases(args), } }, }, @@ -711,3 +726,121 @@ test('with waitUntil, lease return completes before release shadow settles', asy await drain.drain() expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) }) + +test('markAccountDeleting replaces UserMeter lease shadows from authoritative D1', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const drain = createWaitUntilDrain() + + // Stale shadow left behind by a failed release (no matching D1 lease). + await meterA.shadowAcquireWriteLease({ + token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + holder: 'test:stale-shadow', + acquiredAt: '2026-01-01 00:00:00', + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + + let heldToken = '' + let heldAcquiredAt = '' + let startWrite: () => void = () => undefined + let finishWrite: () => void = () => undefined + const started = new Promise((resolve) => { + startWrite = resolve + }) + const finish = new Promise((resolve) => { + finishWrite = resolve + }) + const operation = withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:active-d1', + env: meter.env, + waitUntil: drain.waitUntil, + async write() { + const [lease] = await listActiveAccountWriteLeases(db, 'user-a') + heldToken = lease!.token + heldAcquiredAt = lease!.acquired_at + startWrite() + await finish + return 'lost' + }, + }) + await started + await drain.drain() + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 2 }) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: meter.env, + waitUntil: drain.waitUntil, + }), + ).resolves.toBe(1) + await drain.drain() + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2099-01-01 00:00:00', + }) + // Stale shadow removed; active D1 lease preserved under the tombstone. + expect(await meterA.listWriteLeases({})).toEqual({ + leases: [ + { + token: heldToken, + holder: 'test:active-d1', + acquiredAt: heldAcquiredAt, + }, + ], + nextStartAfter: null, + truncated: false, + }) + expect(await meterA.exportCounters({})).toMatchObject({ + deletionShadow: { + deletingAt: '2099-01-01 00:00:00', + activeWriteLeaseCount: 1, + writeLeases: [{ acquiredAt: heldAcquiredAt }], + }, + }) + + await repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: heldToken, + expectedAcquiredAt: heldAcquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + waitUntil: drain.waitUntil, + }) + await drain.drain() + finishWrite() + await expect(operation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + + // Failed release after D1 drain can leave a stale shadow; retry clears it. + await meterA.shadowReplaceDeletionState({ + deletingAt: '2099-01-01 00:00:00', + leases: [ + { + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:post-drain-stale', + acquiredAt: '2099-01-01 00:01:00', + }, + ], + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: meter.env, + }), + ).resolves.toBe(0) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2099-01-01 00:00:00', + }) +}) diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index 495da004e0..00ba34f431 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -77,40 +77,49 @@ export async function markAccountDeleting(input: { if ((result.meta.changes ?? 0) !== 1) { throw new Error('Account could not be marked for deletion.') } - const row = await input.db + const userRow = await input.db .prepare( - `SELECT COUNT(*) AS count - FROM account_write_leases - WHERE user_id = ( - SELECT stable_user_id FROM users WHERE id = ? - ) AND released_at IS NULL`, + `SELECT stable_user_id, deleting_at + FROM users + WHERE id = ?`, ) .bind(input.dbUserId) - .first<{ count: number }>() - if (input.env && userMeterNamespace(input.env)) { - const userRow = await input.db - .prepare( - `SELECT stable_user_id, deleting_at - FROM users - WHERE id = ?`, - ) - .bind(input.dbUserId) - .first<{ stable_user_id: string; deleting_at: string | null }>() - if (userRow?.stable_user_id && userRow.deleting_at) { - const { stable_user_id: stableUserId, deleting_at: deletingAt } = userRow - const markShadowPromise = scheduleAccountDeletionShadow({ - env: input.env, - waitUntil: input.waitUntil, - task: async (env) => { - await userMeterRpc({ env, userId: stableUserId }).shadowMarkDeleting({ - deletingAt, - }) - }, - }) - if (!input.waitUntil) await markShadowPromise - } + .first<{ stable_user_id: string; deleting_at: string | null }>() + const stableUserId = userRow?.stable_user_id + const deletingAt = userRow?.deleting_at + // Authoritative active D1 leases after deleting_at is set (new acquires + // are fenced). Returned count stays D1-authoritative for drain waits. + const activeLeases = stableUserId + ? await listActiveAccountWriteLeases(input.db, stableUserId) + : [] + if ( + input.env && + userMeterNamespace(input.env) && + stableUserId && + deletingAt + ) { + const markShadowPromise = scheduleAccountDeletionShadow({ + env: input.env, + waitUntil: input.waitUntil, + task: async (env) => { + // Replace (not append) so a prior failed shadow release cannot + // leave stale rows after D1 drain; empty list clears shadows. + await userMeterRpc({ + env, + userId: stableUserId, + }).shadowReplaceDeletionState({ + deletingAt, + leases: activeLeases.map((lease) => ({ + token: lease.token, + holder: lease.holder, + acquiredAt: lease.acquired_at, + })), + }) + }, + }) + if (!input.waitUntil) await markShadowPromise } - return Number(row?.count ?? 0) + return activeLeases.length } export async function assertAccountWritableDb( diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 4aa9e87617..4c757e37dc 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -1592,10 +1592,9 @@ test('account export includes user_meter counters, pages them, and warns on trun ] const deletionShadow = { deletingAt: '2026-07-31 03:10:00' as string | null, + activeWriteLeaseCount: 1, writeLeases: [ { - token: '11111111-1111-4111-8111-111111111111', - holder: 'test:export', acquiredAt: '2026-07-31 03:00:00', }, ], diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index e3d9f65aac..c1136ed088 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -213,7 +213,7 @@ function countUserMeterExportEntries(result: UserMeterExportResult): number { result.deletionShadow == null ? 0 : (result.deletionShadow.deletingAt == null ? 0 : 1) + - result.deletionShadow.writeLeases.length + result.deletionShadow.activeWriteLeaseCount return ( result.counters.length + (result.storageBytesShadow == null ? 0 : 1) + diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index 5afc89aa7d..b939ea036d 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -132,7 +132,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray, meta: { changes: 0 } } } + if ( + lower === + 'select stable_user_id, deleting_at from users where id = ?' + ) { + const numericId = Number(params[0]) + results = (rows.users ?? []) + .filter((row) => Number(row['id']) === numericId) + .map((row) => ({ + stable_user_id: row['stable_user_id'], + deleting_at: row['deleting_at'] ?? null, + })) + return { results: results as Array, meta: { changes: 0 } } + } + if ( + lower.startsWith( + 'select token, holder, acquired_at from account_write_leases', + ) + ) { + const stableUserId = params[0] as string + results = (rows.account_write_leases ?? []) + .filter( + (row) => + row['user_id'] === stableUserId && + row['released_at'] == null, + ) + .map((row) => ({ + token: row['token'], + holder: row['holder'], + acquired_at: row['acquired_at'], + })) + .sort((left, right) => { + const byAcquired = String(left.acquired_at).localeCompare( + String(right.acquired_at), + ) + if (byAcquired !== 0) return byAcquired + return String(left.token).localeCompare(String(right.token)) + }) + return { results: results as Array, meta: { changes: 0 } } + } if ( lower === 'select id from saved_packages where user_id = ? and has_app = 1' diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index c4eaf94849..722b5289ea 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -180,16 +180,36 @@ export type UserMeterWriteLeaseShadow = { acquiredAt: string } +/** + * Internal cutover / list shape. Includes token and holder for parity with D1 + * `account_write_leases`; never emit this from account export. + */ export type UserMeterDeletionShadow = { deletingAt: string | null writeLeases: Array } +/** + * Account-export deletion shadow. Omits raw lease token and holder; retains + * only deleting tombstone presence, active lease count, and acquired_at. + */ +export type UserMeterDeletionShadowExport = { + deletingAt: string | null + activeWriteLeaseCount: number + writeLeases: Array<{ acquiredAt: string }> +} + export type UserMeterShadowMarkDeletingResult = { deletingAt: string created: boolean } +export type UserMeterShadowReplaceDeletionStateResult = { + deletingAt: string + created: boolean + leaseCount: number +} + export type UserMeterShadowAcquireWriteLeaseResult = { acquired: boolean } @@ -231,9 +251,11 @@ export type UserMeterExportResult = { /** * Non-authoritative D1 deletion-fence / write-lease shadow. Emitted only * on the first export page (`startAfter` absent); subsequent pages return - * `null` so paged consumers never double-count the inventory. + * `null` so paged consumers never double-count the inventory. Sanitized: + * excludes raw lease token and holder (see + * {@link UserMeterDeletionShadowExport}). */ - deletionShadow: UserMeterDeletionShadow | null + deletionShadow: UserMeterDeletionShadowExport | null nextStartAfter: string | null truncated: boolean } @@ -1390,10 +1412,14 @@ class UserMeterBase extends DurableObject { return rows.map((row) => this.writeLeaseFromRow(row)) } - private readDeletionShadow(): UserMeterDeletionShadow { + private readDeletionShadowExport(): UserMeterDeletionShadowExport { + const writeLeases = this.listAllWriteLeaseRows().map((lease) => ({ + acquiredAt: lease.acquiredAt, + })) return { deletingAt: this.readDeletingAt(), - writeLeases: this.listAllWriteLeaseRows(), + activeWriteLeaseCount: writeLeases.length, + writeLeases, } } @@ -1406,7 +1432,7 @@ class UserMeterBase extends DurableObject { if (existing != null) { return { deletingAt: existing, created: false } } - this.ctx.storage.sql.exec( + const cursor = this.ctx.storage.sql.exec( `INSERT INTO deletion_state (id, deleting_at) VALUES (?, ?) ON CONFLICT(id) DO NOTHING`, @@ -1414,7 +1440,67 @@ class UserMeterBase extends DurableObject { deletingAt, ) const stored = this.readDeletingAt() ?? deletingAt - return { deletingAt: stored, created: stored === deletingAt } + return { deletingAt: stored, created: cursor.rowsWritten > 0 } + } + + /** + * Atomically set/preserve the deleting tombstone and replace the shadow + * lease set with exactly the supplied active D1 leases. Used by + * `markAccountDeleting` so stale unreleased shadow rows cannot linger + * after D1 drain (empty lease list clears all shadow leases). + */ + async shadowReplaceDeletionState(input: { + deletingAt: string + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }): Promise { + const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + const leasesByToken = new Map< + string, + { token: string; holder: string; acquiredAt: string } + >() + for (const lease of input.leases ?? []) { + const token = assertWriteLeaseToken(lease.token) + leasesByToken.set(token, { + token, + holder: assertWriteLeaseHolder(lease.holder), + acquiredAt: assertDeletionTimestamp('acquiredAt', lease.acquiredAt), + }) + } + const leases = [...leasesByToken.values()] + return await this.ctx.blockConcurrencyWhile(async () => { + const existing = this.readDeletingAt() + let created = false + if (existing == null) { + const cursor = this.ctx.storage.sql.exec( + `INSERT INTO deletion_state (id, deleting_at) + VALUES (?, ?) + ON CONFLICT(id) DO NOTHING`, + deletionStateRowId, + deletingAt, + ) + created = cursor.rowsWritten > 0 + } + const stored = this.readDeletingAt() ?? deletingAt + this.ctx.storage.sql.exec(`DELETE FROM account_write_leases`) + for (const lease of leases) { + this.ctx.storage.sql.exec( + `INSERT INTO account_write_leases (token, holder, acquired_at) + VALUES (?, ?, ?)`, + lease.token, + lease.holder, + lease.acquiredAt, + ) + } + return { + deletingAt: stored, + created, + leaseCount: leases.length, + } + }) } /** Expand-phase shadow lease acquire; fencing still uses D1. */ @@ -1642,7 +1728,9 @@ class UserMeterBase extends DurableObject { const packageServiceStatesShadow = includeShadows ? this.listAllPackageServiceRows() : null - const deletionShadow = includeShadows ? this.readDeletionShadow() : null + const deletionShadow = includeShadows + ? this.readDeletionShadowExport() + : null const last = pageRows[pageRows.length - 1] return { counters, @@ -1761,6 +1849,18 @@ export type UserMeterRpc = { shadowMarkDeleting: (input: { deletingAt: string }) => Promise + /** + * Atomically set/preserve deleting tombstone and replace shadow leases + * with the supplied active D1 set (empty clears stale shadows). + */ + shadowReplaceDeletionState: (input: { + deletingAt: string + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) => Promise /** Expand-phase shadow lease acquire; fencing still uses D1. */ shadowAcquireWriteLease: (input: { token: string diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index 7de59f66b3..a43054c3a1 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -562,7 +562,11 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], - deletionShadow: { deletingAt: null, writeLeases: [] }, + deletionShadow: { + deletingAt: null, + activeWriteLeaseCount: 0, + writeLeases: [], + }, nextStartAfter: null, truncated: false, }) @@ -690,7 +694,11 @@ test('UserMeter purge blocks concurrent RPCs across deleteAll and schema restore counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], - deletionShadow: { deletingAt: null, writeLeases: [] }, + deletionShadow: { + deletingAt: null, + activeWriteLeaseCount: 0, + writeLeases: [], + }, nextStartAfter: null, truncated: false, }) @@ -995,6 +1003,7 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a expect(firstPage.packageServiceStatesShadow).toEqual([]) expect(firstPage.deletionShadow).toEqual({ deletingAt: null, + activeWriteLeaseCount: 0, writeLeases: [], }) @@ -1021,7 +1030,11 @@ test('UserMeter storage RPCs, export shadow field, and purge work additively', a counters: [], storageBytesShadow: null, packageServiceStatesShadow: [], - deletionShadow: { deletingAt: null, writeLeases: [] }, + deletionShadow: { + deletingAt: null, + activeWriteLeaseCount: 0, + writeLeases: [], + }, nextStartAfter: null, truncated: false, }) @@ -1370,6 +1383,7 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve expect(firstPage.truncated).toBe(true) expect(firstPage.deletionShadow).toEqual({ deletingAt: '2026-08-01 10:00:00', + activeWriteLeaseCount: 0, writeLeases: [], }) const secondPage = await meterA.exportCounters({ @@ -1392,6 +1406,7 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve packageServiceStatesShadow: [], deletionShadow: { deletingAt: '2026-08-01 10:00:00', + activeWriteLeaseCount: 0, writeLeases: [], }, nextStartAfter: null, @@ -1418,4 +1433,49 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve }), ]), }) + + const replaced = await meterB.shadowReplaceDeletionState({ + deletingAt: '2026-08-01 14:00:00', + leases: [ + { + token: 'ffffffff-ffff-4fff-8fff-ffffffffffff', + holder: 'test:replace', + acquiredAt: '2026-08-01 10:09:00', + }, + ], + }) + expect(replaced).toEqual({ + deletingAt: '2026-08-01 12:00:00', + created: false, + leaseCount: 1, + }) + expect(await meterB.listWriteLeases({})).toEqual({ + leases: [ + { + token: 'ffffffff-ffff-4fff-8fff-ffffffffffff', + holder: 'test:replace', + acquiredAt: '2026-08-01 10:09:00', + }, + ], + nextStartAfter: null, + truncated: false, + }) + expect(await meterB.exportCounters({})).toMatchObject({ + deletionShadow: { + deletingAt: '2026-08-01 12:00:00', + activeWriteLeaseCount: 1, + writeLeases: [{ acquiredAt: '2026-08-01 10:09:00' }], + }, + }) + await expect( + meterB.shadowReplaceDeletionState({ + deletingAt: '2026-08-01 15:00:00', + leases: [], + }), + ).resolves.toEqual({ + deletingAt: '2026-08-01 12:00:00', + created: false, + leaseCount: 0, + }) + expect(await meterB.countActiveWriteLeases()).toEqual({ count: 0 }) }, 30_000) diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index 79b844de68..fab87a0953 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -4,7 +4,7 @@ import { userMeterMirrorUpdatedAtToken, userMeterPackageServiceStateStaleMs, type DailyEntitlementResource, - type UserMeterDeletionShadow, + type UserMeterDeletionShadowExport, type UserMeterPackageServiceState, type UserMeterPackageServiceStatus, type UserMeterStorageBytesState, @@ -150,10 +150,14 @@ export function createInMemoryUserMeterEnv() { }) } - function readDeletionShadow(): UserMeterDeletionShadow { + function readDeletionShadowExport(): UserMeterDeletionShadowExport { + const writeLeases = listWriteLeasesSorted().map((lease) => ({ + acquiredAt: lease.acquiredAt, + })) return { deletingAt: deletion.deletingAt, - writeLeases: listWriteLeasesSorted(), + activeWriteLeaseCount: writeLeases.length, + writeLeases, } } @@ -452,6 +456,36 @@ export function createInMemoryUserMeterEnv() { deletion.deletingAt = input.deletingAt return { deletingAt: input.deletingAt, created: true } }, + async shadowReplaceDeletionState(input: { + deletingAt: string + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) { + const created = deletion.deletingAt == null + if (created) deletion.deletingAt = input.deletingAt + deletion.leases.clear() + const leasesByToken = new Map< + string, + { holder: string; acquiredAt: string } + >() + for (const lease of input.leases ?? []) { + leasesByToken.set(lease.token, { + holder: lease.holder, + acquiredAt: lease.acquiredAt, + }) + } + for (const [token, row] of leasesByToken) { + deletion.leases.set(token, row) + } + return { + deletingAt: deletion.deletingAt ?? input.deletingAt, + created, + leaseCount: deletion.leases.size, + } + }, async shadowAcquireWriteLease(input: { token: string holder: string @@ -593,7 +627,9 @@ export function createInMemoryUserMeterEnv() { const packageServiceStatesShadow = includeShadows ? listPackageServiceStatesSorted() : null - const deletionShadow = includeShadows ? readDeletionShadow() : null + const deletionShadow = includeShadows + ? readDeletionShadowExport() + : null return { counters, storageBytesShadow, From d20696d8ec429b054a256115f52b807b48362568 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 10:45:50 +0000 Subject: [PATCH 23/39] feat(account): move write leases into UserMeter Co-authored-by: Kent C. Dodds --- .../src/account/deletion-state.node.test.ts | 1354 ++++++++++++----- packages/worker/src/account/deletion-state.ts | 738 ++++++--- .../src/app/account-deletion.node.test.ts | 53 +- ...ccount-write-lease-middleware.node.test.ts | 25 +- .../worker/src/entitlements/user-meter-do.ts | 675 ++++++-- .../entitlements/user-meter.workers.test.ts | 167 +- .../admin/admin-account-write-lease-list.ts | 1 + .../worker/src/test-support/user-meter.ts | 251 ++- 8 files changed, 2412 insertions(+), 852 deletions(-) diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index 5838b76585..8ebf165fe2 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -1,11 +1,10 @@ import { DatabaseSync } from 'node:sqlite' -import { expect, test, vi } from 'vitest' +import { expect, test } from 'vitest' import { createD1FromSqlite } from '#worker/test-support/create-d1-from-sqlite.ts' import { userMeterRpc, type UserMeterEnv, } from '#worker/entitlements/user-meter-client.ts' -import { silenceExpectedConsoleWarns } from '#worker/test-support/console-spies.ts' import { createInMemoryUserMeterEnv, createWaitUntilDrain, @@ -306,7 +305,129 @@ test('waitUntil moves lease release off the response path', async () => { await releasePromise }) -test('optional env shadows mark/acquire/release/repair into UserMeter without changing D1 authority', async () => { +function createDeferred() { + let resolve: () => void = () => undefined + const promise = new Promise((next) => { + resolve = next + }) + return { promise, resolve: () => resolve() } +} + +function createGatedDoEnv(input: { + releaseGate?: Promise + finalizeGate?: Promise + finalizeFailOnceError?: string +}) { + const base = createInMemoryUserMeterEnv() + const namespace = base.env.USER_METER! + const gates = { releaseEntered: false, finalizeEntered: false } + let finalizeAttempts = 0 + const env = { + USER_METER: { + idFromName: namespace.idFromName.bind(namespace), + get(id: DurableObjectId) { + const stub = namespace.get(id) + return new Proxy(stub, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver) + if (prop === 'releaseWriteLease' && input.releaseGate) { + return async (args: { token: string }) => { + gates.releaseEntered = true + await input.releaseGate + return target.releaseWriteLease(args) + } + } + if (prop === 'finalizeWriteLeaseRepair') { + return async (args: { + token: string + repairId: string + expectedAcquiredAt: string + }) => { + if (input.finalizeGate) { + gates.finalizeEntered = true + await input.finalizeGate + } + if (input.finalizeFailOnceError) { + finalizeAttempts += 1 + if (finalizeAttempts === 1) { + throw new Error(input.finalizeFailOnceError) + } + } + return target.finalizeWriteLeaseRepair(args) + } + } + return typeof value === 'function' ? value.bind(target) : value + }, + }) + }, + }, + } as unknown as UserMeterEnv + return { + env, + gates, + get finalizeAttempts() { + return finalizeAttempts + }, + meterFor(userId: string) { + return userMeterRpc({ env, userId }) + }, + } +} + +async function waitFor(predicate: () => Promise, label: string) { + const deadline = Date.now() + 2_000 + while (Date.now() < deadline) { + if (await predicate()) return + await new Promise((resolve) => setTimeout(resolve, 5)) + } + throw new Error(`Timed out waiting for ${label}`) +} + +function holdDoWriteLease(input: { + db: D1Database + env: UserMeterEnv + holder: string + stableUserId?: string + waitUntil?: (promise: Promise) => void +}) { + const stableUserId = input.stableUserId ?? 'user-a' + const finish = createDeferred() + const started = createDeferred() + let token = '' + let acquiredAt = '' + const operation = withAccountWriteLease({ + db: input.db, + stableUserId, + holder: input.holder, + env: input.env, + waitUntil: input.waitUntil, + async write() { + const [lease] = await listActiveAccountWriteLeases( + input.db, + stableUserId, + input.env, + ) + token = lease!.token + acquiredAt = lease!.acquired_at + started.resolve() + await finish.promise + return 'lost' + }, + }) + return { + operation, + started: started.promise, + finish: finish.resolve, + get token() { + return token + }, + get acquiredAt() { + return acquiredAt + }, + } +} + +test('env makes UserMeter authoritative for acquire/held/release with D1 deleting_at gate', async () => { const { sqlite, db } = createLeaseTestDb() addLeaseRepairsTable(sqlite) const meter = createInMemoryUserMeterEnv() @@ -317,45 +438,33 @@ test('optional env shadows mark/acquire/release/repair into UserMeter without ch await withAccountWriteLease({ db, stableUserId: 'user-a', - holder: 'test:shadowed', + holder: 'test:do-authority', env: meter.env, waitUntil: drain.waitUntil, async write() { + expect(countLeaseRows(sqlite)).toBe(1) + expect( + await listActiveAccountWriteLeases(db, 'user-a', meter.env), + ).toHaveLength(1) expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) return 'ok' }, }) await drain.drain() - expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + expect( + await listActiveAccountWriteLeases(db, 'user-a', meter.env), + ).toHaveLength(0) + expect(countLeaseRows(sqlite)).toBe(0) expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) - let heldToken = '' - let heldAcquiredAt = '' - let startWrite: () => void = () => undefined - let finishWrite: () => void = () => undefined - const started = new Promise((resolve) => { - startWrite = resolve - }) - const finish = new Promise((resolve) => { - finishWrite = resolve - }) - const operation = withAccountWriteLease({ + const held = holdDoWriteLease({ db, - stableUserId: 'user-a', - holder: 'test:repair-shadow', env: meter.env, + holder: 'test:do-repair', waitUntil: drain.waitUntil, - async write() { - const [lease] = await listActiveAccountWriteLeases(db, 'user-a') - heldToken = lease!.token - heldAcquiredAt = lease!.acquired_at - startWrite() - await finish - return 'lost' - }, }) - await started - await drain.drain() + await held.started expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) await expect( @@ -364,10 +473,8 @@ test('optional env shadows mark/acquire/release/repair into UserMeter without ch dbUserId: 1, now: new Date('2099-01-01T00:00:00.000Z'), env: meter.env, - waitUntil: drain.waitUntil, }), ).resolves.toBe(1) - await drain.drain() expect(await meterA.readDeletionState()).toEqual({ deletingAt: '2099-01-01 00:00:00', }) @@ -375,17 +482,22 @@ test('optional env shadows mark/acquire/release/repair into UserMeter without ch await repairAccountWriteLease({ db, stableUserId: 'user-a', - token: heldToken, - expectedAcquiredAt: heldAcquiredAt, + token: held.token, + expectedAcquiredAt: held.acquiredAt, repairedByUserId: 'admin-user', reason: 'Inspected worker crash and confirmed process termination.', env: meter.env, - waitUntil: drain.waitUntil, }) - await drain.drain() - finishWrite() - await expect(operation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + expect( + sqlite.prepare(`SELECT reason FROM account_write_lease_repairs`).get(), + ).toEqual({ + reason: 'Inspected worker crash and confirmed process termination.', + }) await expect( withAccountWriteLease({ @@ -411,303 +523,280 @@ test('optional env shadows mark/acquire/release/repair into UserMeter without ch expect(await meterB.readDeletionState()).toEqual({ deletingAt: null }) }) -function createFailingShadowEnv(): UserMeterEnv { - return { - USER_METER: { - idFromName: (name: string) => ({ name, toString: () => name }), - get: () => ({ - async shadowMarkDeleting() { - throw new Error('shadow mark failed') - }, - async shadowReplaceDeletionState() { - throw new Error('shadow replace failed') - }, - async shadowAcquireWriteLease() { - throw new Error('shadow acquire failed') - }, - async shadowReleaseWriteLease() { - throw new Error('shadow release failed') - }, - }), - }, - } as unknown as UserMeterEnv -} - -test('UserMeter shadow failures never alter D1 lease or deletion results', async () => { - silenceExpectedConsoleWarns(['account-deletion-user-meter-shadow-failed']) - const { db } = createLeaseTestDb() - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const failingEnv = createFailingShadowEnv() - - let finishRelease: () => void = () => undefined - const releaseGate = new Promise((resolve) => { - finishRelease = resolve - }) - const originalBatch = db.batch.bind(db) - let releaseBatchSeen = false - db.batch = (async (statements: Parameters[0]) => { - const sql = statements - .map((statement) => { - const withQuery = statement as { query?: string } - return typeof withQuery.query === 'string' ? withQuery.query : '' - }) - .join('\n') - const isRelease = - sql.includes('MAX(active_write_count - 1') || - sql.includes('DELETE FROM account_write_leases') - if (isRelease && !releaseBatchSeen) { - releaseBatchSeen = true - await releaseGate - } - return await originalBatch(statements) - }) as D1Database['batch'] - - const pending: Array> = [] - await expect( - withAccountWriteLease({ - db, - stableUserId: 'user-a', - holder: 'test:shadow-fail', - env: failingEnv, - waitUntil(promise) { - pending.push(promise) - }, - async write() { - return 'mutated' - }, - }), - ).resolves.toBe('mutated') - expect(pending.length).toBeGreaterThan(0) - expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) - finishRelease() - await Promise.all(pending) - expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) - - await expect( - markAccountDeleting({ - db, - dbUserId: 1, - env: failingEnv, - }), - ).resolves.toBe(0) - expect(warn).toHaveBeenCalledWith( - 'account-deletion-user-meter-shadow-failed', - expect.any(Error), - ) - warn.mockRestore() -}) - -function createDeferred() { - let resolve: () => void = () => undefined - const promise = new Promise((next) => { - resolve = next - }) - return { promise, resolve: () => resolve() } -} - -function createGatedShadowEnv(input: { - releaseGate: Promise - markGate?: Promise -}) { - const base = createInMemoryUserMeterEnv() - const namespace = base.env.USER_METER! - const gates = { - markEntered: false, - releaseEntered: false, - } - const env = { - USER_METER: { - idFromName: namespace.idFromName.bind(namespace), - get(id: DurableObjectId) { - const stub = namespace.get(id) - return { - shadowMarkDeleting: (args: { deletingAt: string }) => - stub.shadowMarkDeleting(args), - shadowReplaceDeletionState: async (args: { - deletingAt: string - leases?: ReadonlyArray<{ - token: string - holder: string - acquiredAt: string - }> - }) => { - gates.markEntered = true - if (input.markGate) await input.markGate - return stub.shadowReplaceDeletionState(args) - }, - shadowAcquireWriteLease: (args: { - token: string - holder: string - acquiredAt: string - }) => stub.shadowAcquireWriteLease(args), - shadowReleaseWriteLease: async (args: { token: string }) => { - gates.releaseEntered = true - await input.releaseGate - return stub.shadowReleaseWriteLease(args) - }, - countActiveWriteLeases: () => stub.countActiveWriteLeases(), - readDeletionState: () => stub.readDeletionState(), - listWriteLeases: ( - args: { pageSize?: number; startAfter?: string | null } = {}, - ) => stub.listWriteLeases(args), - } - }, - }, - } as unknown as UserMeterEnv - return { - env, - gates, - meterFor(userId: string) { - return userMeterRpc({ env, userId }) - }, - } -} - -async function waitFor(predicate: () => Promise, label: string) { - const deadline = Date.now() + 2_000 - while (Date.now() < deadline) { - if (await predicate()) return - await new Promise((resolve) => setTimeout(resolve, 5)) - } - throw new Error(`Timed out waiting for ${label}`) -} - -test('without waitUntil, lease return awaits release shadow', async () => { - const { db } = createLeaseTestDb() - const release = createDeferred() - const gated = createGatedShadowEnv({ releaseGate: release.promise }) - const meterA = gated.meterFor('user-a') +test('email-style omitted env keeps exact D1 lease path', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) - let settled = false - const operation = withAccountWriteLease({ + await withAccountWriteLease({ db, stableUserId: 'user-a', - holder: 'test:await-release-shadow', - env: gated.env, + holder: 'test:email-d1', async write() { + expect(countLeaseRows(sqlite)).toBe(1) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) return 'ok' }, - }).then((value) => { - settled = true - return value }) - await waitFor( - async () => - gated.gates.releaseEntered && - (await listActiveAccountWriteLeases(db, 'user-a')).length === 0 && - (await meterA.countActiveWriteLeases()).count === 1, - 'D1 release with gated shadow lease still held', - ) - expect(settled).toBe(false) - release.resolve() - await expect(operation).resolves.toBe('ok') + expect(countLeaseRows(sqlite)).toBe(0) expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) }) -test('without waitUntil, mark and repair await their shadows', async () => { +test('mixed D1 + DO leases drain/count/dedupe and admin union list', async () => { const { sqlite, db } = createLeaseTestDb() addLeaseRepairsTable(sqlite) - const release = createDeferred() - const mark = createDeferred() - const gated = createGatedShadowEnv({ - releaseGate: release.promise, - markGate: mark.promise, - }) - const meterA = gated.meterFor('user-a') + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) const drain = createWaitUntilDrain() - let heldToken = '' - let heldAcquiredAt = '' - let startWrite: () => void = () => undefined - let finishWrite: () => void = () => undefined - const started = new Promise((resolve) => { - startWrite = resolve + let d1Token = '' + let d1AcquiredAt = '' + let startD1: () => void = () => undefined + let finishD1: () => void = () => undefined + const d1Started = new Promise((resolve) => { + startD1 = resolve }) - const finish = new Promise((resolve) => { - finishWrite = resolve + const d1Finish = new Promise((resolve) => { + finishD1 = resolve }) - const held = withAccountWriteLease({ + const d1Operation = withAccountWriteLease({ db, stableUserId: 'user-a', - holder: 'test:await-mark-repair', - env: gated.env, + holder: 'test:email-d1', waitUntil: drain.waitUntil, async write() { const [lease] = await listActiveAccountWriteLeases(db, 'user-a') - heldToken = lease!.token - heldAcquiredAt = lease!.acquired_at - startWrite() - await finish - return 'lost' + d1Token = lease!.token + d1AcquiredAt = lease!.acquired_at + startD1() + await d1Finish + return 'd1' }, }) - await started - await drain.drain() + await d1Started - let markSettled = false - const markPromise = markAccountDeleting({ - db, - dbUserId: 1, - now: new Date('2099-01-01T00:00:00.000Z'), - env: gated.env, - }).then((count) => { - markSettled = true - return count + let doToken = '' + let doAcquiredAt = '' + let startDo: () => void = () => undefined + let finishDo: () => void = () => undefined + const doStarted = new Promise((resolve) => { + startDo = resolve }) - await waitFor( - async () => gated.gates.markEntered && !markSettled, - 'mark shadow gate', - ) - expect(await meterA.readDeletionState()).toEqual({ deletingAt: null }) - mark.resolve() - await expect(markPromise).resolves.toBe(1) - expect(await meterA.readDeletionState()).toEqual({ - deletingAt: '2099-01-01 00:00:00', + const doFinish = new Promise((resolve) => { + finishDo = resolve }) - - let repairSettled = false - const repairPromise = repairAccountWriteLease({ + const doOperation = withAccountWriteLease({ db, stableUserId: 'user-a', - token: heldToken, - expectedAcquiredAt: heldAcquiredAt, - repairedByUserId: 'admin-user', - reason: 'Inspected worker crash and confirmed process termination.', - env: gated.env, - }).then((result) => { - repairSettled = true - return result + holder: 'test:do-authority', + env: meter.env, + waitUntil: drain.waitUntil, + async write() { + const [lease] = ( + await listActiveAccountWriteLeases(db, 'user-a', meter.env) + ).filter((row) => row.holder === 'test:do-authority') + doToken = lease!.token + doAcquiredAt = lease!.acquired_at + startDo() + await doFinish + return 'do' + }, }) - await waitFor( - async () => - gated.gates.releaseEntered && - (await meterA.countActiveWriteLeases()).count === 1 && - !repairSettled, - 'repair shadow gate', - ) - release.resolve() - await expect(repairPromise).resolves.toEqual( - expect.objectContaining({ repaired: true }), - ) - expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) - finishWrite() - await expect(held).rejects.toBeInstanceOf(AccountWriteLeaseLostError) - await drain.drain() -}) + await doStarted -test('with waitUntil, lease return completes before release shadow settles', async () => { - const { db } = createLeaseTestDb() - const release = createDeferred() - const gated = createGatedShadowEnv({ releaseGate: release.promise }) - const meterA = gated.meterFor('user-a') - const drain = createWaitUntilDrain() + const union = await listActiveAccountWriteLeases(db, 'user-a', meter.env) + expect(union).toHaveLength(2) + expect(union.map((lease) => lease.holder).sort()).toEqual([ + 'test:do-authority', + 'test:email-d1', + ]) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toEqual( + expect.arrayContaining([ + expect.objectContaining({ token: d1Token, holder: 'test:email-d1' }), + expect.objectContaining({ token: doToken, holder: 'test:do-authority' }), + ]), + ) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(2) await expect( - withAccountWriteLease({ + markAccountDeleting({ db, - stableUserId: 'user-a', - holder: 'test:detach-release-shadow', - env: gated.env, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: meter.env, + }), + ).resolves.toBe(2) + expect(await meterA.listWriteLeases({})).toEqual({ + leases: expect.arrayContaining([ + expect.objectContaining({ + token: d1Token, + holder: 'test:email-d1', + authority: 'legacy', + }), + expect.objectContaining({ + token: doToken, + holder: 'test:do-authority', + authority: 'do', + }), + ]), + nextStartAfter: null, + truncated: false, + }) + + await repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: d1Token, + expectedAcquiredAt: d1AcquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected email writer crash and confirmed process termination.', + env: meter.env, + }) + finishD1() + await expect(d1Operation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: meter.env, + }), + ).resolves.toBe(1) + + await repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: doToken, + expectedAcquiredAt: doAcquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected DO writer crash and confirmed process termination.', + env: meter.env, + }) + finishDo() + await expect(doOperation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) + await drain.drain() + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: meter.env, + }), + ).resolves.toBe(0) +}) + +test('nested/detached/waitUntil parity for DO-authority leases', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + + const nested = await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:outer-do', + env: meter.env, + async write() { + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + return await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:nested-do', + env: meter.env, + async write() { + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + const active = await listActiveAccountWriteLeases( + db, + 'user-a', + meter.env, + ) + expect(active).toHaveLength(1) + expect(active[0]).toEqual( + expect.objectContaining({ holder: 'test:outer-do' }), + ) + return 'nested' + }, + }) + }, + }) + expect(nested).toBe('nested') + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + + let detached: Promise = Promise.resolve() + let releaseOuter: () => void = () => undefined + const outerReleased = new Promise((resolve) => { + releaseOuter = resolve + }) + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:outer-do', + env: meter.env, + async write() { + detached = (async () => { + await outerReleased + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:detached-do', + env: meter.env, + async write() { + const active = await listActiveAccountWriteLeases( + db, + 'user-a', + meter.env, + ) + expect(active).toHaveLength(1) + expect(active[0]).toEqual( + expect.objectContaining({ holder: 'test:detached-do' }), + ) + }, + }) + })() + }, + }) + releaseOuter() + await detached + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + expect(countLeaseRows(sqlite)).toBe(0) + + const release = createDeferred() + const gated = createGatedDoEnv({ releaseGate: release.promise }) + const gatedMeter = gated.meterFor('user-a') + const drain = createWaitUntilDrain() + + let settled = false + const awaitedRelease = withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:await-do-release', + env: gated.env, + async write() { + return 'ok' + }, + }).then((value) => { + settled = true + return value + }) + await waitFor( + async () => + gated.gates.releaseEntered && + (await gatedMeter.countActiveWriteLeases()).count === 1 && + !settled, + 'DO release gate', + ) + release.resolve() + await expect(awaitedRelease).resolves.toBe('ok') + expect(await gatedMeter.countActiveWriteLeases()).toEqual({ count: 0 }) + + const release2 = createDeferred() + const gated2 = createGatedDoEnv({ releaseGate: release2.promise }) + const gatedMeter2 = gated2.meterFor('user-a') + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:detach-do-release', + env: gated2.env, waitUntil: drain.waitUntil, async write() { return 'ok' @@ -716,34 +805,258 @@ test('with waitUntil, lease return completes before release shadow settles', asy ).resolves.toBe('ok') await waitFor( async () => - gated.gates.releaseEntered && - (await listActiveAccountWriteLeases(db, 'user-a')).length === 0 && - (await meterA.countActiveWriteLeases()).count === 1, - 'detached release shadow still held', + gated2.gates.releaseEntered && + (await gatedMeter2.countActiveWriteLeases()).count === 1, + 'detached DO release still held', ) - - release.resolve() + release2.resolve() await drain.drain() + expect(await gatedMeter2.countActiveWriteLeases()).toEqual({ count: 0 }) +}) + +test('D1 deleting_at gate fails closed after purge tombstone', async () => { + const { db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + + await markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: meter.env, + }) + await meterA.purge() + expect(await meterA.readDeletionState()).toEqual({ + deletingAt: '2099-01-01 00:00:00', + }) expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + env: meter.env, + async write() { + return 'blocked' + }, + }), + ).rejects.toBeInstanceOf(AccountDeletionInProgressError) }) -test('markAccountDeleting replaces UserMeter lease shadows from authoritative D1', async () => { +test('DO repair prepare/audit/finalize is idempotent and lease-lost aware', async () => { const { sqlite, db } = createLeaseTestDb() addLeaseRepairsTable(sqlite) const meter = createInMemoryUserMeterEnv() const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) - const drain = createWaitUntilDrain() + const held = holdDoWriteLease({ + db, + env: meter.env, + holder: 'test:do-repair-protocol', + }) + await held.started + + const prepared = await meterA.prepareWriteLeaseRepair({ + token: held.token, + expectedAcquiredAt: held.acquiredAt, + }) + expect(prepared).toEqual( + expect.objectContaining({ + prepared: true, + token: held.token, + acquiredAt: held.acquiredAt, + }), + ) + const repairId = + prepared.prepared === true ? prepared.repairId : 'missing-repair-id' + const retried = await meterA.prepareWriteLeaseRepair({ + token: held.token, + expectedAcquiredAt: held.acquiredAt, + }) + expect(retried).toEqual(expect.objectContaining({ prepared: true, repairId })) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: true, + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + }), + ).resolves.toEqual({ repaired: true, repairId }) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: false, + }) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + expect( + sqlite + .prepare( + `SELECT id, lease_token, lease_acquired_at FROM account_write_lease_repairs`, + ) + .get(), + ).toEqual({ + id: repairId, + lease_token: held.token, + lease_acquired_at: held.acquiredAt, + }) + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + }), + ).resolves.toEqual({ repaired: true, repairId }) + expect( + sqlite + .prepare(`SELECT COUNT(*) AS count FROM account_write_lease_repairs`) + .get(), + ).toEqual({ count: 1 }) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) +}) + +test('USER_METER failures fail closed when env is supplied', async () => { + const { db } = createLeaseTestDb() + const failingEnv = { + USER_METER: { + idFromName: (name: string) => ({ name, toString: () => name }), + get: () => ({ + async acquireWriteLease() { + throw new Error('do acquire failed') + }, + async markDeleting() { + throw new Error('do mark failed') + }, + }), + }, + } as unknown as UserMeterEnv + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + env: failingEnv, + async write() { + return 'mutated' + }, + }), + ).rejects.toThrow('do acquire failed') + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: failingEnv, + }), + ).rejects.toThrow('do mark failed') + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + env: {}, + async write() { + return 'blocked' + }, + }), + ).rejects.toThrow('USER_METER Durable Object binding is not configured.') +}) + +test('markAccountDeleting clears stale legacy rows and preserves DO-authority leases', async () => { + const { db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) - // Stale shadow left behind by a failed release (no matching D1 lease). await meterA.shadowAcquireWriteLease({ token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', - holder: 'test:stale-shadow', + holder: 'test:stale-legacy', acquiredAt: '2026-01-01 00:00:00', }) - expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + await meterA.acquireWriteLease({ + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:do-held', + acquiredAt: '2026-01-01 00:01:00', + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 2 }) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + env: meter.env, + }), + ).resolves.toBe(1) + expect(await meterA.listWriteLeases({})).toEqual({ + leases: [ + { + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:do-held', + acquiredAt: '2026-01-01 00:01:00', + authority: 'do', + }, + ], + nextStartAfter: null, + truncated: false, + }) + + await meterA.shadowReplaceDeletionState({ + deletingAt: '2099-01-01 00:00:00', + leases: [ + { + token: 'cccccccc-cccc-4ccc-8ccc-cccccccccccc', + holder: 'test:post-drain-stale', + acquiredAt: '2099-01-01 00:01:00', + }, + ], + }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 2 }) + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + env: meter.env, + }), + ).resolves.toBe(1) + expect(await meterA.listWriteLeases({})).toEqual({ + leases: [ + { + token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', + holder: 'test:do-held', + acquiredAt: '2026-01-01 00:01:00', + authority: 'do', + }, + ], + nextStartAfter: null, + truncated: false, + }) +}) + +test('old Phase-A D1-only mark counts a new DO-authority lease via D1 mirror', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const drain = createWaitUntilDrain() - let heldToken = '' - let heldAcquiredAt = '' let startWrite: () => void = () => undefined let finishWrite: () => void = () => undefined const started = new Promise((resolve) => { @@ -755,92 +1068,401 @@ test('markAccountDeleting replaces UserMeter lease shadows from authoritative D1 const operation = withAccountWriteLease({ db, stableUserId: 'user-a', - holder: 'test:active-d1', + holder: 'test:do-mirror-for-phase-a', env: meter.env, waitUntil: drain.waitUntil, async write() { - const [lease] = await listActiveAccountWriteLeases(db, 'user-a') - heldToken = lease!.token - heldAcquiredAt = lease!.acquired_at startWrite() await finish - return 'lost' + return 'held' }, }) await started - await drain.drain() - expect(await meterA.countActiveWriteLeases()).toEqual({ count: 2 }) + expect(countLeaseRows(sqlite)).toBe(1) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) await expect( markAccountDeleting({ db, dbUserId: 1, now: new Date('2099-01-01T00:00:00.000Z'), - env: meter.env, - waitUntil: drain.waitUntil, }), ).resolves.toBe(1) + + finishWrite() + await expect(operation).resolves.toBe('held') await drain.drain() - expect(await meterA.readDeletionState()).toEqual({ - deletingAt: '2099-01-01 00:00:00', - }) - // Stale shadow removed; active D1 lease preserved under the tombstone. - expect(await meterA.listWriteLeases({})).toEqual({ - leases: [ - { - token: heldToken, - holder: 'test:active-d1', - acquiredAt: heldAcquiredAt, + expect(countLeaseRows(sqlite)).toBe(0) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) +}) + +test('DO acquire fails closed and releases DO when D1 mirror response is lost unconfirmed', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const originalBatch = db.batch.bind(db) + db.batch = (async () => { + throw new Error('simulated lost D1 batch response') + }) as D1Database['batch'] + const originalPrepare = db.prepare.bind(db) + db.prepare = ((query: string) => { + const statement = originalPrepare(query) + if ( + query.includes('INSERT INTO account_write_leases') || + query.includes('active_write_count + 1') + ) { + return { + bind(...params: Array) { + const bound = statement.bind(...params) + return { + async run() { + return { meta: { changes: 0 } } + }, + first: bound.first.bind(bound), + all: bound.all.bind(bound), + } + }, + } + } + return statement + }) as D1Database['prepare'] + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:mirror-fail', + env: meter.env, + async write() { + return 'should-not-run' }, - ], - nextStartAfter: null, - truncated: false, + }), + ).rejects.toBeInstanceOf(AccountDeletionInProgressError) + expect(countLeaseRows(sqlite)).toBe(0) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + db.batch = originalBatch + db.prepare = originalPrepare +}) + +test('DO repair clears D1 mirror, stays audit-first, and retries after finalize without double audit/count', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const held = holdDoWriteLease({ + db, + env: meter.env, + holder: 'test:do-repair-mirror', }) - expect(await meterA.exportCounters({})).toMatchObject({ - deletionShadow: { - deletingAt: '2099-01-01 00:00:00', - activeWriteLeaseCount: 1, - writeLeases: [{ acquiredAt: heldAcquiredAt }], - }, + await held.started + expect(countLeaseRows(sqlite)).toBe(1) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 1 }) + + const prepared = await meterA.prepareWriteLeaseRepair({ + token: held.token, + expectedAcquiredAt: held.acquiredAt, }) + expect(prepared.prepared).toBe(true) + const repairId = + prepared.prepared === true ? prepared.repairId : 'missing-repair-id' - await repairAccountWriteLease({ + sqlite + .prepare( + `INSERT INTO account_write_lease_repairs ( + id, target_user_id, lease_token, lease_holder, + lease_acquired_at, repaired_by_user_id, reason, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .run( + repairId, + 'user-a', + held.token, + 'test:do-repair-mirror', + held.acquiredAt, + 'admin-user', + 'Inspected worker crash and confirmed process termination.', + '2099-01-01 00:00:00', + ) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: true, + }) + expect(countLeaseRows(sqlite)).toBe(1) + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + }), + ).resolves.toEqual({ repaired: true, repairId }) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: false, + }) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + expect( + sqlite + .prepare(`SELECT COUNT(*) AS count FROM account_write_lease_repairs`) + .get(), + ).toEqual({ count: 1 }) + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: meter.env, + }), + ).resolves.toEqual({ repaired: true, repairId }) + expect( + sqlite + .prepare(`SELECT COUNT(*) AS count FROM account_write_lease_repairs`) + .get(), + ).toEqual({ count: 1 }) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) +}) + +test('authoritative release cleans D1 mirror after DO release under waitUntil', async () => { + const { sqlite, db } = createLeaseTestDb() + const release = createDeferred() + const gated = createGatedDoEnv({ releaseGate: release.promise }) + const gatedMeter = gated.meterFor('user-a') + const drain = createWaitUntilDrain() + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:release-mirror-order', + env: gated.env, + waitUntil: drain.waitUntil, + async write() { + expect(countLeaseRows(sqlite)).toBe(1) + return 'ok' + }, + }), + ).resolves.toBe('ok') + await waitFor( + async () => + gated.gates.releaseEntered && + (await gatedMeter.countActiveWriteLeases()).count === 1 && + countLeaseRows(sqlite) === 1, + 'DO release gated while D1 mirror still held', + ) + + release.resolve() + await drain.drain() + expect(await gatedMeter.countActiveWriteLeases()).toEqual({ count: 0 }) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) +}) + +test('repair keeps D1 mirror until finalize; Phase-A mark counts during gate then drains', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const finalize = createDeferred() + const gated = createGatedDoEnv({ finalizeGate: finalize.promise }) + const gatedMeter = gated.meterFor('user-a') + const held = holdDoWriteLease({ + db, + env: gated.env, + holder: 'test:repair-finalize-gate', + }) + await held.started + + let repairSettled = false + const repairPromise = repairAccountWriteLease({ db, stableUserId: 'user-a', - token: heldToken, - expectedAcquiredAt: heldAcquiredAt, + token: held.token, + expectedAcquiredAt: held.acquiredAt, repairedByUserId: 'admin-user', reason: 'Inspected worker crash and confirmed process termination.', + env: gated.env, + }).then((result) => { + repairSettled = true + return result + }) + await waitFor( + async () => + gated.gates.finalizeEntered && + (await gatedMeter.assertWriteLeaseHeld({ token: held.token })).held && + countLeaseRows(sqlite) === 1 && + !repairSettled, + 'finalize gated with DO + D1 mirror still held', + ) + + await expect( + markAccountDeleting({ + db, + dbUserId: 1, + now: new Date('2099-01-01T00:00:00.000Z'), + }), + ).resolves.toBe(1) + + finalize.resolve() + await expect(repairPromise).resolves.toEqual( + expect.objectContaining({ repaired: true }), + ) + expect(await gatedMeter.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: false, + }) + expect(countLeaseRows(sqlite)).toBe(0) + await expect(markAccountDeleting({ db, dbUserId: 1 })).resolves.toBe(0) + + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) +}) + +test('lost finalize response retries clear stale D1 mirror and return stable repairId', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const held = holdDoWriteLease({ + db, env: meter.env, - waitUntil: drain.waitUntil, + holder: 'test:stale-mirror-after-finalize', }) - await drain.drain() - finishWrite() - await expect(operation).rejects.toBeInstanceOf(AccountWriteLeaseLostError) - expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + await held.started - // Failed release after D1 drain can leave a stale shadow; retry clears it. - await meterA.shadowReplaceDeletionState({ - deletingAt: '2099-01-01 00:00:00', - leases: [ - { - token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', - holder: 'test:post-drain-stale', - acquiredAt: '2099-01-01 00:01:00', - }, - ], + const prepared = await meterA.prepareWriteLeaseRepair({ + token: held.token, + expectedAcquiredAt: held.acquiredAt, }) - expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) + const repairId = + prepared.prepared === true ? prepared.repairId : 'missing-repair-id' + sqlite + .prepare( + `INSERT INTO account_write_lease_repairs ( + id, target_user_id, lease_token, lease_holder, + lease_acquired_at, repaired_by_user_id, reason, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .run( + repairId, + 'user-a', + held.token, + 'test:stale-mirror-after-finalize', + held.acquiredAt, + 'admin-user', + 'Inspected worker crash and confirmed process termination.', + '2099-01-01 00:00:00', + ) + await meterA.finalizeWriteLeaseRepair({ + token: held.token, + repairId, + expectedAcquiredAt: held.acquiredAt, + }) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: false, + }) + expect(countLeaseRows(sqlite)).toBe(1) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 1 }) await expect( - markAccountDeleting({ + repairAccountWriteLease({ db, - dbUserId: 1, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', env: meter.env, }), - ).resolves.toBe(0) - expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) - expect(await meterA.readDeletionState()).toEqual({ - deletingAt: '2099-01-01 00:00:00', + ).resolves.toEqual({ repaired: true, repairId }) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + expect( + sqlite + .prepare(`SELECT COUNT(*) AS count FROM account_write_lease_repairs`) + .get(), + ).toEqual({ count: 1 }) + + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) +}) + +test('finalize failure without commit leaves D1 mirror held for retry', async () => { + const { sqlite, db } = createLeaseTestDb() + addLeaseRepairsTable(sqlite) + const gated = createGatedDoEnv({ + finalizeFailOnceError: 'simulated finalize transport failure', + }) + const meterA = gated.meterFor('user-a') + const held = holdDoWriteLease({ + db, + env: gated.env, + holder: 'test:finalize-fail-closed', + }) + await held.started + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: gated.env, + }), + ).rejects.toThrow('simulated finalize transport failure') + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: true, }) + expect(countLeaseRows(sqlite)).toBe(1) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 1 }) + + await expect( + repairAccountWriteLease({ + db, + stableUserId: 'user-a', + token: held.token, + expectedAcquiredAt: held.acquiredAt, + repairedByUserId: 'admin-user', + reason: 'Inspected worker crash and confirmed process termination.', + env: gated.env, + }), + ).resolves.toEqual(expect.objectContaining({ repaired: true })) + expect(await meterA.assertWriteLeaseHeld({ token: held.token })).toEqual({ + held: false, + }) + expect(countLeaseRows(sqlite)).toBe(0) + expect(gated.finalizeAttempts).toBe(2) + + held.finish() + await expect(held.operation).rejects.toBeInstanceOf( + AccountWriteLeaseLostError, + ) }) diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index 00ba34f431..ea6b7cddfa 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -39,30 +39,349 @@ export type AccountWriteLease = { acquiredAt: string } -const accountDeletionShadowFailedLog = - 'account-deletion-user-meter-shadow-failed' +type ListedAccountWriteLease = { + token: string + holder: string + acquired_at: string +} -function scheduleAccountDeletionShadow(input: { - env?: UserMeterEnv - waitUntil?: (promise: Promise) => void - task: (env: UserMeterEnv) => Promise -}): Promise { - const env = input.env - if (!env || !userMeterNamespace(env)) return Promise.resolve() - const tracked = input.task(env).catch((error: unknown) => { - console.warn(accountDeletionShadowFailedLog, error) +function requireUserMeterEnv(env: UserMeterEnv) { + if (!userMeterNamespace(env)) { + throw new Error('USER_METER Durable Object binding is not configured.') + } + return env +} + +async function listD1ActiveAccountWriteLeases( + db: D1Database, + stableUserId: string, +): Promise> { + const rows = await db + .prepare( + `SELECT token, holder, acquired_at + FROM account_write_leases + WHERE user_id = ? AND released_at IS NULL + ORDER BY acquired_at, token`, + ) + .bind(stableUserId) + .all() + return rows.results ?? [] +} + +async function listAllDoAuthorityWriteLeases( + env: UserMeterEnv, + stableUserId: string, +): Promise> { + const meter = userMeterRpc({ env, userId: stableUserId }) + const leases: Array = [] + let startAfter: string | null = null + for (;;) { + const page = await meter.listDoAuthorityWriteLeases({ + pageSize: 500, + startAfter, + }) + for (const lease of page.leases) { + leases.push({ + token: lease.token, + holder: lease.holder, + acquired_at: lease.acquiredAt, + }) + } + if (!page.nextStartAfter) break + startAfter = page.nextStartAfter + } + return leases +} + +function unionActiveWriteLeases( + d1Leases: ReadonlyArray, + doLeases: ReadonlyArray, +): Array { + const byToken = new Map() + for (const lease of d1Leases) { + byToken.set(lease.token, lease) + } + for (const lease of doLeases) { + if (!byToken.has(lease.token)) byToken.set(lease.token, lease) + } + return [...byToken.values()].sort((left, right) => { + const byAcquired = left.acquired_at.localeCompare(right.acquired_at) + if (byAcquired !== 0) return byAcquired + return left.token.localeCompare(right.token) }) - if (input.waitUntil) input.waitUntil(tracked) - return tracked +} + +/** + * Temporary rollout mirror: insert the exact DO lease into D1 so old Phase-A + * isolates that only read D1 still observe every DO-authority acquire. + * Uses the same atomic batch + lost-response reconcile as the D1 lease path. + */ +async function acquireD1WriteLeaseMirror(input: { + db: D1Database + lease: AccountWriteLease +}): Promise { + const prepareLeaseInsert = () => + input.db + .prepare( + `INSERT INTO account_write_leases ( + token, user_id, holder, acquired_at + ) + SELECT ?, stable_user_id, ?, ? + FROM users + WHERE stable_user_id = ? AND deleting_at IS NULL`, + ) + .bind( + input.lease.token, + input.lease.holder, + input.lease.acquiredAt, + input.lease.stableUserId, + ) + const prepareWriteCountIncrement = () => + input.db + .prepare( + `UPDATE users + SET active_write_count = active_write_count + 1 + WHERE stable_user_id = ? + AND EXISTS ( + SELECT 1 FROM account_write_leases + WHERE token = ? AND user_id = ? + AND released_at IS NULL + )`, + ) + .bind( + input.lease.stableUserId, + input.lease.token, + input.lease.stableUserId, + ) + let inserted: D1Result + let counted: D1Result + try { + const results = await input.db.batch([ + prepareLeaseInsert(), + prepareWriteCountIncrement(), + ]) + if (!results[0] || !results[1]) { + throw new Error('Account write lease batch returned incomplete results.') + } + inserted = results[0] + counted = results[1] + } catch { + const held = await input.db + .prepare( + `SELECT 1 AS held FROM account_write_leases + WHERE token = ? AND user_id = ? AND released_at IS NULL`, + ) + .bind(input.lease.token, input.lease.stableUserId) + .first<{ held: number }>() + if (held?.held === 1) { + inserted = { meta: { changes: 1 } } as D1Result + counted = { meta: { changes: 1 } } as D1Result + } else { + inserted = await prepareLeaseInsert().run() + counted = await prepareWriteCountIncrement().run() + } + } + return ( + (inserted?.meta.changes ?? 0) === 1 && (counted?.meta.changes ?? 0) === 1 + ) +} + +/** + * Release the temporary D1 rollout mirror. Idempotent when the row is already + * gone (no count decrement). Prefer calling after DO release so a stale D1 + * mirror can overblock old Phase-A marks but never underblock. + */ +async function releaseD1WriteLeaseMirror(input: { + db: D1Database + stableUserId: string + token: string + expectedAcquiredAt?: string +}): Promise<{ released: boolean }> { + const prepareWriteCountDecrement = () => + input.expectedAcquiredAt == null + ? input.db + .prepare( + `UPDATE users + SET active_write_count = MAX(active_write_count - 1, 0) + WHERE stable_user_id = ? + AND EXISTS ( + SELECT 1 FROM account_write_leases + WHERE token = ? AND user_id = ? + AND released_at IS NULL + )`, + ) + .bind(input.stableUserId, input.token, input.stableUserId) + : input.db + .prepare( + `UPDATE users + SET active_write_count = MAX(active_write_count - 1, 0) + WHERE stable_user_id = ? + AND EXISTS ( + SELECT 1 FROM account_write_leases + WHERE token = ? AND user_id = ? + AND acquired_at = ? AND released_at IS NULL + )`, + ) + .bind( + input.stableUserId, + input.token, + input.stableUserId, + input.expectedAcquiredAt, + ) + const prepareLeaseDelete = () => + input.expectedAcquiredAt == null + ? input.db + .prepare( + `DELETE FROM account_write_leases + WHERE token = ? AND user_id = ? AND released_at IS NULL`, + ) + .bind(input.token, input.stableUserId) + : input.db + .prepare( + `DELETE FROM account_write_leases + WHERE token = ? AND user_id = ? AND acquired_at = ? + AND released_at IS NULL`, + ) + .bind(input.token, input.stableUserId, input.expectedAcquiredAt) + try { + const results = await input.db.batch([ + prepareWriteCountDecrement(), + prepareLeaseDelete(), + ]) + return { released: (results[1]?.meta.changes ?? 0) === 1 } + } catch { + const held = await input.db + .prepare( + input.expectedAcquiredAt == null + ? `SELECT 1 AS held FROM account_write_leases + WHERE token = ? AND user_id = ? AND released_at IS NULL` + : `SELECT 1 AS held FROM account_write_leases + WHERE token = ? AND user_id = ? AND acquired_at = ? + AND released_at IS NULL`, + ) + .bind( + ...(input.expectedAcquiredAt == null + ? [input.token, input.stableUserId] + : [input.token, input.stableUserId, input.expectedAcquiredAt]), + ) + .first<{ held: number }>() + if (held?.held !== 1) return { released: false } + const deleted = await prepareLeaseDelete().run() + if ((deleted.meta.changes ?? 0) === 1) { + await input.db + .prepare( + `UPDATE users + SET active_write_count = MAX(active_write_count - 1, 0) + WHERE stable_user_id = ?`, + ) + .bind(input.stableUserId) + .run() + return { released: true } + } + return { released: false } + } +} + +async function insertOrVerifyDoRepairAudit(input: { + db: D1Database + repairId: string + stableUserId: string + token: string + holder: string + acquiredAt: string + repairedByUserId: string + reason: string + now: string +}) { + const inserted = await input.db + .prepare( + `INSERT OR IGNORE INTO account_write_lease_repairs ( + id, target_user_id, lease_token, lease_holder, + lease_acquired_at, repaired_by_user_id, reason, created_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + input.repairId, + input.stableUserId, + input.token, + input.holder, + input.acquiredAt, + input.repairedByUserId, + input.reason, + input.now, + ) + .run() + if ((inserted.meta.changes ?? 0) === 1) return + const existing = await input.db + .prepare( + `SELECT target_user_id, lease_token, lease_holder, lease_acquired_at, + repaired_by_user_id, reason + FROM account_write_lease_repairs + WHERE id = ?`, + ) + .bind(input.repairId) + .first<{ + target_user_id: string + lease_token: string + lease_holder: string + lease_acquired_at: string + repaired_by_user_id: string + reason: string + }>() + if ( + !existing || + existing.target_user_id !== input.stableUserId || + existing.lease_token !== input.token || + existing.lease_holder !== input.holder || + existing.lease_acquired_at !== input.acquiredAt || + existing.repaired_by_user_id !== input.repairedByUserId || + existing.reason !== input.reason + ) { + throw new Error('Active account write lease did not match repair request.') + } +} + +async function findMatchingRepairAudit(input: { + db: D1Database + stableUserId: string + token: string + expectedAcquiredAt: string + repairedByUserId: string + reason: string +}) { + return await input.db + .prepare( + `SELECT id + FROM account_write_lease_repairs + WHERE target_user_id = ? + AND lease_token = ? + AND lease_acquired_at = ? + AND repaired_by_user_id = ? + AND reason = ? + ORDER BY created_at DESC, id DESC + LIMIT 1`, + ) + .bind( + input.stableUserId, + input.token, + input.expectedAcquiredAt, + input.repairedByUserId, + input.reason, + ) + .first<{ id: string }>() } export async function markAccountDeleting(input: { db: D1Database dbUserId: number now?: Date - /** Optional expand-phase UserMeter shadow after a successful D1 mark. */ + /** + * When supplied, UserMeter is authoritative for the deletion drain count + * (fail closed). When omitted, retain the exact D1 lease-count path. + */ env?: UserMeterEnv - /** Prefer `ctx.waitUntil`; awaited (non-rejecting) when omitted. */ + /** Kept for API compatibility; DO mark is awaited. */ waitUntil?: (promise: Promise) => void }) { const now = utcSqliteTimestamp(input.now ?? new Date()) @@ -87,39 +406,28 @@ export async function markAccountDeleting(input: { .first<{ stable_user_id: string; deleting_at: string | null }>() const stableUserId = userRow?.stable_user_id const deletingAt = userRow?.deleting_at - // Authoritative active D1 leases after deleting_at is set (new acquires - // are fenced). Returned count stays D1-authoritative for drain waits. - const activeLeases = stableUserId - ? await listActiveAccountWriteLeases(input.db, stableUserId) + const activeD1Leases = stableUserId + ? await listD1ActiveAccountWriteLeases(input.db, stableUserId) : [] - if ( - input.env && - userMeterNamespace(input.env) && - stableUserId && - deletingAt - ) { - const markShadowPromise = scheduleAccountDeletionShadow({ - env: input.env, - waitUntil: input.waitUntil, - task: async (env) => { - // Replace (not append) so a prior failed shadow release cannot - // leave stale rows after D1 drain; empty list clears shadows. - await userMeterRpc({ - env, - userId: stableUserId, - }).shadowReplaceDeletionState({ - deletingAt, - leases: activeLeases.map((lease) => ({ - token: lease.token, - holder: lease.holder, - acquiredAt: lease.acquired_at, - })), - }) - }, - }) - if (!input.waitUntil) await markShadowPromise + if (input.env === undefined) { + return activeD1Leases.length + } + const env = requireUserMeterEnv(input.env) + if (!stableUserId || !deletingAt) { + throw new Error('Account could not be marked for deletion.') } - return activeLeases.length + const marked = await userMeterRpc({ + env, + userId: stableUserId, + }).markDeleting({ + deletingAt, + legacyLeases: activeD1Leases.map((lease) => ({ + token: lease.token, + holder: lease.holder, + acquiredAt: lease.acquired_at, + })), + }) + return marked.leaseCount } export async function assertAccountWritableDb( @@ -143,12 +451,12 @@ export async function assertAccountWritable(env: Env, stableUserId: string) { * Live lease frames keyed by stable user id, propagated down the current * async call chain. Nested {@link withAccountWriteLease} calls for the same * user reuse the outer lease while its frame is still active instead of - * paying another acquire/release round trip to D1 (~5 statements each): the - * outer lease spans the nested write, so deletion stays blocked for exactly - * as long as it does today. MCP requests, app requests, job runs, and - * package invocations all take a lease at their boundary, so before this - * reuse a single execute call that invoked one package export paid for two - * full leases. + * paying another acquire/release round trip (~5 statements each on D1, or a + * DO RPC pair when `env` is supplied): the outer lease spans the nested + * write, so deletion stays blocked for exactly as long as it does today. MCP + * requests, app requests, job runs, and package invocations all take a lease + * at their boundary, so before this reuse a single execute call that invoked + * one package export paid for two full leases. * * Frames deactivate when the outer lease releases. Detached work (for * example `waitUntil` callbacks spawned inside `write`) inherits this @@ -165,7 +473,12 @@ export async function withAccountWriteLease(input: { stableUserId: string holder?: string waitUntil?: (promise: Promise) => void - /** Optional expand-phase UserMeter shadow after D1 acquire/release. */ + /** + * When supplied, UserMeter is authoritative for acquire/held/release + * (fail closed) after a D1 `deleting_at` point gate, with a temporary D1 + * rollout mirror of the same lease token. When omitted, retain the exact + * D1 lease path (email transition). + */ env?: UserMeterEnv write: () => Promise }) { @@ -178,19 +491,85 @@ export async function withAccountWriteLease(input: { nextHeldLeases.set(input.stableUserId, frame) try { return await heldAccountWriteLeaseStorage.run(nextHeldLeases, async () => - acquireAccountWriteLeaseAndWrite({ ...input, frame }), + input.env === undefined + ? acquireD1AccountWriteLeaseAndWrite({ ...input, frame }) + : acquireDoAccountWriteLeaseAndWrite({ + ...input, + env: requireUserMeterEnv(input.env), + frame, + }), ) } finally { frame.active = false } } -async function acquireAccountWriteLeaseAndWrite(input: { +async function acquireDoAccountWriteLeaseAndWrite(input: { + db: D1Database + stableUserId: string + holder?: string + waitUntil?: (promise: Promise) => void + env: UserMeterEnv + frame: AccountWriteLeaseFrame + write: () => Promise +}) { + await assertAccountWritableDb(input.db, input.stableUserId) + const lease: AccountWriteLease = { + token: crypto.randomUUID(), + stableUserId: input.stableUserId, + holder: input.holder ?? 'unspecified', + acquiredAt: utcSqliteTimestamp(), + } + const meter = userMeterRpc({ + env: input.env, + userId: lease.stableUserId, + }) + const acquired = await meter.acquireWriteLease({ + token: lease.token, + holder: lease.holder, + acquiredAt: lease.acquiredAt, + }) + if (!acquired.acquired) { + throw new AccountDeletionInProgressError() + } + try { + const mirrored = await acquireD1WriteLeaseMirror({ + db: input.db, + lease, + }) + if (!mirrored) { + throw new AccountDeletionInProgressError() + } + } catch (error) { + await meter.releaseWriteLease({ token: lease.token }).catch(() => {}) + throw error + } + try { + const result = await input.write() + const held = await meter.assertWriteLeaseHeld({ token: lease.token }) + if (!held.held) throw new AccountWriteLeaseLostError() + return result + } finally { + input.frame.active = false + // DO release before D1 mirror clear (stale D1 may overblock; never underblock). + const releasePromise = (async () => { + await meter.releaseWriteLease({ token: lease.token }) + await releaseD1WriteLeaseMirror({ + db: input.db, + stableUserId: lease.stableUserId, + token: lease.token, + }) + })() + if (input.waitUntil) input.waitUntil(releasePromise) + else await releasePromise + } +} + +async function acquireD1AccountWriteLeaseAndWrite(input: { db: D1Database stableUserId: string holder?: string waitUntil?: (promise: Promise) => void - env?: UserMeterEnv frame: AccountWriteLeaseFrame write: () => Promise }) { @@ -228,82 +607,10 @@ async function acquireAccountWriteLeaseAndWrite(input: { holder: input.holder ?? 'unspecified', acquiredAt: utcSqliteTimestamp(), } - const prepareLeaseInsert = () => - input.db - .prepare( - `INSERT INTO account_write_leases ( - token, user_id, holder, acquired_at - ) - SELECT ?, stable_user_id, ?, ? - FROM users - WHERE stable_user_id = ? AND deleting_at IS NULL`, - ) - .bind(lease.token, lease.holder, lease.acquiredAt, lease.stableUserId) - const prepareWriteCountIncrement = () => - input.db - .prepare( - `UPDATE users - SET active_write_count = active_write_count + 1 - WHERE stable_user_id = ? - AND EXISTS ( - SELECT 1 FROM account_write_leases - WHERE token = ? AND user_id = ? - AND released_at IS NULL - )`, - ) - .bind(lease.stableUserId, lease.token, lease.stableUserId) - let inserted: D1Result - let counted: D1Result - try { - const results = await input.db.batch([ - prepareLeaseInsert(), - prepareWriteCountIncrement(), - ]) - if (!results[0] || !results[1]) { - throw new Error('Account write lease batch returned incomplete results.') - } - inserted = results[0] - counted = results[1] - } catch { - // A D1 batch can commit and still lose its response. Reconcile by token; - // if it did not commit, retain the old sequential fallback so test and - // local adapters that reject batch calls still preserve deletion fencing. - const held = await input.db - .prepare( - `SELECT 1 AS held FROM account_write_leases - WHERE token = ? AND user_id = ? AND released_at IS NULL`, - ) - .bind(lease.token, lease.stableUserId) - .first<{ held: number }>() - if (held?.held === 1) { - inserted = { meta: { changes: 1 } } as D1Result - counted = { meta: { changes: 1 } } as D1Result - } else { - inserted = await prepareLeaseInsert().run() - counted = await prepareWriteCountIncrement().run() - } - } - if ( - (inserted?.meta.changes ?? 0) !== 1 || - (counted?.meta.changes ?? 0) !== 1 - ) { + const mirrored = await acquireD1WriteLeaseMirror({ db: input.db, lease }) + if (!mirrored) { throw new AccountDeletionInProgressError() } - // Hold acquire so a detached release shadow cannot land first. - const acquireShadowPromise = scheduleAccountDeletionShadow({ - env: input.env, - waitUntil: input.waitUntil, - task: async (env) => { - await userMeterRpc({ - env, - userId: lease.stableUserId, - }).shadowAcquireWriteLease({ - token: lease.token, - holder: lease.holder, - acquiredAt: lease.acquiredAt, - }) - }, - }) try { const result = await input.write() const held = await input.db @@ -317,88 +624,29 @@ async function acquireAccountWriteLeaseAndWrite(input: { return result } finally { input.frame.active = false - const release = async () => { - const prepareWriteCountDecrement = () => - input.db - .prepare( - `UPDATE users - SET active_write_count = MAX(active_write_count - 1, 0) - WHERE stable_user_id = ? - AND EXISTS ( - SELECT 1 FROM account_write_leases - WHERE token = ? AND user_id = ? - AND released_at IS NULL - )`, - ) - .bind(lease.stableUserId, lease.token, lease.stableUserId) - const prepareLeaseDelete = () => - input.db - .prepare( - `DELETE FROM account_write_leases - WHERE token = ? AND user_id = ? AND released_at IS NULL`, - ) - .bind(lease.token, lease.stableUserId) - try { - await input.db.batch([ - prepareWriteCountDecrement(), - prepareLeaseDelete(), - ]) - } catch { - const held = await input.db - .prepare( - `SELECT 1 AS held FROM account_write_leases - WHERE token = ? AND user_id = ? AND released_at IS NULL`, - ) - .bind(lease.token, lease.stableUserId) - .first<{ held: number }>() - if (held?.held !== 1) return - const deleted = await prepareLeaseDelete().run() - if ((deleted.meta.changes ?? 0) === 1) { - await input.db - .prepare( - `UPDATE users - SET active_write_count = MAX(active_write_count - 1, 0) - WHERE stable_user_id = ?`, - ) - .bind(lease.stableUserId) - .run() - } - } - } - const releasePromise = release() + const releasePromise = releaseD1WriteLeaseMirror({ + db: input.db, + stableUserId: lease.stableUserId, + token: lease.token, + }) if (input.waitUntil) input.waitUntil(releasePromise) else await releasePromise - // Ordered after D1 release; await when waitUntil is omitted. - const releaseShadowPromise = scheduleAccountDeletionShadow({ - env: input.env, - waitUntil: input.waitUntil, - task: async (env) => { - await releasePromise - await acquireShadowPromise - await userMeterRpc({ - env, - userId: lease.stableUserId, - }).shadowReleaseWriteLease({ token: lease.token }) - }, - }) - if (!input.waitUntil) await releaseShadowPromise } } export async function listActiveAccountWriteLeases( db: D1Database, stableUserId: string, + env?: UserMeterEnv, ) { - const rows = await db - .prepare( - `SELECT token, holder, acquired_at - FROM account_write_leases - WHERE user_id = ? AND released_at IS NULL - ORDER BY acquired_at, token`, - ) - .bind(stableUserId) - .all<{ token: string; holder: string; acquired_at: string }>() - return rows.results ?? [] + const d1Leases = await listD1ActiveAccountWriteLeases(db, stableUserId) + if (env === undefined) return d1Leases + const requiredEnv = requireUserMeterEnv(env) + const doLeases = await listAllDoAuthorityWriteLeases( + requiredEnv, + stableUserId, + ) + return unionActiveWriteLeases(d1Leases, doLeases) } export async function repairAccountWriteLease(input: { @@ -408,14 +656,81 @@ export async function repairAccountWriteLease(input: { expectedAcquiredAt: string repairedByUserId: string reason: string - /** Optional expand-phase UserMeter shadow after a successful D1 repair. */ + /** + * When supplied, DO-authority leases use prepare → audit → finalize DO → + * clear D1 mirror (fail closed). Never clear the D1 mirror while the DO + * lease remains held. D1-only leases keep the existing atomic batch. When + * omitted, retain the exact D1 repair path. + */ env?: UserMeterEnv - /** Prefer `ctx.waitUntil`; awaited (non-rejecting) when omitted. */ + /** Kept for API compatibility; DO finalize is awaited. */ waitUntil?: (promise: Promise) => void }) { if (input.reason.trim().length < 10) { throw new Error('Lease repair requires a detailed audit reason.') } + const reason = input.reason.trim() + if (input.env !== undefined) { + const env = requireUserMeterEnv(input.env) + const meter = userMeterRpc({ env, userId: input.stableUserId }) + const prepared = await meter.prepareWriteLeaseRepair({ + token: input.token, + expectedAcquiredAt: input.expectedAcquiredAt, + }) + if (prepared.prepared) { + const now = utcSqliteTimestamp() + await insertOrVerifyDoRepairAudit({ + db: input.db, + repairId: prepared.repairId, + stableUserId: input.stableUserId, + token: prepared.token, + holder: prepared.holder, + acquiredAt: prepared.acquiredAt, + repairedByUserId: input.repairedByUserId, + reason, + now, + }) + // Finalize DO before D1 mirror clear; failed finalize leaves both held. + await meter.finalizeWriteLeaseRepair({ + token: prepared.token, + repairId: prepared.repairId, + expectedAcquiredAt: prepared.acquiredAt, + }) + await releaseD1WriteLeaseMirror({ + db: input.db, + stableUserId: input.stableUserId, + token: prepared.token, + expectedAcquiredAt: prepared.acquiredAt, + }) + return { repaired: true as const, repairId: prepared.repairId } + } + // Lost-finalize retry: matching audit + absent DO lease clears stale D1 mirror. + const existingAudit = await findMatchingRepairAudit({ + db: input.db, + stableUserId: input.stableUserId, + token: input.token, + expectedAcquiredAt: input.expectedAcquiredAt, + repairedByUserId: input.repairedByUserId, + reason, + }) + if (existingAudit) { + const stillHeld = await meter.assertWriteLeaseHeld({ + token: input.token, + }) + if (!stillHeld.held) { + await releaseD1WriteLeaseMirror({ + db: input.db, + stableUserId: input.stableUserId, + token: input.token, + expectedAcquiredAt: input.expectedAcquiredAt, + }) + return { repaired: true as const, repairId: existingAudit.id } + } + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + } const now = utcSqliteTimestamp() const repairId = crypto.randomUUID() const [audited, decremented, deleted] = await input.db.batch([ @@ -433,7 +748,7 @@ export async function repairAccountWriteLease(input: { .bind( repairId, input.repairedByUserId, - input.reason.trim(), + reason, now, input.token, input.stableUserId, @@ -471,16 +786,5 @@ export async function repairAccountWriteLease(input: { ) { throw new Error('Active account write lease did not match repair request.') } - const repairShadowPromise = scheduleAccountDeletionShadow({ - env: input.env, - waitUntil: input.waitUntil, - task: async (env) => { - await userMeterRpc({ - env, - userId: input.stableUserId, - }).shadowReleaseWriteLease({ token: input.token }) - }, - }) - if (!input.waitUntil) await repairShadowPromise return { repaired: true as const, repairId } } diff --git a/packages/worker/src/app/account-deletion.node.test.ts b/packages/worker/src/app/account-deletion.node.test.ts index 6b2a3af81f..01eb593fe9 100644 --- a/packages/worker/src/app/account-deletion.node.test.ts +++ b/packages/worker/src/app/account-deletion.node.test.ts @@ -37,7 +37,18 @@ function createTestDb( } { const rows: RowMap = {} for (const [key, value] of Object.entries(initial)) { - rows[key] = value.map((row) => ({ ...row })) + rows[key] = value.map((row) => { + const copy = { ...row } + if ( + key === 'users' && + (copy['stable_user_id'] == null || copy['stable_user_id'] === '') + ) { + const id = Number(copy['id']) + if (id === 1) copy['stable_user_id'] = 'user-aaa' + else if (id === 2) copy['stable_user_id'] = 'user-bbb' + } + return copy + }) } function deleteByPredicate( @@ -1414,7 +1425,9 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async for (const run of rows.package_codemod_runs ?? []) { expect(String(run['filters_json'])).not.toContain(userAaa) } - expect(rows.users).toEqual([{ id: 2, email: 'b@example.com' }]) + expect(rows.users).toEqual([ + { id: 2, email: 'b@example.com', stable_user_id: 'user-bbb' }, + ]) expect(result.deletedRowCounts.password_resets).toBe(2) expect(result.deletedRowCounts.user_roles).toBe(1) @@ -1795,19 +1808,6 @@ test('account deletion reports missing Durable Object / blob bindings and remain ]) }, }, - { - envOverrides: { USER_METER: undefined }, - cleanupError: - 'USER_METER binding was unavailable; the user meter Durable Object was not purged.', - partialResult: { - clearedDurableObjects: expect.objectContaining({ - userMeters: 0, - }), - }, - seed: { - users: [{ id: 1, email: 'a@example.com', stable_user_id: 'user-aaa' }], - }, - }, { envOverrides: { MAILBOX: undefined }, cleanupError: @@ -1845,12 +1845,31 @@ test('account deletion reports missing Durable Object / blob bindings and remain }), ]) } + + const { db: missingMeterDb, rows: missingMeterRows } = createTestDb({ + users: [{ id: 1, email: 'a@example.com', stable_user_id: 'user-aaa' }], + }) + await expect( + deleteUserAccount({ + env: createSuccessfulDeletionEnv(missingMeterDb, { + USER_METER: undefined, + }), + dbUserId: 1, + mcpUserId: 'user-aaa', + }), + ).rejects.toThrow('USER_METER Durable Object binding is not configured.') + expect(missingMeterRows.users).toEqual([ + expect.objectContaining({ + id: 1, + deleting_at: expect.any(String), + }), + ]) }) test('deleteUserAccount fails closed when preflight inventory cannot be read', async () => { const { db, rows } = createTestDb( { - users: [{ id: 1, email: 'a@example.com' }], + users: [{ id: 1, email: 'a@example.com', stable_user_id: 'user-aaa' }], mcp_memories: [{ id: 'memory-a', user_id: 'user-aaa' }], jobs: [{ id: 'job-a', user_id: 'user-aaa', storage_id: 'job:job-a' }], }, @@ -1858,10 +1877,12 @@ test('deleteUserAccount fails closed when preflight inventory cannot be read', a ) const deleteVectors = vi.fn(async () => undefined) const clearStorage = vi.fn(async () => undefined) + const userMeter = createInMemoryUserMeterEnv() await expect( deleteUserAccount({ env: { APP_DB: db, + USER_METER: userMeter.env.USER_METER, CAPABILITY_VECTOR_INDEX: { deleteByIds: deleteVectors }, STORAGE_RUNNER: { idFromName: (name: string) => name as unknown as DurableObjectId, diff --git a/packages/worker/src/app/account-write-lease-middleware.node.test.ts b/packages/worker/src/app/account-write-lease-middleware.node.test.ts index 749fcf950a..c8490e43d4 100644 --- a/packages/worker/src/app/account-write-lease-middleware.node.test.ts +++ b/packages/worker/src/app/account-write-lease-middleware.node.test.ts @@ -1,4 +1,6 @@ import { expect, test, vi } from 'vitest' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' const authMock = vi.hoisted(() => ({ loadResolvedRequestAuth: vi.fn(), @@ -19,21 +21,32 @@ test('authenticated delayed mutation holds web lease through handler completion' mcpUser: { userId: 'user-a' }, }, }) - let active = 0 + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) const db = { prepare(query: string) { return { bind() { return { + async first() { + if (query.includes('deleting_at')) { + return { deleting_at: null } + } + return null + }, async run() { - if (query.includes('active_write_count + 1')) active += 1 - if (query.includes('MAX(active_write_count - 1')) active -= 1 return { meta: { changes: 1 } } }, + async all() { + return { results: [], meta: { changes: 0 } } + }, } }, } }, + async batch() { + return [] + }, } as unknown as D1Database let startHandler: () => void = () => undefined let finishHandler: () => void = () => undefined @@ -45,6 +58,7 @@ test('authenticated delayed mutation holds web lease through handler completion' }) const middleware = createAccountWriteLeaseMiddleware({ APP_DB: db, + USER_METER: meter.env.USER_METER, } as Env) const responsePromise = middleware( { @@ -56,14 +70,15 @@ test('authenticated delayed mutation holds web lease through handler completion' context: new Map(), } as never, async () => { + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) startHandler() await finish return Response.json({ ok: true }) }, ) await started - expect(active).toBe(1) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 1 }) finishHandler() await expect(responsePromise).resolves.toMatchObject({ status: 200 }) - expect(active).toBe(0) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) }) diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 722b5289ea..4a3eff0eac 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -26,10 +26,10 @@ export const userMeterDailyCounterRetentionDays = 7 const metaSchemaVersionKey = 'schema_version' /** Bump when initializeSchema DDL changes; warm objects skip DDL. */ -const userMeterSchemaVersion = 6 +const userMeterSchemaVersion = 7 /** Singleton row id for expand-phase D1 storage-byte shadow (schema v4). */ const storageBytesStateRowId = 1 -/** Singleton row id for expand-phase D1 deletion-fence shadow (schema v6). */ +/** Singleton row id for deletion fence / write leases (schema v6+). */ const deletionStateRowId = 1 /** Matches D1 `packageServiceStateStaleMs` (24h); cutover-support RPC only. */ export const userMeterPackageServiceStateStaleMs = 24 * 60 * 60 * 1000 @@ -40,6 +40,7 @@ const maxInboundDeliveryIdLength = 256 const maxPackageServiceIdLength = 256 const maxWriteLeaseTokenLength = 64 const maxWriteLeaseHolderLength = 256 +const maxWriteLeaseRepairIdLength = 64 const maxDeletionTimestampLength = 64 const inboundReceiveResource = 'email_receives_per_day' satisfies DailyEntitlementResource @@ -53,6 +54,8 @@ const packageServiceShadowStatuses = [ 'stopped', 'error', ] as const +/** Write-lease authority: `do` is UserMeter-authoritative; `legacy` mirrors D1. */ +export const userMeterWriteLeaseAuthorities = ['do', 'legacy'] as const export type UserMeterPackageServiceStatus = (typeof packageServiceShadowStatuses)[number] @@ -174,10 +177,22 @@ export type UserMeterPackageServiceBootstrapResult = { skipped: number } +export type UserMeterWriteLeaseAuthority = + (typeof userMeterWriteLeaseAuthorities)[number] + +export function isUserMeterWriteLeaseAuthority( + authority: string, +): authority is UserMeterWriteLeaseAuthority { + return (userMeterWriteLeaseAuthorities as ReadonlyArray).includes( + authority, + ) +} + export type UserMeterWriteLeaseShadow = { token: string holder: string acquiredAt: string + authority: UserMeterWriteLeaseAuthority } /** @@ -190,7 +205,7 @@ export type UserMeterDeletionShadow = { } /** - * Account-export deletion shadow. Omits raw lease token and holder; retains + * Account-export deletion inventory. Omits raw lease token and holder; retains * only deleting tombstone presence, active lease count, and acquired_at. */ export type UserMeterDeletionShadowExport = { @@ -210,14 +225,47 @@ export type UserMeterShadowReplaceDeletionStateResult = { leaseCount: number } +export type UserMeterMarkDeletingResult = { + deletingAt: string + created: boolean + /** Deduped union of preserved DO-authority rows plus D1 legacy snapshot. */ + leaseCount: number +} + export type UserMeterShadowAcquireWriteLeaseResult = { acquired: boolean } +export type UserMeterAcquireWriteLeaseResult = { + acquired: boolean +} + export type UserMeterShadowReleaseWriteLeaseResult = { released: boolean } +export type UserMeterReleaseWriteLeaseResult = { + released: boolean +} + +export type UserMeterAssertWriteLeaseHeldResult = { + held: boolean +} + +export type UserMeterPrepareWriteLeaseRepairResult = + | { + prepared: true + repairId: string + token: string + holder: string + acquiredAt: string + } + | { prepared: false } + +export type UserMeterFinalizeWriteLeaseRepairResult = { + finalized: boolean +} + export type UserMeterWriteLeaseListResult = { leases: Array nextStartAfter: string | null @@ -249,11 +297,10 @@ export type UserMeterExportResult = { */ packageServiceStatesShadow: Array | null /** - * Non-authoritative D1 deletion-fence / write-lease shadow. Emitted only - * on the first export page (`startAfter` absent); subsequent pages return - * `null` so paged consumers never double-count the inventory. Sanitized: - * excludes raw lease token and holder (see - * {@link UserMeterDeletionShadowExport}). + * Sanitized deletion-fence / write-lease inventory. Emitted only on the + * first export page (`startAfter` absent); subsequent pages return `null` + * so paged consumers never double-count. Excludes raw lease token and + * holder (see {@link UserMeterDeletionShadowExport}). */ deletionShadow: UserMeterDeletionShadowExport | null nextStartAfter: string | null @@ -303,6 +350,8 @@ type WriteLeaseSqlRow = { token: string holder: string acquired_at: string + authority: string + pending_repair_id: string | null } function assertDailyResource(resource: string): DailyEntitlementResource { @@ -433,6 +482,30 @@ function assertWriteLeaseHolder(holder: string): string { return holder } +function assertWriteLeaseAuthority( + authority: string, +): UserMeterWriteLeaseAuthority { + if (!isUserMeterWriteLeaseAuthority(authority)) { + throw new Error( + `UserMeter write lease authority must be one of ${userMeterWriteLeaseAuthorities.join(', ')}; got ${JSON.stringify(authority)}.`, + ) + } + return authority +} + +function assertWriteLeaseRepairId(repairId: string): string { + if ( + typeof repairId !== 'string' || + repairId.length === 0 || + repairId.length > maxWriteLeaseRepairIdLength + ) { + throw new Error( + `UserMeter write lease repairId must be a non-empty string up to ${maxWriteLeaseRepairIdLength} characters.`, + ) + } + return repairId +} + function retentionCutoffDay(now: Date): string { const cutoff = new Date(now) cutoff.setUTCDate( @@ -607,7 +680,7 @@ class UserMeterBase extends DurableObject { `CREATE INDEX IF NOT EXISTS idx_package_service_states_status_source ON package_service_states (status, source_updated_at)`, ) - // Expand-phase shadow of D1 deletion fence / write leases (schema v6). + // Deletion fence / write leases (schema v6+; v7 adds authority/pending_repair_id). this.ctx.storage.sql.exec(` CREATE TABLE IF NOT EXISTS deletion_state ( id INTEGER PRIMARY KEY NOT NULL CHECK (id = 1), @@ -618,13 +691,37 @@ class UserMeterBase extends DurableObject { CREATE TABLE IF NOT EXISTS account_write_leases ( token TEXT PRIMARY KEY NOT NULL, holder TEXT NOT NULL, - acquired_at TEXT NOT NULL + acquired_at TEXT NOT NULL, + authority TEXT NOT NULL DEFAULT 'legacy', + pending_repair_id TEXT ) `) + if (version != null && version < 7) { + try { + this.ctx.storage.sql.exec( + `ALTER TABLE account_write_leases + ADD COLUMN authority TEXT NOT NULL DEFAULT 'legacy'`, + ) + } catch { + // Column already present on a partially migrated object. + } + try { + this.ctx.storage.sql.exec( + `ALTER TABLE account_write_leases + ADD COLUMN pending_repair_id TEXT`, + ) + } catch { + // Column already present on a partially migrated object. + } + } this.ctx.storage.sql.exec( `CREATE INDEX IF NOT EXISTS idx_account_write_leases_acquired_token ON account_write_leases (acquired_at, token)`, ) + this.ctx.storage.sql.exec( + `CREATE INDEX IF NOT EXISTS idx_account_write_leases_authority_acquired_token + ON account_write_leases (authority, acquired_at, token)`, + ) this.ctx.storage.sql.exec( `INSERT INTO user_meter_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value`, @@ -1381,22 +1478,36 @@ class UserMeterBase extends DurableObject { } private writeLeaseFromRow(row: WriteLeaseSqlRow): UserMeterWriteLeaseShadow { + const authorityRaw = String(row.authority ?? 'legacy') return { token: String(row.token), holder: String(row.holder), acquiredAt: String(row.acquired_at), + authority: isUserMeterWriteLeaseAuthority(authorityRaw) + ? authorityRaw + : 'legacy', } } - private readWriteLease(token: string): UserMeterWriteLeaseShadow | null { + private pendingRepairIdFromRow(row: WriteLeaseSqlRow): string | null { + const pending = row.pending_repair_id + return typeof pending === 'string' && pending.length > 0 ? pending : null + } + + private readWriteLeaseRow(token: string): WriteLeaseSqlRow | null { const row = this.ctx.storage.sql .exec( - `SELECT token, holder, acquired_at + `SELECT token, holder, acquired_at, authority, pending_repair_id FROM account_write_leases WHERE token = ?`, token, ) .toArray()[0] + return row ?? null + } + + private readWriteLease(token: string): UserMeterWriteLeaseShadow | null { + const row = this.readWriteLeaseRow(token) if (!row) return null return this.writeLeaseFromRow(row) } @@ -1404,7 +1515,7 @@ class UserMeterBase extends DurableObject { private listAllWriteLeaseRows(): Array { const rows = this.ctx.storage.sql .exec( - `SELECT token, holder, acquired_at + `SELECT token, holder, acquired_at, authority, pending_repair_id FROM account_write_leases ORDER BY acquired_at ASC, token ASC`, ) @@ -1412,22 +1523,19 @@ class UserMeterBase extends DurableObject { return rows.map((row) => this.writeLeaseFromRow(row)) } - private readDeletionShadowExport(): UserMeterDeletionShadowExport { - const writeLeases = this.listAllWriteLeaseRows().map((lease) => ({ - acquiredAt: lease.acquiredAt, - })) - return { - deletingAt: this.readDeletingAt(), - activeWriteLeaseCount: writeLeases.length, - writeLeases, - } + private countWriteLeases(): number { + const row = this.ctx.storage.sql + .exec<{ + count: number + }>(`SELECT COUNT(*) AS count FROM account_write_leases`) + .toArray()[0] + return Math.max(0, Number(row?.count ?? 0)) } - /** Expand-phase shadow of D1 `users.deleting_at` (first write wins). */ - async shadowMarkDeleting(input: { + private insertOrPreserveDeletingAt(deletingAt: string): { deletingAt: string - }): Promise { - const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + created: boolean + } { const existing = this.readDeletingAt() if (existing != null) { return { deletingAt: existing, created: false } @@ -1439,30 +1547,58 @@ class UserMeterBase extends DurableObject { deletionStateRowId, deletingAt, ) - const stored = this.readDeletingAt() ?? deletingAt - return { deletingAt: stored, created: cursor.rowsWritten > 0 } + return { + deletingAt: this.readDeletingAt() ?? deletingAt, + created: cursor.rowsWritten > 0, + } } /** - * Atomically set/preserve the deleting tombstone and replace the shadow - * lease set with exactly the supplied active D1 leases. Used by - * `markAccountDeleting` so stale unreleased shadow rows cannot linger - * after D1 drain (empty lease list clears all shadow leases). + * Replace only legacy/D1-mirror rows with the supplied snapshot; preserve + * DO-authority rows (and any pending repair state on those rows). */ - async shadowReplaceDeletionState(input: { - deletingAt: string - leases?: ReadonlyArray<{ + private replaceLegacyWriteLeases( + leases: ReadonlyArray<{ token: string holder: string acquiredAt: string - }> - }): Promise { - const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + }>, + ): number { + this.ctx.storage.sql.exec( + `DELETE FROM account_write_leases WHERE authority = 'legacy'`, + ) + for (const lease of leases) { + this.ctx.storage.sql.exec( + `INSERT INTO account_write_leases ( + token, holder, acquired_at, authority, pending_repair_id + ) + VALUES (?, ?, ?, 'legacy', NULL) + ON CONFLICT(token) DO UPDATE SET + holder = excluded.holder, + acquired_at = excluded.acquired_at, + authority = 'legacy', + pending_repair_id = NULL + WHERE account_write_leases.authority = 'legacy'`, + lease.token, + lease.holder, + lease.acquiredAt, + ) + } + return this.countWriteLeases() + } + + private normalizeLegacyLeaseSnapshot( + leases: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }>, + ) { const leasesByToken = new Map< string, { token: string; holder: string; acquiredAt: string } >() - for (const lease of input.leases ?? []) { + for (const lease of leases) { const token = assertWriteLeaseToken(lease.token) leasesByToken.set(token, { token, @@ -1470,40 +1606,146 @@ class UserMeterBase extends DurableObject { acquiredAt: assertDeletionTimestamp('acquiredAt', lease.acquiredAt), }) } - const leases = [...leasesByToken.values()] + return [...leasesByToken.values()] + } + + private listWriteLeasesPage(input: { + pageSize?: number + startAfter?: string | null + authority?: UserMeterWriteLeaseAuthority + }): UserMeterWriteLeaseListResult { + const pageSize = normalizePageSize(input.pageSize) + const cursor = + typeof input.startAfter === 'string' && input.startAfter.length > 0 + ? decodeWriteLeaseCursor(input.startAfter) + : null + const authority = input.authority + const rows = ( + authority + ? cursor + ? this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at, authority, pending_repair_id + FROM account_write_leases + WHERE authority = ? + AND ( + acquired_at > ? + OR (acquired_at = ? AND token > ?) + ) + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + authority, + cursor.acquiredAt, + cursor.acquiredAt, + cursor.token, + pageSize + 1, + ) + : this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at, authority, pending_repair_id + FROM account_write_leases + WHERE authority = ? + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + authority, + pageSize + 1, + ) + : cursor + ? this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at, authority, pending_repair_id + FROM account_write_leases + WHERE acquired_at > ? + OR (acquired_at = ? AND token > ?) + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + cursor.acquiredAt, + cursor.acquiredAt, + cursor.token, + pageSize + 1, + ) + : this.ctx.storage.sql.exec( + `SELECT token, holder, acquired_at, authority, pending_repair_id + FROM account_write_leases + ORDER BY acquired_at ASC, token ASC + LIMIT ?`, + pageSize + 1, + ) + ).toArray() + const truncated = rows.length > pageSize + const pageRows = truncated ? rows.slice(0, pageSize) : rows + const leases = pageRows.map((row) => this.writeLeaseFromRow(row)) + const last = pageRows[pageRows.length - 1] + return { + leases, + nextStartAfter: + truncated && last + ? encodeWriteLeaseCursor({ + acquiredAt: String(last.acquired_at), + token: String(last.token), + }) + : null, + truncated, + } + } + + private readDeletionShadowExport(): UserMeterDeletionShadowExport { + const writeLeases = this.listAllWriteLeaseRows().map((lease) => ({ + acquiredAt: lease.acquiredAt, + })) + return { + deletingAt: this.readDeletingAt(), + activeWriteLeaseCount: writeLeases.length, + writeLeases, + } + } + + /** Legacy/compat mark of D1 `users.deleting_at` (first write wins). */ + async shadowMarkDeleting(input: { + deletingAt: string + }): Promise { + const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + return this.insertOrPreserveDeletingAt(deletingAt) + } + + /** Legacy/compat replace of legacy rows from a D1 snapshot. */ + async shadowReplaceDeletionState(input: { + deletingAt: string + leases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }): Promise { + return this.markDeleting({ + deletingAt: input.deletingAt, + legacyLeases: input.leases, + }) + } + + /** + * Authoritative mark: preserve tombstone + DO-authority rows, replace only + * legacy rows from the D1 snapshot, return deduped union lease count. + */ + async markDeleting(input: { + deletingAt: string + legacyLeases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }): Promise { + const deletingAt = assertDeletionTimestamp('deletingAt', input.deletingAt) + const leases = this.normalizeLegacyLeaseSnapshot(input.legacyLeases ?? []) return await this.ctx.blockConcurrencyWhile(async () => { - const existing = this.readDeletingAt() - let created = false - if (existing == null) { - const cursor = this.ctx.storage.sql.exec( - `INSERT INTO deletion_state (id, deleting_at) - VALUES (?, ?) - ON CONFLICT(id) DO NOTHING`, - deletionStateRowId, - deletingAt, - ) - created = cursor.rowsWritten > 0 - } - const stored = this.readDeletingAt() ?? deletingAt - this.ctx.storage.sql.exec(`DELETE FROM account_write_leases`) - for (const lease of leases) { - this.ctx.storage.sql.exec( - `INSERT INTO account_write_leases (token, holder, acquired_at) - VALUES (?, ?, ?)`, - lease.token, - lease.holder, - lease.acquiredAt, - ) - } + const marked = this.insertOrPreserveDeletingAt(deletingAt) + const leaseCount = this.replaceLegacyWriteLeases(leases) return { - deletingAt: stored, - created, - leaseCount: leases.length, + deletingAt: marked.deletingAt, + created: marked.created, + leaseCount, } }) } - /** Expand-phase shadow lease acquire; fencing still uses D1. */ + /** Legacy/compat lease acquire (`authority='legacy'`). */ async shadowAcquireWriteLease(input: { token: string holder: string @@ -1512,11 +1754,14 @@ class UserMeterBase extends DurableObject { const token = assertWriteLeaseToken(input.token) const holder = assertWriteLeaseHolder(input.holder) const acquiredAt = assertDeletionTimestamp('acquiredAt', input.acquiredAt) - if (this.readWriteLease(token)) return { acquired: true } + const existing = this.readWriteLease(token) + if (existing) return { acquired: true } if (this.readDeletingAt() != null) return { acquired: false } const cursor = this.ctx.storage.sql.exec( - `INSERT INTO account_write_leases (token, holder, acquired_at) - VALUES (?, ?, ?) + `INSERT INTO account_write_leases ( + token, holder, acquired_at, authority, pending_repair_id + ) + VALUES (?, ?, ?, 'legacy', NULL) ON CONFLICT(token) DO NOTHING`, token, holder, @@ -1527,7 +1772,38 @@ class UserMeterBase extends DurableObject { } } - /** Expand-phase shadow lease release / repair mirror; fencing uses D1. */ + /** Authoritative lease acquire (`authority='do'`). */ + async acquireWriteLease(input: { + token: string + holder: string + acquiredAt: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const holder = assertWriteLeaseHolder(input.holder) + const acquiredAt = assertDeletionTimestamp('acquiredAt', input.acquiredAt) + const existing = this.readWriteLease(token) + if (existing) { + return { acquired: existing.authority === 'do' } + } + if (this.readDeletingAt() != null) return { acquired: false } + const cursor = this.ctx.storage.sql.exec( + `INSERT INTO account_write_leases ( + token, holder, acquired_at, authority, pending_repair_id + ) + VALUES (?, ?, ?, 'do', NULL) + ON CONFLICT(token) DO NOTHING`, + token, + holder, + acquiredAt, + ) + const held = this.readWriteLease(token) + return { + acquired: + (cursor.rowsWritten > 0 || held != null) && held?.authority === 'do', + } + } + + /** Legacy/compat lease release (any authority). */ async shadowReleaseWriteLease(input: { token: string }): Promise { @@ -1539,74 +1815,155 @@ class UserMeterBase extends DurableObject { return { released: cursor.rowsWritten > 0 } } - /** Cutover-support deletion tombstone read; Phase A uses D1. */ + /** Authoritative release of a DO-authority lease. */ + async releaseWriteLease(input: { + token: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const cursor = this.ctx.storage.sql.exec( + `DELETE FROM account_write_leases + WHERE token = ? AND authority = 'do'`, + token, + ) + return { released: cursor.rowsWritten > 0 } + } + + /** Post-write held check; pending repair still counts as held. */ + async assertWriteLeaseHeld(input: { + token: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const lease = this.readWriteLease(token) + return { held: lease?.authority === 'do' } + } + + /** Prepare audit-safe DO repair; retries reuse the pending `repairId`. */ + async prepareWriteLeaseRepair(input: { + token: string + expectedAcquiredAt: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const expectedAcquiredAt = assertDeletionTimestamp( + 'expectedAcquiredAt', + input.expectedAcquiredAt, + ) + return await this.ctx.blockConcurrencyWhile(async () => { + const row = this.readWriteLeaseRow(token) + if (!row) return { prepared: false as const } + const lease = this.writeLeaseFromRow(row) + if (lease.authority !== 'do') return { prepared: false as const } + if (lease.acquiredAt !== expectedAcquiredAt) { + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + const existingRepairId = this.pendingRepairIdFromRow(row) + const repairId = existingRepairId ?? crypto.randomUUID() + if (!existingRepairId) { + this.ctx.storage.sql.exec( + `UPDATE account_write_leases + SET pending_repair_id = ? + WHERE token = ? AND authority = 'do' AND acquired_at = ? + AND (pending_repair_id IS NULL OR pending_repair_id = '')`, + repairId, + token, + expectedAcquiredAt, + ) + } + const heldRow = this.readWriteLeaseRow(token) + if (!heldRow) return { prepared: false as const } + const held = this.writeLeaseFromRow(heldRow) + if (held.authority !== 'do') return { prepared: false as const } + const pending = this.pendingRepairIdFromRow(heldRow) + if (!pending) { + throw new Error('Account write lease repair could not be prepared.') + } + return { + prepared: true as const, + repairId: pending, + token: held.token, + holder: held.holder, + acquiredAt: held.acquiredAt, + } + }) + } + + /** Finalize exact pending DO repair; idempotent when the lease is already gone. */ + async finalizeWriteLeaseRepair(input: { + token: string + repairId: string + expectedAcquiredAt: string + }): Promise { + const token = assertWriteLeaseToken(input.token) + const repairId = assertWriteLeaseRepairId(input.repairId) + const expectedAcquiredAt = assertDeletionTimestamp( + 'expectedAcquiredAt', + input.expectedAcquiredAt, + ) + return await this.ctx.blockConcurrencyWhile(async () => { + const row = this.readWriteLeaseRow(token) + if (!row) return { finalized: true } + const lease = this.writeLeaseFromRow(row) + if ( + lease.authority !== 'do' || + lease.acquiredAt !== expectedAcquiredAt || + this.pendingRepairIdFromRow(row) !== repairId + ) { + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + const cursor = this.ctx.storage.sql.exec( + `DELETE FROM account_write_leases + WHERE token = ? + AND authority = 'do' + AND acquired_at = ? + AND pending_repair_id = ?`, + token, + expectedAcquiredAt, + repairId, + ) + if (cursor.rowsWritten < 1 && this.readWriteLeaseRow(token) != null) { + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + return { finalized: true } + }) + } + + /** Deletion tombstone read (D1 `deleting_at` remains the permanent gate). */ async readDeletionState(): Promise<{ deletingAt: string | null }> { return { deletingAt: this.readDeletingAt() } } - /** Cutover-support paged shadow lease list; Phase A list uses D1. */ + /** Paged lease list (legacy + DO authority). */ async listWriteLeases( input: { pageSize?: number startAfter?: string | null } = {}, ): Promise { - const pageSize = normalizePageSize(input.pageSize) - const cursor = - typeof input.startAfter === 'string' && input.startAfter.length > 0 - ? decodeWriteLeaseCursor(input.startAfter) - : null - const rows = ( - cursor - ? this.ctx.storage.sql.exec( - `SELECT token, holder, acquired_at - FROM account_write_leases - WHERE acquired_at > ? - OR (acquired_at = ? AND token > ?) - ORDER BY acquired_at ASC, token ASC - LIMIT ?`, - cursor.acquiredAt, - cursor.acquiredAt, - cursor.token, - pageSize + 1, - ) - : this.ctx.storage.sql.exec( - `SELECT token, holder, acquired_at - FROM account_write_leases - ORDER BY acquired_at ASC, token ASC - LIMIT ?`, - pageSize + 1, - ) - ).toArray() - const truncated = rows.length > pageSize - const pageRows = truncated ? rows.slice(0, pageSize) : rows - const leases = pageRows.map((row) => this.writeLeaseFromRow(row)) - const last = pageRows[pageRows.length - 1] - return { - leases, - nextStartAfter: - truncated && last - ? encodeWriteLeaseCursor({ - acquiredAt: String(last.acquired_at), - token: String(last.token), - }) - : null, - truncated, - } + return this.listWriteLeasesPage(input) + } + + /** Paged DO-authority lease list for admin union inventory. */ + async listDoAuthorityWriteLeases( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ): Promise { + return this.listWriteLeasesPage({ ...input, authority: 'do' }) } - /** Cutover-support active lease count; Phase A drain uses D1. */ + /** Active lease count (legacy + DO authority; pending repair still counts). */ async countActiveWriteLeases(): Promise { - const row = this.ctx.storage.sql - .exec<{ - count: number - }>(`SELECT COUNT(*) AS count FROM account_write_leases`) - .toArray()[0] - return { count: Math.max(0, Number(row?.count ?? 0)) } + return { count: this.countWriteLeases() } } /** - * Cutover-support cold seed. Leases apply first (same guards as + * Cold seed. Leases apply first (same guards as * {@link shadowAcquireWriteLease}), then the deleting tombstone. */ async bootstrapDeletionState(input: { @@ -1615,6 +1972,7 @@ class UserMeterBase extends DurableObject { token: string holder: string acquiredAt: string + authority?: string }> }): Promise { let leasesApplied = 0 @@ -1625,7 +1983,14 @@ class UserMeterBase extends DurableObject { leasesSkipped += 1 continue } - const result = await this.shadowAcquireWriteLease(lease) + const authority = + typeof lease.authority === 'string' + ? assertWriteLeaseAuthority(lease.authority) + : 'legacy' + const result = + authority === 'do' + ? await this.acquireWriteLease(lease) + : await this.shadowAcquireWriteLease(lease) if (result.acquired) leasesApplied += 1 else leasesSkipped += 1 } @@ -1641,7 +2006,6 @@ class UserMeterBase extends DurableObject { async purge(): Promise<{ ok: true }> { await this.ctx.blockConcurrencyWhile(async () => { - // Preserve deleting tombstone across deleteAll for later cutover safety. const deletingAt = this.readDeletingAt() await this.ctx.storage.deleteAll() this.initializeSchema() @@ -1845,13 +2209,13 @@ export type UserMeterRpc = { updatedAt?: string }> }) => Promise - /** Expand-phase shadow of D1 users.deleting_at; fence still uses D1. */ + /** Legacy/compat mark of D1 users.deleting_at (first write wins). */ shadowMarkDeleting: (input: { deletingAt: string }) => Promise /** - * Atomically set/preserve deleting tombstone and replace shadow leases - * with the supplied active D1 set (empty clears stale shadows). + * Legacy/compat replace: set/preserve tombstone and replace only legacy + * rows with the supplied D1 set (preserves DO-authority rows). */ shadowReplaceDeletionState: (input: { deletingAt: string @@ -1861,32 +2225,75 @@ export type UserMeterRpc = { acquiredAt: string }> }) => Promise - /** Expand-phase shadow lease acquire; fencing still uses D1. */ + /** + * Authoritative deletion mark with D1 legacy lease snapshot. Preserves + * DO-authority rows and returns the deduped union lease count. + */ + markDeleting: (input: { + deletingAt: string + legacyLeases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) => Promise + /** Legacy/compat lease acquire (`authority='legacy'`). */ shadowAcquireWriteLease: (input: { token: string holder: string acquiredAt: string }) => Promise - /** Expand-phase shadow lease release / repair mirror; fencing uses D1. */ + /** Authoritative lease acquire (`authority='do'`). */ + acquireWriteLease: (input: { + token: string + holder: string + acquiredAt: string + }) => Promise + /** Legacy/compat lease release (any authority). */ shadowReleaseWriteLease: (input: { token: string }) => Promise - /** Cutover-support deletion tombstone read; Phase A uses D1. */ + /** Authoritative release of a DO-authority lease. */ + releaseWriteLease: (input: { + token: string + }) => Promise + /** Post-write held check; pending repair still counts as held. */ + assertWriteLeaseHeld: (input: { + token: string + }) => Promise + /** Prepare an audit-safe DO lease repair; retries reuse repairId. */ + prepareWriteLeaseRepair: (input: { + token: string + expectedAcquiredAt: string + }) => Promise + /** Finalize an exact pending DO repair by deleting the lease. */ + finalizeWriteLeaseRepair: (input: { + token: string + repairId: string + expectedAcquiredAt: string + }) => Promise + /** Deletion tombstone read (D1 deleting_at remains the permanent gate). */ readDeletionState: () => Promise<{ deletingAt: string | null }> - /** Cutover-support paged shadow lease list; Phase A list uses D1. */ + /** Paged lease list (legacy + DO authority). */ listWriteLeases: (input: { pageSize?: number startAfter?: string | null }) => Promise - /** Cutover-support active lease count; Phase A drain uses D1. */ + /** Paged DO-authority lease list for admin union inventory. */ + listDoAuthorityWriteLeases: (input: { + pageSize?: number + startAfter?: string | null + }) => Promise + /** Active lease count (legacy + DO authority; pending repair still counts). */ countActiveWriteLeases: () => Promise - /** Cutover-support cold seed with monotonic mark / acquire guards. */ + /** Cold seed with monotonic mark / acquire guards. */ bootstrapDeletionState: (input: { deletingAt?: string | null leases?: ReadonlyArray<{ token: string holder: string acquiredAt: string + authority?: string }> }) => Promise purge: () => Promise<{ ok: true }> diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index a43054c3a1..823adf7f6d 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -1277,36 +1277,40 @@ test('UserMeter package-service shadow upserts are monotonic, isolated, exportab }) }, 30_000) -test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserves tombstone across purge', async () => { +test('UserMeter deletion leases distinguish do/legacy authority, repair, export, and purge tombstone', async () => { const userA = await seedFreeUser('meter-deletion-a') const userB = await seedFreeUser('meter-deletion-b') const meterA = userMeterRpc({ env, userId: userA.userId }) const meterB = userMeterRpc({ env, userId: userB.userId }) - const firstMark = await meterA.shadowMarkDeleting({ + const firstMark = await meterA.markDeleting({ deletingAt: '2026-08-01 10:00:00', + legacyLeases: [], }) expect(firstMark).toEqual({ deletingAt: '2026-08-01 10:00:00', created: true, + leaseCount: 0, }) - const secondMark = await meterA.shadowMarkDeleting({ + const secondMark = await meterA.markDeleting({ deletingAt: '2026-08-01 11:00:00', + legacyLeases: [], }) expect(secondMark).toEqual({ deletingAt: '2026-08-01 10:00:00', created: false, + leaseCount: 0, }) await expect( - meterB.shadowAcquireWriteLease({ + meterB.acquireWriteLease({ token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', holder: 'test:writer', acquiredAt: '2026-08-01 10:05:00', }), ).resolves.toEqual({ acquired: true }) await expect( - meterB.shadowAcquireWriteLease({ + meterB.acquireWriteLease({ token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', holder: 'test:writer', acquiredAt: '2026-08-01 10:05:00', @@ -1315,11 +1319,23 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve await expect( meterB.shadowAcquireWriteLease({ token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', - holder: 'test:other', + holder: 'test:legacy', acquiredAt: '2026-08-01 10:06:00', }), ).resolves.toEqual({ acquired: true }) expect(await meterB.countActiveWriteLeases()).toEqual({ count: 2 }) + expect(await meterB.listDoAuthorityWriteLeases({})).toEqual({ + leases: [ + { + token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', + holder: 'test:writer', + acquiredAt: '2026-08-01 10:05:00', + authority: 'do', + }, + ], + nextStartAfter: null, + truncated: false, + }) const listed = await meterB.listWriteLeases({ pageSize: 1 }) expect(listed.leases).toHaveLength(1) @@ -1332,14 +1348,14 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve expect(listedRest.truncated).toBe(false) await expect( - meterB.shadowReleaseWriteLease({ + meterB.releaseWriteLease({ token: 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa', }), ).resolves.toEqual({ released: true }) expect(await meterB.countActiveWriteLeases()).toEqual({ count: 1 }) await expect( - meterA.shadowAcquireWriteLease({ + meterA.acquireWriteLease({ token: 'cccccccc-cccc-4ccc-8ccc-cccccccccccc', holder: 'test:blocked', acquiredAt: '2026-08-01 10:07:00', @@ -1351,13 +1367,14 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve leases: [ { token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', - holder: 'test:other', + holder: 'test:legacy', acquiredAt: '2026-08-01 10:06:00', }, { token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', - holder: 'test:boot', + holder: 'test:boot-do', acquiredAt: '2026-08-01 10:08:00', + authority: 'do', }, ], }) @@ -1367,6 +1384,84 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve leasesSkipped: 1, }) + const marked = await meterB.markDeleting({ + deletingAt: '2026-08-01 12:30:00', + legacyLeases: [ + { + token: 'eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee', + holder: 'test:d1-snapshot', + acquiredAt: '2026-08-01 10:09:00', + }, + ], + }) + expect(marked).toEqual({ + deletingAt: '2026-08-01 12:00:00', + created: false, + leaseCount: 2, + }) + expect(await meterB.listWriteLeases({})).toEqual({ + leases: [ + { + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + holder: 'test:boot-do', + acquiredAt: '2026-08-01 10:08:00', + authority: 'do', + }, + { + token: 'eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee', + holder: 'test:d1-snapshot', + acquiredAt: '2026-08-01 10:09:00', + authority: 'legacy', + }, + ], + nextStartAfter: null, + truncated: false, + }) + + const prepared = await meterB.prepareWriteLeaseRepair({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + expectedAcquiredAt: '2026-08-01 10:08:00', + }) + expect(prepared).toEqual( + expect.objectContaining({ + prepared: true, + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + acquiredAt: '2026-08-01 10:08:00', + }), + ) + const repairId = + prepared.prepared === true ? prepared.repairId : 'missing-repair-id' + await expect( + meterB.prepareWriteLeaseRepair({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + expectedAcquiredAt: '2026-08-01 10:08:00', + }), + ).resolves.toEqual(expect.objectContaining({ prepared: true, repairId })) + await expect( + meterB.assertWriteLeaseHeld({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + }), + ).resolves.toEqual({ held: true }) + await expect( + meterB.finalizeWriteLeaseRepair({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + repairId, + expectedAcquiredAt: '2026-08-01 10:08:00', + }), + ).resolves.toEqual({ finalized: true }) + await expect( + meterB.assertWriteLeaseHeld({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + }), + ).resolves.toEqual({ held: false }) + await expect( + meterB.finalizeWriteLeaseRepair({ + token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', + repairId, + expectedAcquiredAt: '2026-08-01 10:08:00', + }), + ).resolves.toEqual({ finalized: true }) + const day = '2026-08-01' for (const resource of [ 'email_receives_per_day', @@ -1413,8 +1508,8 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve truncated: false, }) await expect( - meterA.shadowAcquireWriteLease({ - token: 'eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee', + meterA.acquireWriteLease({ + token: 'ffffffff-ffff-4fff-8fff-ffffffffffff', holder: 'test:post-purge', acquiredAt: '2026-08-01 13:00:00', }), @@ -1423,54 +1518,10 @@ test('UserMeter deletion shadow is monotonic, isolated, exportable, and preserve expect(await meterB.readDeletionState()).toEqual({ deletingAt: '2026-08-01 12:00:00', }) - expect(await meterB.listWriteLeases({})).toMatchObject({ - leases: expect.arrayContaining([ - expect.objectContaining({ - token: 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb', - }), - expect.objectContaining({ - token: 'dddddddd-dddd-4ddd-8ddd-dddddddddddd', - }), - ]), - }) - - const replaced = await meterB.shadowReplaceDeletionState({ - deletingAt: '2026-08-01 14:00:00', - leases: [ - { - token: 'ffffffff-ffff-4fff-8fff-ffffffffffff', - holder: 'test:replace', - acquiredAt: '2026-08-01 10:09:00', - }, - ], - }) - expect(replaced).toEqual({ - deletingAt: '2026-08-01 12:00:00', - created: false, - leaseCount: 1, - }) - expect(await meterB.listWriteLeases({})).toEqual({ - leases: [ - { - token: 'ffffffff-ffff-4fff-8fff-ffffffffffff', - holder: 'test:replace', - acquiredAt: '2026-08-01 10:09:00', - }, - ], - nextStartAfter: null, - truncated: false, - }) - expect(await meterB.exportCounters({})).toMatchObject({ - deletionShadow: { - deletingAt: '2026-08-01 12:00:00', - activeWriteLeaseCount: 1, - writeLeases: [{ acquiredAt: '2026-08-01 10:09:00' }], - }, - }) await expect( - meterB.shadowReplaceDeletionState({ + meterB.markDeleting({ deletingAt: '2026-08-01 15:00:00', - leases: [], + legacyLeases: [], }), ).resolves.toEqual({ deletingAt: '2026-08-01 12:00:00', diff --git a/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-list.ts b/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-list.ts index 81be561735..34efc1f074 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-list.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-account-write-lease-list.ts @@ -35,6 +35,7 @@ export const adminAccountWriteLeaseListCapability = defineDomainCapability( leases: await listActiveAccountWriteLeases( ctx.env.APP_DB, args.stable_user_id, + ctx.env, ), }), ) diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index fab87a0953..8b274cb760 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -1,6 +1,7 @@ import { isDailyEntitlementResource, isUserMeterPackageServiceStatus, + isUserMeterWriteLeaseAuthority, userMeterMirrorUpdatedAtToken, userMeterPackageServiceStateStaleMs, type DailyEntitlementResource, @@ -8,6 +9,7 @@ import { type UserMeterPackageServiceState, type UserMeterPackageServiceStatus, type UserMeterStorageBytesState, + type UserMeterWriteLeaseAuthority, type UserMeterWriteLeaseShadow, } from '#worker/entitlements/user-meter-do.ts' import { type UserMeterEnv } from '#worker/entitlements/user-meter-client.ts' @@ -24,6 +26,8 @@ type PackageServiceRow = { type WriteLeaseRow = { holder: string acquiredAt: string + authority: UserMeterWriteLeaseAuthority + pendingRepairId: string | null } type DeletionState = { deletingAt: string | null @@ -137,11 +141,13 @@ export function createInMemoryUserMeterEnv() { token, holder: row.holder, acquiredAt: row.acquiredAt, + authority: row.authority, } } - function listWriteLeasesSorted() { + function listWriteLeasesSorted(authority?: UserMeterWriteLeaseAuthority) { return [...deletion.leases.entries()] + .filter(([, row]) => authority == null || row.authority === authority) .map(([token, row]) => writeLeaseState(token, row)) .sort((left, right) => { const byAcquired = left.acquiredAt.localeCompare(right.acquiredAt) @@ -150,6 +156,85 @@ export function createInMemoryUserMeterEnv() { }) } + function replaceLegacyLeases( + leases: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }>, + ) { + for (const [token, row] of [...deletion.leases.entries()]) { + if (row.authority === 'legacy') deletion.leases.delete(token) + } + const leasesByToken = new Map< + string, + { holder: string; acquiredAt: string } + >() + for (const lease of leases) { + leasesByToken.set(lease.token, { + holder: lease.holder, + acquiredAt: lease.acquiredAt, + }) + } + for (const [token, row] of leasesByToken) { + const existing = deletion.leases.get(token) + if (existing?.authority === 'do') continue + deletion.leases.set(token, { + holder: row.holder, + acquiredAt: row.acquiredAt, + authority: 'legacy', + pendingRepairId: null, + }) + } + return deletion.leases.size + } + + function pageWriteLeases( + input: { + pageSize?: number + startAfter?: string | null + }, + authority?: UserMeterWriteLeaseAuthority, + ) { + const pageSize = + typeof input.pageSize === 'number' && Number.isFinite(input.pageSize) + ? Math.min(Math.max(Math.trunc(input.pageSize), 1), 500) + : 100 + const all = listWriteLeasesSorted(authority) + let startIndex = 0 + if (typeof input.startAfter === 'string' && input.startAfter.length > 0) { + try { + const parsed = JSON.parse(input.startAfter) as unknown + if ( + Array.isArray(parsed) && + parsed.length === 2 && + typeof parsed[0] === 'string' && + typeof parsed[1] === 'string' + ) { + const acquiredAt = parsed[0] + const token = parsed[1] + startIndex = + all.findIndex( + (row) => row.acquiredAt === acquiredAt && row.token === token, + ) + 1 + } + } catch { + startIndex = 0 + } + } + const page = all.slice(startIndex, startIndex + pageSize) + const truncated = startIndex + pageSize < all.length + const last = page[page.length - 1] + return { + leases: page, + nextStartAfter: + truncated && last + ? JSON.stringify([last.acquiredAt, last.token]) + : null, + truncated, + } + } + function readDeletionShadowExport(): UserMeterDeletionShadowExport { const writeLeases = listWriteLeasesSorted().map((lease) => ({ acquiredAt: lease.acquiredAt, @@ -466,24 +551,28 @@ export function createInMemoryUserMeterEnv() { }) { const created = deletion.deletingAt == null if (created) deletion.deletingAt = input.deletingAt - deletion.leases.clear() - const leasesByToken = new Map< - string, - { holder: string; acquiredAt: string } - >() - for (const lease of input.leases ?? []) { - leasesByToken.set(lease.token, { - holder: lease.holder, - acquiredAt: lease.acquiredAt, - }) - } - for (const [token, row] of leasesByToken) { - deletion.leases.set(token, row) + const leaseCount = replaceLegacyLeases(input.leases ?? []) + return { + deletingAt: deletion.deletingAt ?? input.deletingAt, + created, + leaseCount, } + }, + async markDeleting(input: { + deletingAt: string + legacyLeases?: ReadonlyArray<{ + token: string + holder: string + acquiredAt: string + }> + }) { + const created = deletion.deletingAt == null + if (created) deletion.deletingAt = input.deletingAt + const leaseCount = replaceLegacyLeases(input.legacyLeases ?? []) return { deletingAt: deletion.deletingAt ?? input.deletingAt, created, - leaseCount: deletion.leases.size, + leaseCount, } }, async shadowAcquireWriteLease(input: { @@ -496,12 +585,84 @@ export function createInMemoryUserMeterEnv() { deletion.leases.set(input.token, { holder: input.holder, acquiredAt: input.acquiredAt, + authority: 'legacy', + pendingRepairId: null, + }) + return { acquired: true } + }, + async acquireWriteLease(input: { + token: string + holder: string + acquiredAt: string + }) { + const existing = deletion.leases.get(input.token) + if (existing) return { acquired: existing.authority === 'do' } + if (deletion.deletingAt != null) return { acquired: false } + deletion.leases.set(input.token, { + holder: input.holder, + acquiredAt: input.acquiredAt, + authority: 'do', + pendingRepairId: null, }) return { acquired: true } }, async shadowReleaseWriteLease(input: { token: string }) { return { released: deletion.leases.delete(input.token) } }, + async releaseWriteLease(input: { token: string }) { + const existing = deletion.leases.get(input.token) + if (!existing || existing.authority !== 'do') { + return { released: false } + } + return { released: deletion.leases.delete(input.token) } + }, + async assertWriteLeaseHeld(input: { token: string }) { + const existing = deletion.leases.get(input.token) + return { held: existing != null && existing.authority === 'do' } + }, + async prepareWriteLeaseRepair(input: { + token: string + expectedAcquiredAt: string + }) { + const existing = deletion.leases.get(input.token) + if (!existing || existing.authority !== 'do') { + return { prepared: false as const } + } + if (existing.acquiredAt !== input.expectedAcquiredAt) { + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + if (!existing.pendingRepairId) { + existing.pendingRepairId = crypto.randomUUID() + } + return { + prepared: true as const, + repairId: existing.pendingRepairId, + token: input.token, + holder: existing.holder, + acquiredAt: existing.acquiredAt, + } + }, + async finalizeWriteLeaseRepair(input: { + token: string + repairId: string + expectedAcquiredAt: string + }) { + const existing = deletion.leases.get(input.token) + if (!existing) return { finalized: true } + if ( + existing.authority !== 'do' || + existing.acquiredAt !== input.expectedAcquiredAt || + existing.pendingRepairId !== input.repairId + ) { + throw new Error( + 'Active account write lease did not match repair request.', + ) + } + deletion.leases.delete(input.token) + return { finalized: true } + }, async readDeletionState() { return { deletingAt: deletion.deletingAt } }, @@ -511,46 +672,15 @@ export function createInMemoryUserMeterEnv() { startAfter?: string | null } = {}, ) { - const pageSize = - typeof input.pageSize === 'number' && Number.isFinite(input.pageSize) - ? Math.min(Math.max(Math.trunc(input.pageSize), 1), 500) - : 100 - const all = listWriteLeasesSorted() - let startIndex = 0 - if ( - typeof input.startAfter === 'string' && - input.startAfter.length > 0 - ) { - try { - const parsed = JSON.parse(input.startAfter) as unknown - if ( - Array.isArray(parsed) && - parsed.length === 2 && - typeof parsed[0] === 'string' && - typeof parsed[1] === 'string' - ) { - const acquiredAt = parsed[0] - const token = parsed[1] - startIndex = - all.findIndex( - (row) => row.acquiredAt === acquiredAt && row.token === token, - ) + 1 - } - } catch { - startIndex = 0 - } - } - const page = all.slice(startIndex, startIndex + pageSize) - const truncated = startIndex + pageSize < all.length - const last = page[page.length - 1] - return { - leases: page, - nextStartAfter: - truncated && last - ? JSON.stringify([last.acquiredAt, last.token]) - : null, - truncated, - } + return pageWriteLeases(input) + }, + async listDoAuthorityWriteLeases( + input: { + pageSize?: number + startAfter?: string | null + } = {}, + ) { + return pageWriteLeases(input, 'do') }, async countActiveWriteLeases() { return { count: deletion.leases.size } @@ -561,6 +691,7 @@ export function createInMemoryUserMeterEnv() { token: string holder: string acquiredAt: string + authority?: string }> }) { let leasesApplied = 0 @@ -570,7 +701,15 @@ export function createInMemoryUserMeterEnv() { leasesSkipped += 1 continue } - const result = await this.shadowAcquireWriteLease(lease) + const authority = + typeof lease.authority === 'string' && + isUserMeterWriteLeaseAuthority(lease.authority) + ? lease.authority + : 'legacy' + const result = + authority === 'do' + ? await this.acquireWriteLease(lease) + : await this.shadowAcquireWriteLease(lease) if (result.acquired) leasesApplied += 1 else leasesSkipped += 1 } From 6c2effd3bc7487ebacee061676a2a452f136bf72 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 10:45:56 +0000 Subject: [PATCH 24/39] docs(account): describe UserMeter lease authority Co-authored-by: Kent C. Dodds --- .../account-write-lease-repair.md | 17 ++- .../contributing/architecture/data-storage.md | 35 +++-- .../contributing/architecture/entitlements.md | 132 +++++++++--------- packages/worker/src/account/export.ts | 6 +- .../worker/src/account/user-owned-surfaces.ts | 2 +- 5 files changed, 102 insertions(+), 90 deletions(-) diff --git a/docs/contributing/account-write-lease-repair.md b/docs/contributing/account-write-lease-repair.md index ed8f9ea0f5..e2308c28e9 100644 --- a/docs/contributing/account-write-lease-repair.md +++ b/docs/contributing/account-write-lease-repair.md @@ -4,6 +4,10 @@ Account write leases do not expire. A crashed writer intentionally blocks account deletion until an administrator verifies that the process is gone and releases its exact token with an audit reason. +Active leases may live in D1 (email / no-`env` writers) and/or UserMeter +(`authority='do'` when callers supply `USER_METER`). Admin list unions both, +deduped by token. + First inspect active leases: ```javascript @@ -24,7 +28,12 @@ await kody.admin_account_write_lease_repair({ }) ``` -The repair writes `account_write_lease_repairs` before releasing the token. -Wrong-user, stale timestamp, already-released token, or short reason requests -fail closed. Retry account deletion only after inspection shows no active -leases. +The repair writes `account_write_lease_repairs` before releasing the token. D1 +leases use an atomic audit-before-delete batch. DO-authority leases use prepare +(stable `repairId`, lease stays held) → audit insert/verify → finalize DO +deletion → clear the temporary D1 rollout mirror. The D1 mirror is never cleared +while the DO lease remains held. Retries after a lost finalize response succeed +when the matching audit exists and the DO lease is already gone (idempotently +clearing any stale D1 mirror). Wrong user, stale timestamp, or short reason +requests fail closed. Retry account deletion only after inspection shows no +active leases. diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 1f2eb593bd..61663731bb 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -563,7 +563,7 @@ Naming matches `RunLog` and `JobManager`: one object per untrimmed stable MCP column inside the DO because the object identity is the user. SQLite ownership (schema version tracked in `user_meter_meta`; current version -**6**): +**7**): - `daily_counters` — authoritative UTC-day counters for `email_sends_per_day`, `email_receives_per_day`, `execute_calls_per_day`, and @@ -590,29 +590,28 @@ SQLite ownership (schema version tracked in `user_meter_meta`; current version (`listPackageServiceStates`, `countRunningPackageServices`, `bootstrapPackageServiceStates`) mirror D1 semantics for future parity review only. -- `deletion_state` / `account_write_leases` — **shadow** of D1 - `users.deleting_at` and active `account_write_leases` (singleton - `deleting_at`; lease rows `token` / `holder` / `acquired_at`). Added in schema - v6. Populated by optional best-effort dual-writes from - `account/deletion-state.ts` after successful D1 mark/acquire/release/repair; - never read for fencing, list, repair, or drain in expand-phase slice 5 Phase - A. On mark, after D1 `deleting_at` is set, the helper loads active D1 leases - and calls `shadowReplaceDeletionState` (DO-serialized tombstone set/preserve + - exact lease-set replace) so stale unreleased shadow rows are cleared when D1 - drain reaches zero and active D1 leases are preserved. `purge()` preserves an - existing deleting tombstone across `deleteAll` so a later cutover cannot - reopen writes. Cutover-support RPCs (`readDeletionState`, `listWriteLeases`, - `countActiveWriteLeases`, `bootstrapDeletionState`) exist for future parity - review only; account export emits a sanitized `deletionShadow` without raw - token/holder. +- `deletion_state` / `account_write_leases` — deletion tombstone plus write + leases (singleton `deleting_at`; lease rows `token` / `holder` / `acquired_at` + / `authority` (`do`|`legacy`) / `pending_repair_id`). Added as shadow in + schema v6; schema v7 is the authority cutover. When callers supply + `USER_METER`, `authority='do'` rows are authoritative for acquire/held/release + and admin union list; `authority='legacy'` rows are the D1 lease snapshot + replaced by `markDeleting`. During rolling deploy, DO-authority acquires also + mirror the same token into D1 `account_write_leases` so Phase-A marks that + only read D1 still observe them (temporary; not the final hot path). D1 + `users.deleting_at` remains the permanent point gate; email paths omit `env` + and keep exact D1 leases. `purge()` preserves an existing deleting tombstone + across `deleteAll`. Account export emits a sanitized `deletionShadow` without + raw token/holder. Retention is self-enforced inside the DO: every read/write path opportunistically deletes counter and claim rows older than seven UTC days (`userMeterDailyCounterRetentionDays`). Enforcement only needs the current day; the window covers timezone edge cases, recent account exports, and inbound retries. Shadow storage-byte and package-service liveness rows are not -time-pruned. Deletion-fence lease shadows are bounded by the D1-backed mark -replace path above rather than time retention. +time-pruned. Deletion-fence legacy lease rows are bounded by the D1 snapshot +replace on `markDeleting` rather than time retention; DO-authority rows clear on +release/repair/purge. **Expand-phase D1 mirrors (daily counters only):** enforcement and point reads are authoritative in UserMeter for daily counters. D1 diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 5ea3edf7eb..09cdd9a251 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -130,11 +130,12 @@ v5 adds an optional per-service shadow table as **best-effort future-cutover support only** — see [Package service liveness — UserMeter shadow](#package-service-liveness--usermeter-shadow-expand-phase-slice-4-phase-a). -**D1 account-deletion write fencing** (`users.deleting_at` / -`account_write_leases`) stays authoritative for fencing, list, repair, and -drain. UserMeter schema v6 adds an optional deletion-fence / write-lease shadow -as **best-effort future-cutover support only** — see -[Account-deletion write fencing — UserMeter shadow](#account-deletion-write-fencing--usermeter-shadow-expand-phase-slice-5-phase-a). +**Account-deletion write fencing** uses a split authority model after slice 5 +Phase B: D1 `users.deleting_at` remains the permanent point gate; callers that +supply `USER_METER` treat UserMeter (`authority='do'`) as authoritative for +lease acquire/held/release/drain count, while email/transition callers that omit +`env` keep the exact D1 `account_write_leases` path. See +[Account-deletion write fencing — UserMeter authority](#account-deletion-write-fencing--usermeter-authority-expand-phase-slice-5-phase-b). StorageRunner bucket `estimatedBytes` and the per-bucket inventory in `user_storage_buckets` stay a **separate** quota component. StorageRunner write @@ -282,60 +283,63 @@ remains on D1. **Account purge:** `UserMeter.purge()` clears package-service shadow rows with the rest of DO state via `deleteAll`. -### Account-deletion write fencing — UserMeter shadow (expand phase slice 5, Phase A) - -D1 `users.deleting_at`, `account_write_leases`, and -`account_write_lease_repairs` remain the **sole authority** for deletion -fencing, lease acquire/release, active-lease list, audited repair, and deletion -drain waits in this PR. Nothing in Phase A switches `withAccountWriteLease`, -`markAccountDeleting`, `listActiveAccountWriteLeases`, or -`repairAccountWriteLease` onto UserMeter reads — public errors, ALS nested-lease -reuse, and D1 `waitUntil` release semantics stay identical when optional `env` -is omitted. - -UserMeter schema **v6** adds a singleton `deletion_state.deleting_at` tombstone -plus per-token `account_write_leases` rows (`token`, `holder`, `acquired_at`) as -**best-effort shadow / future-cutover support only**. User scope is the DO -identity — there is no `user_id` column. Shadow rows are never read for fencing, -list, repair, or drain in this slice. - -**Optional dual-write from `account/deletion-state.ts`:** after a **successful** -D1 mark/acquire/release/repair, callers may pass optional `env?: UserMeterEnv` -(and `waitUntil` when available). The helper schedules a best-effort UserMeter -shadow (`shadowReplaceDeletionState` on mark, `shadowAcquireWriteLease`, -`shadowReleaseWriteLease`). After D1 `deleting_at` is set, mark loads the -authoritative active D1 lease rows and calls `shadowReplaceDeletionState`, which -under DO serialization sets/preserves the tombstone then deletes and re-inserts -exactly that lease set — so a prior failed shadow release cannot leave stale -rows, and a deletion retry with zero D1 leases clears the shadow set. Acquire -shadows start during the write; ordered release shadows run after authoritative -D1 release. When `waitUntil` is provided, shadows detach through it; when -omitted, mark/release/repair shadows are awaited (still non-rejecting) so -middleware and other non-`waitUntil` callers cannot return with a stale shadow -lease or missing tombstone. Missing `USER_METER` is a no-op; failures log -`account-deletion-user-meter-shadow-failed` and never alter D1 results or -errors. Mark is monotonic (first `deleting_at` wins). Acquire is rejected when a -deleting tombstone is present and is idempotent for an already-held token. -Cutover/list RPCs (`listWriteLeases`, `bootstrapDeletionState`) may still carry -token/holder for parity review. - -**Account export:** `UserMeter.exportCounters` returns additive `deletionShadow` -on the first page only (`startAfter` absent); later pages return `null`. The -export shape is sanitized (`deletingAt`, `activeWriteLeaseCount`, and -`writeLeases` with `acquiredAt` only — no raw token or holder). Section totals -count the tombstone (when set) plus `activeWriteLeaseCount`; the field is -explicitly non-authoritative — authoritative fencing remains on D1. - -**Account purge:** `UserMeter.purge()` clears counters, claims, storage shadow, -package-service shadow, and write-lease shadow via `deleteAll`, then restores -any pre-existing deleting tombstone so a later authority cutover cannot reopen -writes for a purged account. - -Non-email production call sites that hold `env` pass it into deletion-state APIs -(`withAccountWriteLease`, `markAccountDeleting`, `repairAccountWriteLease`) and -forward existing `waitUntil` when available. Email paths stay untouched in this -slice (`system:email` skip unchanged). Assert-only readers -(`assertAccountWritable*`, `listActiveAccountWriteLeases`) remain D1-only. +### Account-deletion write fencing — UserMeter authority (expand phase slice 5, Phase B) + +UserMeter schema **v7** stores `account_write_leases.authority` (`do` | +`legacy`; rows created under v6 default to `legacy`) and `pending_repair_id` for +audit-safe DO repair. + +**Split authority:** + +- D1 `users.deleting_at` remains the **permanent point gate** (auth projection / + purge failures fail closed). Assert-only readers (`assertAccountWritable*`) + stay D1-only. D1 `account_write_leases` / `account_write_lease_repairs` tables + are **not** retired. +- When `env.USER_METER` is supplied (all non-email Phase-A-wired callers), + UserMeter is **authoritative** for lease acquire / held / release via + `acquireWriteLease` / `assertWriteLeaseHeld` / `releaseWriteLease` + (`authority='do'`). Missing binding or DO failures **fail closed**. +- When `env` is omitted (email transition), retain the **exact D1 lease path** + (including `active_write_count`, ALS nested reuse, and `waitUntil` release). + +**Temporary rolling-version D1 mirror (not the final hot path):** while older +isolates may still run Phase-A mark (D1-only lease counts), every DO-authority +acquire also inserts the **same** `token` / `holder` / `acquiredAt` into D1 +`account_write_leases` and bumps `active_write_count` with the existing atomic +batch + lost-response reconcile. If that D1 mirror cannot be confirmed, the DO +lease is released and acquire fails closed. Authoritative release clears the DO +lease first, then the D1 mirror/`active_write_count` (same `waitUntil` +detachment semantics) so a stale D1 row may overblock old marks but never +underblock. After old isolates drain, remove this mirror (and later the D1 lease +inventory) from the hot path. + +**`markAccountDeleting`:** always `COALESCE`s D1 `deleting_at` first, then loads +live D1 leases (including temporary DO mirrors). With `env`, calls authoritative +`markDeleting` which sets/preserves the DO tombstone, replaces **only legacy** +rows with that exact D1 snapshot, preserves DO-authority rows, and returns the +deduped-by-token union count used for drain waits. Without `env`, returns the D1 +lease count unchanged (so Phase-A marks observe DO acquires via the mirror). + +**Admin list / repair:** `listActiveAccountWriteLeases(db, userId, env?)` with +`env` unions live D1 leases + DO-authority leases (dedupe by token, same +`acquired_at, token` order); without `env` exact D1 behavior. DO repair is +audit-first and idempotent: prepare (stable `repairId`, lease stays held) → +insert/verify D1 audit → finalize exact pending DO repair → then idempotently +clear matching D1 mirror/count. The D1 mirror is never cleared while the DO +lease remains held. Finalize failure without commit leaves DO + D1 mirror held +and fails closed; retry resumes the pending repair. Retry after finalize commit +/ lost response returns success when the matching audit exists and the DO lease +is absent (clears any stale D1 mirror; never falls through to a mismatched D1 +atomic batch). Post-write held checks treat pending repair as held until +finalize, then surface `AccountWriteLeaseLostError`. Pure D1 leases keep the +existing atomic audit-before-release batch. + +**Account export / purge:** first-page sanitized `deletionShadow` still omits +raw token/holder. `purge()` clears leases/counters/shadows via `deleteAll` then +restores any deleting tombstone; D1 `deleting_at` remains the gate. + +Public errors, ALS nested-lease reuse, holder strings, and `waitUntil` release +detachment stay parity-compatible with the pre-cutover surface. ### Future package-service authority flip (contract follow-up) @@ -361,11 +365,11 @@ Only then do reads, reserves, and reconciliation switch to UserMeter-first with D1 as mirror/cursor. Until that flip, shadow divergence is acceptable; D1 remains the contract. -**Remaining UserMeter expand roadmap:** slice 5 Phase A — account-deletion write -fencing shadow (this PR; D1 authority unchanged; non-email dual-write wired). -Package-service and storage authority flips are separate high-risk contract -follow-ups after soak/parity review. Deletion-fence authority flip is a separate -high-risk contract follow-up after shadow/D1 parity review. +**Current UserMeter cutover:** slice 5 Phase B moves non-email write leases into +UserMeter authority while retaining the D1 `deleting_at` gate and temporary +same-token rollout mirrors; email keeps its D1 lease path. Package-service and +storage authority flips remain separate high-risk contract follow-ups after +soak/parity review. **Daily-counter mirror retirement:** dropping D1 `entitlement_daily_counters` waits until reporting-off-D1 work merges and mirror parity is verified in diff --git a/packages/worker/src/account/export.ts b/packages/worker/src/account/export.ts index c1136ed088..7e6f9176b7 100644 --- a/packages/worker/src/account/export.ts +++ b/packages/worker/src/account/export.ts @@ -312,9 +312,9 @@ export type AccountExportSectionResult = { */ packageServiceStatesShadow?: UserMeterExportResult['packageServiceStatesShadow'] /** - * Non-authoritative UserMeter deletion-fence / write-lease shadow. Present - * only on the first `user_meter` page (`startAfter` absent); later pages - * set it to `null`. + * Sanitized UserMeter deletion-fence / write-lease inventory (no raw token + * or holder). Present only on the first `user_meter` page (`startAfter` + * absent); later pages set it to `null`. */ deletionShadow?: UserMeterExportResult['deletionShadow'] } diff --git a/packages/worker/src/account/user-owned-surfaces.ts b/packages/worker/src/account/user-owned-surfaces.ts index b939ea036d..420e96ca1b 100644 --- a/packages/worker/src/account/user-owned-surfaces.ts +++ b/packages/worker/src/account/user-owned-surfaces.ts @@ -132,7 +132,7 @@ export const accountUserOwnedDurableObjectSurfaces: ReadonlyArray Date: Sat, 1 Aug 2026 10:56:46 +0000 Subject: [PATCH 25/39] test(account): provide authoritative meter fixtures Co-authored-by: Kent C. Dodds --- .../worker/src/community/avatar.node.test.ts | 93 ++++++++---- .../src/community/community-icon.node.test.ts | 93 +++++++----- packages/worker/src/jobs/service.node.test.ts | 134 +++++++++++------- .../src/mcp/memory/service.node.test.ts | 51 ++++--- .../src/mcp/run-kody-registry.node.test.ts | 32 ++++- .../src/package-registry/service.node.test.ts | 50 ++++--- .../incidental-runtime-warnings.ts | 1 + .../worker/src/test-support/user-meter.ts | 41 ++++++ 8 files changed, 343 insertions(+), 152 deletions(-) diff --git a/packages/worker/src/community/avatar.node.test.ts b/packages/worker/src/community/avatar.node.test.ts index e5f80d5033..036862fe48 100644 --- a/packages/worker/src/community/avatar.node.test.ts +++ b/packages/worker/src/community/avatar.node.test.ts @@ -7,6 +7,63 @@ import { saveUserAvatar, } from './avatar.ts' import { AccountDeletionInProgressError } from '#worker/account/deletion-state.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' + +function createAvatarTestEnv(input: { + db: D1Database + communityAssets: R2Bucket + meter?: ReturnType +}) { + const meter = input.meter ?? createInMemoryUserMeterEnv() + return { + APP_DB: input.db, + COMMUNITY_ASSETS: input.communityAssets, + USER_METER: meter.env.USER_METER, + } as Pick +} + +function createAvatarDeletionRaceDbMock() { + let deleting = false + const db = { + prepare(query: string) { + return { + bind() { + return { + async first() { + if (query.includes('SELECT deleting_at')) { + return { + deleting_at: deleting ? '2026-07-22 22:00:00' : null, + } as T + } + if (query.includes('SELECT avatar_key')) { + return { avatar_key: null } as T + } + return null + }, + async run() { + if ( + query.includes('UPDATE users') && + query.includes('avatar_key') + ) { + return { meta: { changes: deleting ? 0 : 1 } } + } + return { meta: { changes: 1 } } + }, + } + }, + } + }, + async batch() { + return [{ meta: { changes: 1 } }, { meta: { changes: 1 } }] + }, + } as unknown as D1Database + return { + db, + setDeleting(value: boolean) { + deleting = value + }, + } +} function createPngHeader(width: number, height: number) { const bytes = new Uint8Array(24) @@ -164,7 +221,6 @@ test('getUserAvatarObject refuses keys outside the user-avatars prefix', async ( }) test('saveUserAvatar removes an in-flight upload when deletion starts', async () => { - let deleting = false let releasePut: () => void = () => undefined let markPutStarted: () => void = () => undefined const putStarted = new Promise((resolve) => { @@ -174,34 +230,11 @@ test('saveUserAvatar removes an in-flight upload when deletion starts', async () releasePut = resolve }) const deleted: Array = [] - const db = { - prepare(query: string) { - return { - bind() { - return { - async first() { - if (query.includes('SELECT deleting_at')) { - return { - deleting_at: deleting ? '2026-07-22 22:00:00' : null, - } as T - } - if (query.includes('SELECT avatar_key')) { - return { avatar_key: null } as T - } - return null - }, - async run() { - return { meta: { changes: deleting ? 0 : 1 } } - }, - } - }, - } - }, - } as unknown as D1Database + const { db, setDeleting } = createAvatarDeletionRaceDbMock() const save = saveUserAvatar({ - env: { - APP_DB: db, - COMMUNITY_ASSETS: { + env: createAvatarTestEnv({ + db, + communityAssets: { async put() { markPutStarted() await putReleased @@ -211,14 +244,14 @@ test('saveUserAvatar removes an in-flight upload when deletion starts', async () deleted.push(key) }, } as R2Bucket, - }, + }), numericUserId: 1, stableUserId: 'stable-1', bytes: createPngHeader(128, 128), contentType: 'image/png', }) await putStarted - deleting = true + setDeleting(true) releasePut() await expect(save).rejects.toBeInstanceOf(AccountDeletionInProgressError) expect(deleted).toHaveLength(1) diff --git a/packages/worker/src/community/community-icon.node.test.ts b/packages/worker/src/community/community-icon.node.test.ts index 98322b2471..b03b3b6640 100644 --- a/packages/worker/src/community/community-icon.node.test.ts +++ b/packages/worker/src/community/community-icon.node.test.ts @@ -12,6 +12,7 @@ import { renderCommunityIconFallbackPng, } from './community-icon.ts' import { type CommunityListingRecord } from './types.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' const mocks = vi.hoisted(() => ({ readFirstArtifactFileAtCommit: vi.fn(), @@ -153,6 +154,63 @@ function createFakeR2() { return { bucket, values } } +function createCommunityIconTestEnv(input: { + db: D1Database + kv: KVNamespace + bucket: R2Bucket + meter?: ReturnType +}) { + const meter = input.meter ?? createInMemoryUserMeterEnv() + return { + APP_DB: input.db, + BUNDLE_ARTIFACTS_KV: input.kv, + COMMUNITY_ASSETS: input.bucket, + USER_METER: meter.env.USER_METER, + } as Env +} + +function createCommunityIconDeletionRaceDbMock() { + let deleting = false + let leaseAcquires = 0 + const db = { + prepare(query: string) { + const normalized = query.replace(/\s+/g, ' ').trim() + return { + bind() { + return { + async first() { + if (normalized.includes('SELECT deleting_at')) { + return { deleting_at: deleting ? 'now' : null } as T + } + return null + }, + async run() { + if ( + normalized.includes('INSERT INTO account_write_leases') || + normalized.includes( + 'active_write_count = active_write_count + 1', + ) + ) { + leaseAcquires += 1 + if (leaseAcquires === 2) deleting = true + return { meta: { changes: deleting ? 0 : 1 } } + } + return { meta: { changes: 1 } } + }, + } + }, + } + }, + async batch() { + leaseAcquires += 1 + if (leaseAcquires === 2) deleting = true + const changes = deleting ? 0 : 1 + return [{ meta: { changes } }, { meta: { changes } }] + }, + } as unknown as D1Database + return { db, getLeaseAcquires: () => leaseAcquires } +} + const listing = { id: 'listing-1', ownerUserId: 'owner-1', @@ -358,32 +416,7 @@ test('community icon cache write loses the race to account deletion', async () = const png = createPngHeader(128, 128) const { kv, values: kvValues } = createFakeKv() const { bucket } = createFakeR2() - let deleting = false - let leaseAcquires = 0 - const db = { - prepare(query: string) { - return { - bind() { - return { - async first() { - if (query.includes('SELECT deleting_at')) { - return { deleting_at: deleting ? 'now' : null } as T - } - return null - }, - async run() { - if (query.includes('active_write_count + 1')) { - leaseAcquires += 1 - if (leaseAcquires === 2) deleting = true - return { meta: { changes: deleting ? 0 : 1 } } - } - return { meta: { changes: 1 } } - }, - } - }, - } - }, - } as unknown as D1Database + const { db, getLeaseAcquires } = createCommunityIconDeletionRaceDbMock() mocks.readCommunitySnapshot.mockResolvedValue({ version: 1, listingId: listing.id, @@ -399,15 +432,11 @@ test('community icon cache write loses the race to account deletion', async () = bytes: png, }) await getCommunityIconObject({ - env: { - APP_DB: db, - BUNDLE_ARTIFACTS_KV: kv, - COMMUNITY_ASSETS: bucket, - } as Env, + env: createCommunityIconTestEnv({ db, kv, bucket }), listing, iconCommit: listing.pinnedCommit, }) - expect(leaseAcquires).toBe(2) + expect(getLeaseAcquires()).toBe(2) expect(kvValues.size).toBe(0) }) diff --git a/packages/worker/src/jobs/service.node.test.ts b/packages/worker/src/jobs/service.node.test.ts index e3e9f6b30e..cbb0b857d8 100644 --- a/packages/worker/src/jobs/service.node.test.ts +++ b/packages/worker/src/jobs/service.node.test.ts @@ -3,6 +3,10 @@ import { McpCallerError } from '#mcp/caller-error.ts' import { createMcpCallerContext } from '#mcp/context.ts' import { repoBackedModuleEntrypointExportErrorMessage } from '#worker/repo/repo-kody-execution.ts' import { silenceIncidentalRuntimeWarnings } from '#worker/test-support/incidental-runtime-warnings.ts' +import { + createInMemoryUserMeterEnv, + createPermissiveAccountWriteLeaseDbHooks, +} from '#worker/test-support/user-meter.ts' import { isEntitlementLimitError } from '#worker/entitlements/errors.ts' import { planLimits } from '#worker/entitlements/plans.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' @@ -354,6 +358,7 @@ function createDatabase( ['jobs', (initialRows.jobs ?? []).map((row) => ({ ...row }))], ['users', (initialRows.users ?? []).map((row) => ({ ...row }))], ]) + const writeLeaseDb = createPermissiveAccountWriteLeaseDbHooks() const clone = (value: T): T => structuredClone(value) @@ -402,11 +407,20 @@ function createDatabase( } return { + async batch(statements: Array<{ run: () => Promise }>) { + return await writeLeaseDb.runWriteLeaseBatch(statements) + }, prepare(query: string) { return { bind(...params: Array) { return { async first>() { + if (writeLeaseDb.supportsDeletingAtQuery(query)) { + return writeLeaseDb.deletingAtFirstResult() as T + } + if (writeLeaseDb.supportsHeldLeaseQuery(query)) { + return writeLeaseDb.heldLeaseFirstResult() as T + } if (query.includes('SELECT plan, stripe_plan FROM users')) { return selectOne( 'users', @@ -685,6 +699,9 @@ function createDatabase( throw new Error(`Unsupported all query: ${query}`) }, async run() { + if (writeLeaseDb.supportsWriteLeaseRun(query)) { + return writeLeaseDb.writeLeaseRunResult() + } if (query.startsWith('UPDATE users')) { return { meta: { changes: 1, last_row_id: 0 } } } @@ -1098,6 +1115,19 @@ function createDatabase( } as unknown as D1Database } +function createJobServiceTestEnv( + bindings: Record & { + APP_DB: ReturnType + }, + meter?: ReturnType, +) { + const userMeter = meter ?? createInMemoryUserMeterEnv() + return { + ...bindings, + USER_METER: userMeter.env.USER_METER, + } as Env +} + function createBundleArtifactsKv() { const store = new Map() return { @@ -1206,9 +1236,9 @@ function createBaseCallerContext(): PersistedJobCallerContext { } test('package job sync reports scheduler changes for add, update, and remove only', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) const input = { env, userId: 'user-1', @@ -1292,12 +1322,12 @@ test('package job sync reports scheduler changes for add, update, and remove onl }) test('create, update, and delete jobs sync the job manager alarm', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: createBundleArtifactsKv(), - } as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() @@ -1436,11 +1466,11 @@ function buildEntitlementTestJobBody(index: number): JobCreateInput { test('createJob enforces scheduled job entitlements for plan users and denies at the max plan ceiling', async () => { const plannedEmail = 'planned@example.com' const plannedUserId = await createStableUserIdFromEmail(plannedEmail) - const plannedEnv = { + const plannedEnv = createJobServiceTestEnv({ APP_DB: createDatabase({ users: [{ email: plannedEmail, plan: 'free' }], }), - } as Env + }) mockRepoPersistence() const plannedCallerContext = createPlanUserCallerContext({ userId: plannedUserId, @@ -1478,7 +1508,7 @@ test('createJob enforces scheduled job entitlements for plan users and denies at const maxEmail = 'max@example.com' const maxUserId = await createStableUserIdFromEmail(maxEmail) const maxLimit = planLimits.max.maxScheduledJobs - const belowMaxEnv = { + const belowMaxEnv = createJobServiceTestEnv({ APP_DB: createDatabase({ users: [{ email: maxEmail, plan: 'max' }], jobs: Array.from( @@ -1489,7 +1519,7 @@ test('createJob enforces scheduled job entitlements for plan users and denies at }), ), }), - } as Env + }) mockRepoPersistence() const maxCallerContext = createPlanUserCallerContext({ userId: maxUserId, @@ -1501,7 +1531,7 @@ test('createJob enforces scheduled job entitlements for plan users and denies at body: buildEntitlementTestJobBody(0), }) - const atCeilingEnv = { + const atCeilingEnv = createJobServiceTestEnv({ APP_DB: createDatabase({ users: [{ email: maxEmail, plan: 'max' }], jobs: Array.from({ length: maxLimit }, (_, index) => ({ @@ -1509,7 +1539,7 @@ test('createJob enforces scheduled job entitlements for plan users and denies at user_id: maxUserId, })), }), - } as Env + }) const maxError = await createJob({ env: atCeilingEnv, callerContext: maxCallerContext, @@ -1531,9 +1561,9 @@ test('createJob enforces scheduled job entitlements for plan users and denies at }) test('updateJob and deleteJob reject another user trying to mutate or remove a job by id', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) mockRepoPersistence() const ownerCallerContext = createBaseCallerContext() const created = await createJob({ @@ -1600,9 +1630,9 @@ test('updateJob and deleteJob reject another user trying to mutate or remove a j }) test('missing job ids throw McpCallerError from get/inspect/run-now', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) const missingJobId = 'missing-job-id' const userId = createBaseCallerContext().user.userId @@ -1638,9 +1668,9 @@ test('missing job ids throw McpCallerError from get/inspect/run-now', async () = }) test('updateJob clears params, updates timezone, and disables a job', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() const created = await createJob({ @@ -1690,9 +1720,9 @@ test('updateJob clears params, updates timezone, and disables a job', async () = }) test('createJob and updateJob reject modules without a default export as McpCallerError', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() const bareArrowSnippet = 'async () => ({ ok: true })' @@ -1744,9 +1774,9 @@ test('createJob and updateJob reject modules without a default export as McpCall }) test('inspectJobsForUser returns persisted job fields with alarm debug state', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), - } as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() const created = await createJob({ @@ -1853,10 +1883,10 @@ test('inspectJobsForUser returns persisted job fields with alarm debug state', a }) test('getJobInspection reports alarm state, source code, and artifact gaps', async () => { - const env = { + const env = createJobServiceTestEnv({ APP_DB: createDatabase(), BUNDLE_ARTIFACTS_KV: createBundleArtifactsKv(), - } as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() const created = await createJob({ @@ -1990,10 +2020,10 @@ test('getJobInspection reports alarm state, source code, and artifact gaps', asy }) const bundleKv = createBundleArtifactsKv() - const manifestEnv = { + const manifestEnv = createJobServiceTestEnv({ APP_DB: createDatabase(), BUNDLE_ARTIFACTS_KV: bundleKv, - } as Env + }) const missingManifestJob = await createJob({ env: manifestEnv, callerContext, @@ -2054,7 +2084,7 @@ test('executeJobOnce binds writable storage and overrides persisted interactive // Usage rollup writes are best-effort and fail against this fake env. silenceIncidentalRuntimeWarnings() const db = createDatabase() - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', @@ -2095,7 +2125,7 @@ test('executeJobOnce binds writable storage and overrides persisted interactive } }, }, - } as unknown as Env + }) mockRepoPersistence() const callerContext = { ...createBaseCallerContext(), @@ -2250,7 +2280,7 @@ test('executeJobOnce binds writable storage and overrides persisted interactive test('executeJobOnce runs repo-backed one-off jobs from kody.json manifests', async () => { silenceIncidentalRuntimeWarnings() const db = createDatabase() - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', @@ -2291,7 +2321,7 @@ test('executeJobOnce runs repo-backed one-off jobs from kody.json manifests', as } }, }, - } as unknown as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() @@ -2462,7 +2492,7 @@ test('executeJobOnce preserves kody secret and value semantics', async () => { silenceIncidentalRuntimeWarnings() const db = createDatabase() const bundleKv = createBundleArtifactsKv() - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', @@ -2470,7 +2500,7 @@ test('executeJobOnce preserves kody secret and value semantics', async () => { COOKIE_SECRET: 'test-secret-0123456789abcdef0123456789', LOADER: {} as WorkerLoader, REPO_SESSION: {} as DurableObjectNamespace, - } as unknown as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() @@ -2663,13 +2693,13 @@ test('executeJobOnce refreshes repo sessions when base commit moves', async () = BUNDLE_ARTIFACTS_KV: bundleKv, } as Env, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -2882,13 +2912,13 @@ test('executeJobOnce rebuilds stale published job bundles after the source commi 'export default async () => { console.log("canary"); return { version: "new" } }', }, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -2972,13 +3002,13 @@ test('executeJobOnce executes package-backed jobs from published artifacts', asy 'src/custom-job.ts': 'export default async () => ({ ok: true })', }, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3104,13 +3134,13 @@ test('executeJobOnce bypasses typecheck-only failures when the stored repo polic BUNDLE_ARTIFACTS_KV: bundleKv, } as Env, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3283,13 +3313,13 @@ test('executeJobOnce succeeds for repo-backed jobs with repo-session absolute pa BUNDLE_ARTIFACTS_KV: bundleKv, } as Env, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3451,13 +3481,13 @@ test('executeJobOnce fails instead of reusing a stale repo session when discard publishedCommit: 'commit-1', manifestPath: 'package.json', }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3572,13 +3602,13 @@ test('executeJobOnce bundles and runs ESM repo-backed job entrypoints', async () 'src/lib.ts': 'export const value = 1', }, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3773,13 +3803,13 @@ test('executeJobOnce returns an error when kody secret policy would reject execu BUNDLE_ARTIFACTS_KV: bundleKv, } as Env, }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() const job: JobRecord = { version: 1, @@ -3900,14 +3930,14 @@ test('runJobNow retains once jobs for retention cleanup instead of deleting them silenceIncidentalRuntimeWarnings() const db = createDatabase() const bundleKv = createBundleArtifactsKv() - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: bundleKv, LOADER: {} as WorkerLoader, REPO_SESSION: {} as DurableObjectNamespace, - } as Env & { CAPABILITY_VECTOR_INDEX?: Pick } + }) as Env & { CAPABILITY_VECTOR_INDEX?: Pick } mockRepoPersistence() const callerContext = createBaseCallerContext() const jobView = await createJob({ @@ -4069,13 +4099,13 @@ test('runJobNow can use a one-off repo check policy override without changing th publishedCommit: 'commit-run-now-override', manifestPath: 'package.json', }) - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', BUNDLE_ARTIFACTS_KV: createBundleArtifactsKv(), LOADER: {} as WorkerLoader, - } as Env + }) const callerContext = createBaseCallerContext() mockRepoPersistence() const jobView = await createJob({ @@ -4245,7 +4275,7 @@ test('executeJobOnce records job_run usage for success and failure without chang .spyOn(usageModule, 'recordUsage') .mockResolvedValue(undefined) const db = createDatabase() - const env = { + const env = createJobServiceTestEnv({ APP_DB: db, CLOUDFLARE_ACCOUNT_ID: 'acct-test', CLOUDFLARE_API_TOKEN: 'token-test', @@ -4286,7 +4316,7 @@ test('executeJobOnce records job_run usage for success and failure without chang } }, }, - } as unknown as Env + }) mockRepoPersistence() const callerContext = createBaseCallerContext() const jobView = await createJob({ diff --git a/packages/worker/src/mcp/memory/service.node.test.ts b/packages/worker/src/mcp/memory/service.node.test.ts index c785a064aa..bfbd62ef24 100644 --- a/packages/worker/src/mcp/memory/service.node.test.ts +++ b/packages/worker/src/mcp/memory/service.node.test.ts @@ -3,6 +3,10 @@ import { CAPABILITY_EMBEDDING_DIMENSIONS, deterministicEmbedding, } from '#worker/vectorize/embedding.ts' +import { + createInMemoryUserMeterEnv, + createPermissiveAccountWriteLeaseDbHooks, +} from '#worker/test-support/user-meter.ts' import { acknowledgeSurfacedMemories, deleteMemory, @@ -21,6 +25,7 @@ function createMemoryTestDb() { const memories = new Map() const suppressions = new Map() const batches: Array> = [] + const writeLeaseDb = createPermissiveAccountWriteLeaseDbHooks() function suppressionKey( userId: string, @@ -45,6 +50,12 @@ function createMemoryTestDb() { bind(...params: Array) { return { async first() { + if (writeLeaseDb.supportsDeletingAtQuery(query)) { + return writeLeaseDb.deletingAtFirstResult() as T + } + if (writeLeaseDb.supportsHeldLeaseQuery(query)) { + return writeLeaseDb.heldLeaseFirstResult() as T + } if ( normalizedQuery.startsWith( 'select 1 as held from account_write_leases', @@ -122,16 +133,8 @@ function createMemoryTestDb() { return { results: [] as Array, meta: { changes: 0 } } }, async run() { - if ( - normalizedQuery.startsWith( - 'insert into account_write_leases', - ) || - normalizedQuery.startsWith('update account_write_leases') - ) { - return { meta: { changes: 1 } } - } - if (normalizedQuery.startsWith('update users')) { - return { meta: { changes: 1 } } + if (writeLeaseDb.supportsWriteLeaseRun(query)) { + return writeLeaseDb.writeLeaseRunResult() } if (normalizedQuery.startsWith('insert into mcp_memories')) { const [ @@ -290,10 +293,25 @@ function createMemoryTestDb() { return { db, memories, suppressions, batches } } -const env = (db: D1Database) => +const env = ( + db: D1Database, + bindings: Partial< + Pick + > = {}, + meter = createInMemoryUserMeterEnv(), +) => ({ APP_DB: db, - }) satisfies Pick as Pick + USER_METER: meter.env.USER_METER, + ...bindings, + }) satisfies Pick as Pick< + Env, + | 'APP_DB' + | 'USER_METER' + | 'AI' + | 'CAPABILITY_VECTOR_INDEX' + | 'SENTRY_ENVIRONMENT' + > function createDeterministicAiBinding(): Ai { return { @@ -385,7 +403,7 @@ test('memory service upserts, verifies, and soft deletes', async () => { expect(deleted?.status).toBe('deleted') const loaded = await getMemory({ - env: { APP_DB: testDb.db }, + env: runtimeEnv, userId: 'user-123', memoryId: created.memory.id, }) @@ -577,7 +595,7 @@ test('memory service rejects invalid source uris and tolerates missing stored va } as unknown as McpMemoryRow) const loaded = await getMemory({ - env: { APP_DB: testDb.db }, + env: runtimeEnv, userId: 'user-123', memoryId: 'legacy-memory', }) @@ -642,8 +660,7 @@ test('memory search online queries Vectorize first and hydrates vector hits by i namespace?: string filter?: Record }> = [] - const runtimeEnv = { - APP_DB: testDb.db, + const runtimeEnv = env(testDb.db, { SENTRY_ENVIRONMENT: 'production', AI: createDeterministicAiBinding(), CAPABILITY_VECTOR_INDEX: { @@ -661,7 +678,7 @@ test('memory search online queries Vectorize first and hydrates vector hits by i } }, }, - } as unknown as Pick + }) const result = await searchMemoryRecords({ env: runtimeEnv, diff --git a/packages/worker/src/mcp/run-kody-registry.node.test.ts b/packages/worker/src/mcp/run-kody-registry.node.test.ts index 5aaf1c329d..ecc3c16a9b 100644 --- a/packages/worker/src/mcp/run-kody-registry.node.test.ts +++ b/packages/worker/src/mcp/run-kody-registry.node.test.ts @@ -1,5 +1,9 @@ import { expect, test, vi } from 'vitest' import { silenceIncidentalRuntimeWarnings } from '#worker/test-support/incidental-runtime-warnings.ts' +import { + createInMemoryUserMeterEnv, + createPermissiveAccountWriteLeaseDbHooks, +} from '#worker/test-support/user-meter.ts' import { type getCapabilityRegistryForContext } from '#mcp/capabilities/registry.ts' import { buildCapabilityRegistry } from '#mcp/capabilities/build-capability-registry.ts' import { defineDomainCapability } from '#mcp/capabilities/define-domain-capability.ts' @@ -621,14 +625,24 @@ function createJobMutationDatabase(input: { tables.set(name, remaining) return rows.length - remaining.length } + const writeLeaseDb = createPermissiveAccountWriteLeaseDbHooks() return { + async batch(statements: Array<{ run: () => Promise }>) { + return await writeLeaseDb.runWriteLeaseBatch(statements) + }, prepare(query: string) { const normalized = query.replace(/\s+/g, ' ').trim() return { bind(...params: Array) { return { async first>() { + if (writeLeaseDb.supportsDeletingAtQuery(query)) { + return writeLeaseDb.deletingAtFirstResult() as T + } + if (writeLeaseDb.supportsHeldLeaseQuery(query)) { + return writeLeaseDb.heldLeaseFirstResult() as T + } if ( normalized === 'SELECT * FROM jobs WHERE id = ? AND user_id = ?' ) { @@ -680,6 +694,9 @@ function createJobMutationDatabase(input: { throw new Error(`Unsupported all query: ${query}`) }, async run() { + if (writeLeaseDb.supportsWriteLeaseRun(query)) { + return writeLeaseDb.writeLeaseRunResult() + } if (normalized.startsWith('UPDATE users')) { return { meta: { changes: 1, last_row_id: 0 } } } @@ -822,6 +839,17 @@ function createJobMutationKv() { } as unknown as KVNamespace & { deletedKeys: Array } } +function createRunKodyRegistryTestEnv( + bindings: Record, + meter?: ReturnType, +) { + const userMeter = meter ?? createInMemoryUserMeterEnv() + return { + ...bindings, + USER_METER: userMeter.env.USER_METER, + } as Env +} + function createJobRow( job: JobRecord, callerContext: PersistedJobCallerContext, @@ -976,7 +1004,7 @@ test('buildKodyFns updates and deletes jobs through production-shaped bindings', const repoSessionAccesses: Array = [] const jobManagerNames: Array = [] const jobManagerSyncPayloads: Array<{ userId: string; source?: string }> = [] - const env = { + const env = createRunKodyRegistryTestEnv({ APP_DB: db, SENTRY_ENVIRONMENT: 'production', CLOUDFLARE_ACCOUNT_ID: 'acct-test', @@ -1025,7 +1053,7 @@ test('buildKodyFns updates and deletes jobs through production-shaped bindings', } }, }, - } as unknown as Env + }) const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValue( new Response( JSON.stringify({ diff --git a/packages/worker/src/package-registry/service.node.test.ts b/packages/worker/src/package-registry/service.node.test.ts index 83ba2c744e..02645a6728 100644 --- a/packages/worker/src/package-registry/service.node.test.ts +++ b/packages/worker/src/package-registry/service.node.test.ts @@ -7,6 +7,10 @@ import { import { isEntitlementLimitError } from '#worker/entitlements/errors.ts' import { planLimits } from '#worker/entitlements/plans.ts' import { createStableUserIdFromEmail } from '#worker/user-id.ts' +import { + createInMemoryUserMeterEnv, + createPermissiveAccountWriteLeaseDbHooks, +} from '#worker/test-support/user-meter.ts' function mockPackageServiceNamespace(): DurableObjectNamespace { return { @@ -171,18 +175,24 @@ function createEnv( userId = 'user-1', options?: { storageBuckets?: Array<{ userId: string; storageId: string }> + meter?: ReturnType + users?: Array<{ email: string; plan: string | null }> + savedPackageCount?: number }, ) { // Projection refresh asserts finite storage bytes (default/missing plan → // `max`). Stub DB answers storage SUM queries with 0 so unplanned unit // fixtures keep focusing on job/artifact side effects. + const meter = options?.meter ?? createInMemoryUserMeterEnv() return { APP_DB: createEntitlementsDatabase({ - users: [], + users: options?.users ?? [], userId, storageBuckets: options?.storageBuckets, + savedPackageCount: options?.savedPackageCount, }), PACKAGE_SERVICE_INSTANCE: mockPackageServiceNamespace(), + USER_METER: meter.env.USER_METER, } as Env } @@ -1054,15 +1064,19 @@ function createEntitlementsDatabase(input: { const users = input.users ?? [] const savedPackageCount = input.savedPackageCount ?? 0 const storageBuckets = input.storageBuckets ?? [] + const writeLeaseDb = createPermissiveAccountWriteLeaseDbHooks() return { + async batch(statements: Array<{ run: () => Promise }>) { + return await writeLeaseDb.runWriteLeaseBatch(statements) + }, prepare(query: string) { return { bind(...params: Array) { return { async run() { - if (query.includes('UPDATE users')) { - return { meta: { changes: 1 } } + if (writeLeaseDb.supportsWriteLeaseRun(query)) { + return writeLeaseDb.writeLeaseRunResult() } if ( query.includes( @@ -1091,6 +1105,12 @@ function createEntitlementsDatabase(input: { throw new Error(`Unsupported run query: ${query}`) }, async first() { + if (writeLeaseDb.supportsDeletingAtQuery(query)) { + return writeLeaseDb.deletingAtFirstResult() as T + } + if (writeLeaseDb.supportsHeldLeaseQuery(query)) { + return writeLeaseDb.heldLeaseFirstResult() as T + } if (query.includes('SELECT plan, stripe_plan FROM users')) { const user = users.find((row) => row.email === params[0]) return (user ? { plan: user.plan } : null) as T | null @@ -1154,14 +1174,10 @@ test('refreshSavedPackageProjection enforces the saved packages entitlement on i const userId = await createStableUserIdFromEmail(email) const limit = planLimits.pro.maxSavedPackages if (limit === null) throw new Error('Expected a numeric pro package limit.') - const env = { - APP_DB: createEntitlementsDatabase({ - users: [{ email, plan: 'pro' }], - savedPackageCount: limit, - userId, - }), - PACKAGE_SERVICE_INSTANCE: mockPackageServiceNamespace(), - } as Env + const env = createEnv(userId, { + users: [{ email, plan: 'pro' }], + savedPackageCount: limit, + }) mockModule.loadPackageSourceBySourceId.mockResolvedValue({ manifest: { name: '@kentcdodds/shade-automation', @@ -1322,14 +1338,10 @@ test('refreshSavedPackageProjection does not gate the update branch at the limit const userId = await createStableUserIdFromEmail(email) const limit = planLimits.pro.maxSavedPackages if (limit === null) throw new Error('Expected a numeric pro package limit.') - const env = { - APP_DB: createEntitlementsDatabase({ - users: [{ email, plan: 'pro' }], - savedPackageCount: limit, - userId, - }), - PACKAGE_SERVICE_INSTANCE: mockPackageServiceNamespace(), - } as Env + const env = createEnv(userId, { + users: [{ email, plan: 'pro' }], + savedPackageCount: limit, + }) mockModule.loadPackageSourceBySourceId.mockResolvedValue({ manifest: { name: '@kentcdodds/shade-automation', diff --git a/packages/worker/src/test-support/incidental-runtime-warnings.ts b/packages/worker/src/test-support/incidental-runtime-warnings.ts index 0cc4b18d15..5cdfcf8ca8 100644 --- a/packages/worker/src/test-support/incidental-runtime-warnings.ts +++ b/packages/worker/src/test-support/incidental-runtime-warnings.ts @@ -11,6 +11,7 @@ const incidentalRuntimeWarnings = [ 'usage-event-record-failed', 'usage-event-analytics-failed', 'usage-rollup-failed', + 'entitlement-storage-bytes-shadow-failed', 'run-record-begin-failed', 'run-record-start-failed', 'run-record-finish-failed', diff --git a/packages/worker/src/test-support/user-meter.ts b/packages/worker/src/test-support/user-meter.ts index 8b274cb760..cce96a76d5 100644 --- a/packages/worker/src/test-support/user-meter.ts +++ b/packages/worker/src/test-support/user-meter.ts @@ -833,6 +833,47 @@ export function createWaitUntilDrain() { } } +/** Minimal D1 hooks for DO-authoritative {@link withAccountWriteLease} in node tests. */ +export function createPermissiveAccountWriteLeaseDbHooks() { + return { + supportsDeletingAtQuery(query: string) { + return query.includes( + 'SELECT deleting_at FROM users WHERE stable_user_id', + ) + }, + deletingAtFirstResult() { + return { deleting_at: null as string | null } + }, + supportsHeldLeaseQuery(query: string) { + return query.includes('SELECT 1 AS held FROM account_write_leases') + }, + heldLeaseFirstResult() { + return { held: 1 as number } + }, + supportsWriteLeaseRun(query: string) { + const normalized = query.replace(/\s+/g, ' ').trim() + return ( + normalized.includes('INSERT INTO account_write_leases') || + normalized.includes('UPDATE account_write_leases') || + normalized.includes('DELETE FROM account_write_leases') || + normalized.startsWith('UPDATE users') + ) + }, + writeLeaseRunResult() { + return { meta: { changes: 1 } } + }, + async runWriteLeaseBatch( + statements: Array<{ run: () => Promise }>, + ) { + const results = [] + for (const statement of statements) { + results.push(await statement.run()) + } + return results + }, + } +} + /** * Patch `db.prepare` and restore it via `using` even when the body throws. */ From 29dc25747d6c0e40ebce8f829629732edb34f17a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 11:07:23 +0000 Subject: [PATCH 26/39] test(account): wire meter into workers fixtures Co-authored-by: Kent C. Dodds --- packages/worker/src/mcp-auth.workers.test.ts | 88 +++++++++++++++---- .../src/mcp/observability.workers.test.ts | 18 ++-- 2 files changed, 86 insertions(+), 20 deletions(-) diff --git a/packages/worker/src/mcp-auth.workers.test.ts b/packages/worker/src/mcp-auth.workers.test.ts index e52ac0c3e7..38483a76e8 100644 --- a/packages/worker/src/mcp-auth.workers.test.ts +++ b/packages/worker/src/mcp-auth.workers.test.ts @@ -1,3 +1,4 @@ +import { env } from 'cloudflare:workers' import { expect, test } from 'vitest' import { type OAuthHelpers, @@ -13,6 +14,7 @@ import { import { oauthScopes } from './oauth-handlers.ts' import { createStableUserIdFromEmail } from './user-id.ts' import { consoleError } from '#worker/test-support/console-spies.ts' +import { createWaitUntilDrain } from '#worker/test-support/user-meter.ts' function expectAuthenticateHeader( header: string, @@ -132,6 +134,18 @@ function createMockDb(options: MockDbOptions = {}) { options.auditInserts?.push(boundParams) return { meta: { changes: 1 } } } + if (normalized.startsWith('delete from account_write_leases')) { + if (!boundParams.includes(expectedLeaseUserId)) { + throw new Error('Unscoped account write lease delete.') + } + return { meta: { changes: 1 } } + } + if (normalized.startsWith('insert into account_write_leases')) { + if (!boundParams.includes(expectedLeaseUserId)) { + throw new Error('Unscoped account write lease insert.') + } + return { meta: { changes: 1 } } + } if (normalized.startsWith('update users')) { if ( !normalized.includes('where stable_user_id = ?') || @@ -177,6 +191,25 @@ function createMockDb(options: MockDbOptions = {}) { suspended_at: options.suspendedAt ?? null, } satisfies MockAccountRow } + if (normalized.includes('select 1 as held from account_write_leases')) { + return { held: 1 } + } + if (normalized.includes('select deleting_at from users')) { + if (!boundStableUserId) return null + if (options.accountByStableId !== undefined) { + if ( + options.accountByStableId === null || + options.accountByStableId.stable_user_id !== boundStableUserId + ) { + return null + } + return { + deleting_at: options.accountByStableId.deleting_at ?? null, + } + } + if (boundStableUserId !== defaultStableUserId) return null + return { deleting_at: null } + } if (normalized.includes('select suspended_at from users')) { // Validate the bound identity like the adjacent // verification branches so a widened isAccountSuspended @@ -266,6 +299,13 @@ function createMockDb(options: MockDbOptions = {}) { } return { prepare: (query: string) => statementFor(query), + async batch(statements: Array<{ run: () => Promise }>) { + const results = [] + for (const statement of statements) { + results.push(await statement.run()) + } + return results + }, } as unknown as D1Database } @@ -277,16 +317,34 @@ function createEnv( return { APP_DB: createMockDb(dbOptions), OAUTH_PROVIDER: helpers, + USER_METER: env.USER_METER, ...overrides, } as unknown as Env } function createContext() { + const drain = createWaitUntilDrain() return { props: {}, - waitUntil: () => undefined, + waitUntil: drain.waitUntil, passThroughOnException: () => undefined, - } as unknown as ExecutionContext + drain: drain.drain, + } +} + +type TestContext = ReturnType + +async function handleMcpRequestAndDrain( + input: Omit[0], 'ctx'> & { + ctx: TestContext + }, +) { + const response = await handleMcpRequest({ + ...input, + ctx: input.ctx as unknown as ExecutionContext, + }) + await input.ctx.drain() + return response } test('protected resource metadata and auth challenge resolve origin consistently', async () => { @@ -315,7 +373,7 @@ test('protected resource metadata and auth challenge resolve origin consistently buildProtectedResourceMetadata(workersDevOrigin), ) - const requestOriginUnauthorizedResponse = await handleMcpRequest({ + const requestOriginUnauthorizedResponse = await handleMcpRequestAndDrain({ request: new Request(`${requestOrigin}${mcpResourcePath}`), env: createEnv(createHelpers()), ctx: createContext(), @@ -335,7 +393,7 @@ test('protected resource metadata and auth challenge resolve origin consistently requestOrigin, ) - const appBaseUrlUnauthorizedResponse = await handleMcpRequest({ + const appBaseUrlUnauthorizedResponse = await handleMcpRequestAndDrain({ request: new Request(`${workersDevOrigin}${mcpResourcePath}`), env: createEnv(createHelpers(), { APP_BASE_URL: appBaseUrl, @@ -384,7 +442,7 @@ test('mcp request enforces token audience and forwards caller props', async () = userSelects, } - const invalidResponse = await handleMcpRequest({ + const invalidResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -396,7 +454,7 @@ test('mcp request enforces token audience and forwards caller props', async () = }) expect(invalidResponse.status).toBe(401) - const missingAudienceResponse = await handleMcpRequest({ + const missingAudienceResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -409,7 +467,7 @@ test('mcp request enforces token audience and forwards caller props', async () = expect(missingAudienceResponse.status).toBe(401) let receivedProps: unknown = null - const validResponse = await handleMcpRequest({ + const validResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -437,7 +495,7 @@ test('mcp request enforces token audience and forwards caller props', async () = expect(userSelects[0]).toContain('suspended_at') expect(verificationLookups).toHaveLength(0) - const withConnectorResponse = await handleMcpRequest({ + const withConnectorResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -480,7 +538,7 @@ test('mcp request enforces token audience and forwards caller props', async () = }, } as unknown as D1Database await expect( - handleMcpRequest({ + handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -499,7 +557,7 @@ test('mcp request enforces token audience and forwards caller props', async () = 'Failed to load MCP auth user context:', expect.any(Error), ) -}) +}, 15_000) test('mcp request rejects unverified and unidentifiable accounts fail-closed', async () => { const request = new Request(`https://example.com${mcpResourcePath}`, { @@ -528,7 +586,7 @@ test('mcp request rejects unverified and unidentifiable accounts fail-closed', a } // Account exists but email_verified_at is null. - const unverifiedResponse = await handleMcpRequest({ + const unverifiedResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -548,7 +606,7 @@ test('mcp request rejects unverified and unidentifiable accounts fail-closed', a }) // No matching account row at all. - const unknownAccountResponse = await handleMcpRequest({ + const unknownAccountResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -562,7 +620,7 @@ test('mcp request rejects unverified and unidentifiable accounts fail-closed', a expect(unknownAccountResponse.status).toBe(403) // Grant props without an identifiable user. - const noUserResponse = await handleMcpRequest({ + const noUserResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -579,7 +637,7 @@ test('mcp request rejects unverified and unidentifiable accounts fail-closed', a // the rejection is recorded so a suspended principal that keeps calling // stays visible instead of failing silently. const auditInserts: Array> = [] - const suspendedResponse = await handleMcpRequest({ + const suspendedResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ @@ -615,7 +673,7 @@ test('mcp request rejects unverified and unidentifiable accounts fail-closed', a const stableUserId = await createStableUserIdFromEmail(fallbackEmail) const verifiedAt = new Date(0).toISOString() const fallbackUserSelects: Array = [] - const fallbackResponse = await handleMcpRequest({ + const fallbackResponse = await handleMcpRequestAndDrain({ request, env: createEnv( createHelpers({ diff --git a/packages/worker/src/mcp/observability.workers.test.ts b/packages/worker/src/mcp/observability.workers.test.ts index 0c6856940c..af420e6bb7 100644 --- a/packages/worker/src/mcp/observability.workers.test.ts +++ b/packages/worker/src/mcp/observability.workers.test.ts @@ -1,3 +1,4 @@ +import { env } from 'cloudflare:workers' import { expect, test, vi } from 'vitest' import { McpCallerError } from '#mcp/caller-error.ts' import { getStaticRegistry } from '#mcp/capabilities/registry.ts' @@ -24,6 +25,13 @@ vi.mock('#worker/repo/source-sync.ts', () => ({ repoMockModule.syncArtifactSourceSnapshot(...args), })) +function createTestEnv(overrides: Record = {}) { + return { + USER_METER: env.USER_METER, + ...overrides, + } as unknown as Env +} + function resetRepoPersistenceMocks() { repoMockModule.ensureEntitySource.mockReset() repoMockModule.syncArtifactSourceSnapshot.mockReset() @@ -193,7 +201,7 @@ test('package_save logs parse failures, rejects invalid manifests, and logs succ handler( {}, { - env: {} as Env, + env: createTestEnv(), callerContext: createMcpCallerContext({ baseUrl: 'https://example.com', }), @@ -217,7 +225,7 @@ test('package_save logs parse failures, rejects invalid manifests, and logs succ resetRepoPersistenceMocks() const handler = getStaticRegistry().capabilityMap['package_save'].handler const signedInContext = { - env: { + env: createTestEnv({ APP_DB: { prepare() { return { @@ -229,7 +237,7 @@ test('package_save logs parse failures, rejects invalid manifests, and logs succ } }, }, - } as unknown as Env, + }), callerContext: createMcpCallerContext({ baseUrl: 'https://example.com', user: { @@ -352,7 +360,7 @@ test('package_save logs parse failures, rejects invalid manifests, and logs succ ], }, { - env: { + env: createTestEnv({ APP_DB: { prepare(query: string) { return { @@ -513,7 +521,7 @@ test('package_save logs parse failures, rejects invalid manifests, and logs succ data: [Array.from({ length: 384 }, () => 0)], }), }, - } as unknown as Env, + }), callerContext: createMcpCallerContext({ baseUrl: 'https://example.com', user: { From 11814af9b5b599eea3fbfaf67e843850b2fc854a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 11:23:01 +0000 Subject: [PATCH 27/39] fix(account): clean partial lease mirrors Co-authored-by: Kent C. Dodds --- .../src/account/deletion-state.node.test.ts | 103 ++++++++++++++++++ packages/worker/src/account/deletion-state.ts | 8 ++ 2 files changed, 111 insertions(+) diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index 8ebf165fe2..d8f238f5c2 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -1144,6 +1144,109 @@ test('DO acquire fails closed and releases DO when D1 mirror response is lost un db.prepare = originalPrepare }) +test('DO acquire cleans partial D1 mirror when insert commits and count returns false', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const originalBatch = db.batch.bind(db) + db.batch = (async () => { + throw new Error('simulated lost D1 batch response') + }) as D1Database['batch'] + const originalPrepare = db.prepare.bind(db) + db.prepare = ((query: string) => { + const statement = originalPrepare(query) + if (query.includes('active_write_count + 1')) { + return { + bind(...params: Array) { + const bound = statement.bind(...params) + return { + async run() { + return { meta: { changes: 0 } } + }, + first: bound.first.bind(bound), + all: bound.all.bind(bound), + } + }, + } + } + return statement + }) as D1Database['prepare'] + let writeStarted = false + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:mirror-partial-false', + env: meter.env, + async write() { + writeStarted = true + return 'should-not-run' + }, + }), + ).rejects.toBeInstanceOf(AccountDeletionInProgressError) + expect(writeStarted).toBe(false) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + db.batch = originalBatch + db.prepare = originalPrepare +}) + +test('DO acquire cleans partial D1 mirror when sequential count throws after insert', async () => { + const { sqlite, db } = createLeaseTestDb() + const meter = createInMemoryUserMeterEnv() + const meterA = userMeterRpc({ env: meter.env, userId: 'user-a' }) + const originalBatch = db.batch.bind(db) + db.batch = (async () => { + throw new Error('simulated lost D1 batch response') + }) as D1Database['batch'] + const originalPrepare = db.prepare.bind(db) + const countFailure = new Error('simulated sequential count failure') + db.prepare = ((query: string) => { + const statement = originalPrepare(query) + if (query.includes('active_write_count + 1')) { + return { + bind(...params: Array) { + const bound = statement.bind(...params) + return { + async run() { + throw countFailure + }, + first: bound.first.bind(bound), + all: bound.all.bind(bound), + } + }, + } + } + return statement + }) as D1Database['prepare'] + let writeStarted = false + + await expect( + withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:mirror-partial-throw', + env: meter.env, + async write() { + writeStarted = true + return 'should-not-run' + }, + }), + ).rejects.toBe(countFailure) + expect(writeStarted).toBe(false) + expect(countLeaseRows(sqlite)).toBe(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) + expect(await meterA.countActiveWriteLeases()).toEqual({ count: 0 }) + db.batch = originalBatch + db.prepare = originalPrepare +}) + test('DO repair clears D1 mirror, stays audit-first, and retries after finalize without double audit/count', async () => { const { sqlite, db } = createLeaseTestDb() addLeaseRepairsTable(sqlite) diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index ea6b7cddfa..4c48a6ff9e 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -541,6 +541,14 @@ async function acquireDoAccountWriteLeaseAndWrite(input: { throw new AccountDeletionInProgressError() } } catch (error) { + // Mirror acquire can leave a partial D1 row (insert committed, count + // missed/threw). Clear it before DO release; cleanup must not mask + // the original error. A failed clear may overblock old Phase-A marks. + await releaseD1WriteLeaseMirror({ + db: input.db, + stableUserId: lease.stableUserId, + token: lease.token, + }).catch(() => {}) await meter.releaseWriteLease({ token: lease.token }).catch(() => {}) throw error } From 4c0aa230487daa40ce1d8c0a337cb00af5c1d3d7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 14:41:44 +0000 Subject: [PATCH 28/39] feat(admin): add UserMeter parity report Co-authored-by: Kent C. Dodds --- .../src/admin/user-meter-parity.node.test.ts | 599 ++++++++++++++++++ .../worker/src/admin/user-meter-parity.ts | 557 ++++++++++++++++ .../admin-user-meter-parity.node.test.ts | 154 +++++ .../admin/admin-user-meter-parity.ts | 134 ++++ .../src/mcp/capabilities/admin/domain.ts | 7 +- .../mcp/capabilities/registry.node.test.ts | 2 + 6 files changed, 1452 insertions(+), 1 deletion(-) create mode 100644 packages/worker/src/admin/user-meter-parity.node.test.ts create mode 100644 packages/worker/src/admin/user-meter-parity.ts create mode 100644 packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts create mode 100644 packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts diff --git a/packages/worker/src/admin/user-meter-parity.node.test.ts b/packages/worker/src/admin/user-meter-parity.node.test.ts new file mode 100644 index 0000000000..1bd5481e8c --- /dev/null +++ b/packages/worker/src/admin/user-meter-parity.node.test.ts @@ -0,0 +1,599 @@ +import { DatabaseSync } from 'node:sqlite' +import { expect, test } from 'vitest' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { createD1FromSqlite } from '#worker/test-support/create-d1-from-sqlite.ts' +import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' +import { dailyEntitlementResources } from '#worker/entitlements/user-meter-do.ts' +import { loadAdminUserMeterParityReport } from './user-meter-parity.ts' + +const stableUserId = testStableUserIdFromEmail('parity@example.com') +const now = new Date('2026-08-01T12:00:00.000Z') +const day = utcDayKey(now) + +function createParityTestDb() { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec(` + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + stable_user_id TEXT UNIQUE NOT NULL, + username TEXT NOT NULL, + email TEXT NOT NULL, + password_hash TEXT, + d1_storage_bytes INTEGER NOT NULL DEFAULT 0, + deleting_at TEXT, + active_write_count INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z', + updated_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z' + ); + CREATE TABLE entitlement_daily_counters ( + user_id TEXT NOT NULL, + resource TEXT NOT NULL, + day TEXT NOT NULL, + count INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, resource, day) + ); + CREATE TABLE package_service_states ( + user_id TEXT NOT NULL, + package_id TEXT NOT NULL, + service_name TEXT NOT NULL, + status TEXT NOT NULL CHECK (status IN ('running', 'idle', 'stopped', 'error')), + started_at TEXT, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, package_id, service_name) + ); + CREATE TABLE account_write_leases ( + token TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + holder TEXT NOT NULL, + acquired_at TEXT NOT NULL, + released_at TEXT + ); + `) + return { sqlite, db: createD1FromSqlite(sqlite) } +} + +function insertUser( + sqlite: DatabaseSync, + input: { + stableUserId: string + d1StorageBytes?: number + deletingAt?: string | null + }, +) { + sqlite + .prepare( + `INSERT INTO users ( + stable_user_id, username, email, d1_storage_bytes, deleting_at + ) VALUES (?, ?, ?, ?, ?)`, + ) + .run( + input.stableUserId, + 'parity-user', + 'parity@example.com', + input.d1StorageBytes ?? 0, + input.deletingAt ?? null, + ) +} + +async function seedMeterParityBaseline(input: { + meter: ReturnType + stableUserId: string + dailyCounts: Record<(typeof dailyEntitlementResources)[number], number> + storageBytes: number + packageServices?: Array<{ + packageId: string + serviceName: string + status: 'running' | 'idle' | 'stopped' | 'error' + startedAt?: string | null + sourceUpdatedAt: string + }> + deletingAt?: string | null + leases?: Array<{ + token: string + holder: string + acquiredAt: string + authority: 'do' | 'legacy' + }> +}) { + const meter = userMeterRpc({ + env: input.meter.env, + userId: input.stableUserId, + }) + for (const resource of dailyEntitlementResources) { + await meter.initialize({ + resource, + day, + count: input.dailyCounts[resource], + updatedAt: now.toISOString(), + }) + } + await meter.initializeStorageBytes({ + bytes: input.storageBytes, + updatedAt: now.toISOString(), + }) + for (const service of input.packageServices ?? []) { + await meter.upsertPackageServiceState({ + packageId: service.packageId, + serviceName: service.serviceName, + status: service.status, + startedAt: service.startedAt ?? null, + sourceUpdatedAt: service.sourceUpdatedAt, + updatedAt: now.toISOString(), + }) + } + await meter.bootstrapDeletionState({ + deletingAt: input.deletingAt ?? null, + leases: input.leases, + }) +} + +function seedD1ParityBaseline(input: { + sqlite: DatabaseSync + stableUserId: string + dailyCounts: Record<(typeof dailyEntitlementResources)[number], number> + storageBytes: number + packageServices?: Array<{ + packageId: string + serviceName: string + status: 'running' | 'idle' | 'stopped' | 'error' + startedAt?: string | null + updatedAt: string + }> + deletingAt?: string | null + leases?: Array<{ + token: string + holder: string + acquiredAt: string + }> +}) { + insertUser(input.sqlite, { + stableUserId: input.stableUserId, + d1StorageBytes: input.storageBytes, + deletingAt: input.deletingAt ?? null, + }) + for (const resource of dailyEntitlementResources) { + input.sqlite + .prepare( + `INSERT INTO entitlement_daily_counters ( + user_id, resource, day, count, updated_at + ) VALUES (?, ?, ?, ?, ?)`, + ) + .run( + input.stableUserId, + resource, + day, + input.dailyCounts[resource], + now.toISOString(), + ) + } + for (const service of input.packageServices ?? []) { + input.sqlite + .prepare( + `INSERT INTO package_service_states ( + user_id, package_id, service_name, status, started_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?)`, + ) + .run( + input.stableUserId, + service.packageId, + service.serviceName, + service.status, + service.status === 'running' ? (service.startedAt ?? null) : null, + service.updatedAt, + ) + } + for (const lease of input.leases ?? []) { + input.sqlite + .prepare( + `INSERT INTO account_write_leases ( + token, user_id, holder, acquired_at, released_at + ) VALUES (?, ?, ?, ?, NULL)`, + ) + .run(lease.token, input.stableUserId, lease.holder, lease.acquiredAt) + } +} + +function assertNoLeaseSecrets(value: unknown) { + const serialized = JSON.stringify(value) + expect(serialized).not.toMatch(/lease-token|holder-secret|secret-holder/i) + expect(serialized).not.toContain('"token"') + expect(serialized).not.toContain('"holder"') +} + +test('loadAdminUserMeterParityReport reports full parity across daily/storage/services/deletion', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + const dailyCounts = { + email_sends_per_day: 3, + email_receives_per_day: 5, + execute_calls_per_day: 11, + outbound_fetches_per_day: 7, + } + const serviceUpdatedAt = '2026-08-01T11:00:00.000Z' + const leaseToken = 'lease-token-shared' + seedD1ParityBaseline({ + sqlite, + stableUserId, + dailyCounts, + storageBytes: 4096, + packageServices: [ + { + packageId: 'pkg-a', + serviceName: 'web', + status: 'running', + startedAt: '2026-08-01T10:00:00.000Z', + updatedAt: serviceUpdatedAt, + }, + ], + leases: [ + { + token: leaseToken, + holder: 'holder-secret', + acquiredAt: '2026-08-01T09:00:00.000Z', + }, + ], + }) + await seedMeterParityBaseline({ + meter, + stableUserId, + dailyCounts, + storageBytes: 4096, + packageServices: [ + { + packageId: 'pkg-a', + serviceName: 'web', + status: 'running', + startedAt: '2026-08-01T10:00:00.000Z', + sourceUpdatedAt: serviceUpdatedAt, + }, + ], + leases: [ + { + token: leaseToken, + holder: 'holder-secret', + acquiredAt: '2026-08-01T09:00:00.000Z', + authority: 'do', + }, + ], + }) + + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(report).toMatchObject({ + generatedAt: now.toISOString(), + stableUserId, + daily: { + day, + mismatchCount: 0, + }, + storage: { + d1Bytes: 4096, + meterBytes: 4096, + needsBootstrap: false, + delta: 0, + parity: true, + }, + packageServices: { + d1Count: 1, + meterCount: 1, + truncated: false, + mismatchCategories: { + d1Only: 0, + meterOnly: 0, + statusMismatch: 0, + startedAtMismatch: 0, + sourceUpdatedAtMismatch: 0, + }, + running: { + d1FreshRunningCount: 1, + meterRunningCount: 1, + parity: true, + }, + parity: true, + }, + deletion: { + d1DeletingAt: null, + meterDeletingAt: null, + deletingAtParity: true, + d1ActiveLeaseCount: 1, + doAuthorityLeaseCount: 1, + doLegacyLeaseCount: 0, + tokenSetMismatches: { + d1Only: 0, + doOnly: 0, + legacyWithoutD1: 0, + }, + truncated: false, + mirrorLeaseParity: true, + }, + }) + expect(report?.daily.resources).toHaveLength(4) + expect(report?.daily.resources.every((row) => row.parity)).toBe(true) + assertNoLeaseSecrets(report) +}) + +test('loadAdminUserMeterParityReport treats D1-only leases as mirrorLeaseParity pass', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + const dailyCounts = { + email_sends_per_day: 0, + email_receives_per_day: 0, + execute_calls_per_day: 0, + outbound_fetches_per_day: 0, + } + seedD1ParityBaseline({ + sqlite, + stableUserId, + dailyCounts, + storageBytes: 0, + leases: [ + { + token: 'lease-token-email-only', + holder: 'holder-secret', + acquiredAt: '2026-08-01T09:00:00.000Z', + }, + ], + }) + await seedMeterParityBaseline({ + meter, + stableUserId, + dailyCounts, + storageBytes: 0, + }) + + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(report?.deletion).toMatchObject({ + d1ActiveLeaseCount: 1, + doAuthorityLeaseCount: 0, + doLegacyLeaseCount: 0, + tokenSetMismatches: { + d1Only: 1, + doOnly: 0, + legacyWithoutD1: 0, + }, + truncated: false, + mirrorLeaseParity: true, + }) + assertNoLeaseSecrets(report) +}) + +test('loadAdminUserMeterParityReport classifies mismatches and needsBootstrap without writing parity state', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + insertUser(sqlite, { stableUserId, d1StorageBytes: 100 }) + sqlite + .prepare( + `INSERT INTO entitlement_daily_counters ( + user_id, resource, day, count, updated_at + ) VALUES (?, 'email_sends_per_day', ?, 9, ?)`, + ) + .run(stableUserId, day, now.toISOString()) + sqlite + .prepare( + `INSERT INTO package_service_states ( + user_id, package_id, service_name, status, started_at, updated_at + ) VALUES (?, 'pkg-a', 'web', 'running', ?, ?)`, + ) + .run(stableUserId, '2026-08-01T10:00:00.000Z', '2026-08-01T11:00:00.000Z') + sqlite + .prepare( + `INSERT INTO account_write_leases ( + token, user_id, holder, acquired_at, released_at + ) VALUES ('d1-only-token', ?, 'holder-secret', ?, NULL)`, + ) + .run(stableUserId, '2026-08-01T09:00:00.000Z') + + const meterStub = userMeterRpc({ env: meter.env, userId: stableUserId }) + await meterStub.initialize({ + resource: 'email_sends_per_day', + day, + count: 2, + updatedAt: now.toISOString(), + }) + await meterStub.upsertPackageServiceState({ + packageId: 'pkg-b', + serviceName: 'worker', + status: 'idle', + sourceUpdatedAt: '2026-08-01T11:30:00.000Z', + }) + await meterStub.bootstrapDeletionState({ + leases: [ + { + token: 'do-only-token', + holder: 'secret-holder', + acquiredAt: '2026-08-01T08:00:00.000Z', + authority: 'do', + }, + { + token: 'legacy-only-token', + holder: 'secret-holder', + acquiredAt: '2026-08-01T07:00:00.000Z', + authority: 'legacy', + }, + ], + }) + + const beforeRead = await meterStub.read({ + resource: 'email_receives_per_day', + day, + now: now.toISOString(), + }) + expect(beforeRead).toEqual({ outcome: 'needs_bootstrap' }) + + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(report?.daily.mismatchCount).toBe(4) + expect( + report?.daily.resources.find( + (row) => row.resource === 'email_sends_per_day', + ), + ).toMatchObject({ + d1Count: 9, + meterCount: 2, + needsBootstrap: false, + delta: 7, + parity: false, + }) + expect( + report?.daily.resources.find( + (row) => row.resource === 'email_receives_per_day', + ), + ).toMatchObject({ + d1Count: 0, + meterCount: null, + needsBootstrap: true, + delta: null, + parity: false, + }) + expect(report?.storage).toMatchObject({ + d1Bytes: 100, + meterBytes: null, + needsBootstrap: true, + delta: null, + parity: false, + }) + expect(report?.packageServices).toMatchObject({ + d1Count: 1, + meterCount: 1, + mismatchCategories: { + d1Only: 1, + meterOnly: 1, + statusMismatch: 0, + startedAtMismatch: 0, + sourceUpdatedAtMismatch: 0, + }, + running: { + d1FreshRunningCount: 1, + meterRunningCount: 0, + parity: false, + }, + parity: false, + }) + expect(report?.deletion).toMatchObject({ + d1ActiveLeaseCount: 1, + doAuthorityLeaseCount: 1, + doLegacyLeaseCount: 1, + tokenSetMismatches: { + d1Only: 1, + doOnly: 1, + legacyWithoutD1: 1, + }, + mirrorLeaseParity: false, + }) + + const afterRead = await meterStub.read({ + resource: 'email_receives_per_day', + day, + now: now.toISOString(), + }) + expect(afterRead).toEqual({ outcome: 'needs_bootstrap' }) + expect(await meterStub.readStorageBytes()).toEqual({ + outcome: 'needs_bootstrap', + }) + assertNoLeaseSecrets(report) +}) + +test('loadAdminUserMeterParityReport fails mirrorLeaseParity for doOnly or legacyWithoutD1', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + insertUser(sqlite, { stableUserId, d1StorageBytes: 0 }) + const meterStub = userMeterRpc({ env: meter.env, userId: stableUserId }) + await meterStub.initializeStorageBytes({ + bytes: 0, + updatedAt: now.toISOString(), + }) + for (const resource of dailyEntitlementResources) { + await meterStub.initialize({ + resource, + day, + count: 0, + updatedAt: now.toISOString(), + }) + } + await meterStub.bootstrapDeletionState({ + leases: [ + { + token: 'do-only-token', + holder: 'secret-holder', + acquiredAt: '2026-08-01T08:00:00.000Z', + authority: 'do', + }, + ], + }) + + const doOnlyReport = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(doOnlyReport?.deletion).toMatchObject({ + d1ActiveLeaseCount: 0, + doAuthorityLeaseCount: 1, + tokenSetMismatches: { + d1Only: 0, + doOnly: 1, + legacyWithoutD1: 0, + }, + mirrorLeaseParity: false, + }) + + await meterStub.releaseWriteLease({ token: 'do-only-token' }) + await meterStub.bootstrapDeletionState({ + leases: [ + { + token: 'legacy-only-token', + holder: 'secret-holder', + acquiredAt: '2026-08-01T07:00:00.000Z', + authority: 'legacy', + }, + ], + }) + const legacyReport = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(legacyReport?.deletion).toMatchObject({ + d1ActiveLeaseCount: 0, + doAuthorityLeaseCount: 0, + doLegacyLeaseCount: 1, + tokenSetMismatches: { + d1Only: 0, + doOnly: 0, + legacyWithoutD1: 1, + }, + mirrorLeaseParity: false, + }) + assertNoLeaseSecrets(doOnlyReport) + assertNoLeaseSecrets(legacyReport) +}) + +test('loadAdminUserMeterParityReport returns null for missing users', async () => { + const { db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId: testStableUserIdFromEmail('missing@example.com'), + now, + }) + expect(report).toBeNull() +}) diff --git a/packages/worker/src/admin/user-meter-parity.ts b/packages/worker/src/admin/user-meter-parity.ts new file mode 100644 index 0000000000..375d67448e --- /dev/null +++ b/packages/worker/src/admin/user-meter-parity.ts @@ -0,0 +1,557 @@ +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' +import { + userMeterRpc, + type UserMeterEnv, +} from '#worker/entitlements/user-meter-client.ts' +import { + dailyEntitlementResources, + type DailyEntitlementResource, + type UserMeterPackageServiceState, + type UserMeterWriteLeaseShadow, +} from '#worker/entitlements/user-meter-do.ts' +import { + countRunningPackageServices, + packageServiceStateStaleMs, +} from '#worker/entitlements/service.ts' +import { planLimits } from '#worker/entitlements/plans.ts' +import { isStableUserId, normalizeStableUserId } from '#worker/user-id.ts' + +/** Cap paged UserMeter inventory pulls (matches max plan package-service ceiling). */ +export const userMeterParityInventoryMaxRows = planLimits.max.maxPackageServices + +/** Page size for UserMeter list RPCs (matches DO max export/list page). */ +export const userMeterParityPageSize = 500 + +type DailyResourceParity = { + resource: DailyEntitlementResource + d1Count: number + meterCount: number | null + needsBootstrap: boolean + delta: number | null + parity: boolean +} + +type StorageParity = { + d1Bytes: number + meterBytes: number | null + needsBootstrap: boolean + delta: number | null + parity: boolean +} + +type PackageServiceMismatchCategories = { + d1Only: number + meterOnly: number + statusMismatch: number + startedAtMismatch: number + sourceUpdatedAtMismatch: number +} + +type PackageServicesParity = { + d1Count: number + meterCount: number + truncated: boolean + mismatchCategories: PackageServiceMismatchCategories + running: { + d1FreshRunningCount: number + meterRunningCount: number + parity: boolean + } + parity: boolean +} + +type DeletionLeaseTokenMismatches = { + d1Only: number + doOnly: number + legacyWithoutD1: number +} + +type DeletionParity = { + d1DeletingAt: string | null + meterDeletingAt: string | null + deletingAtParity: boolean + d1ActiveLeaseCount: number + doAuthorityLeaseCount: number + doLegacyLeaseCount: number + tokenSetMismatches: DeletionLeaseTokenMismatches + truncated: boolean + /** + * Temporary same-token D1 mirror readiness for DO-authority leases. + * Passes when every DO-authority and legacy meter lease is mirrored in D1 + * (`doOnly` / `legacyWithoutD1` are zero), the meter page walk is complete, + * and D1 active lease count covers DO-authority count. `d1Only` is reported + * but does not fail this gate — expected for email/transition D1 leases + * until that path hands off. + */ + mirrorLeaseParity: boolean +} + +export type AdminUserMeterParityReport = { + generatedAt: string + stableUserId: string + daily: { + day: string + resources: Array + mismatchCount: number + } + storage: StorageParity + packageServices: PackageServicesParity + deletion: DeletionParity +} + +type D1PackageServiceStateRow = { + package_id: string + service_name: string + status: string + started_at: string | null + updated_at: string +} + +type D1WriteLeaseTokenRow = { + token: string +} + +function countDeltaParity(input: { + d1Value: number + meterValue: number | null + needsBootstrap: boolean +}): { delta: number | null; parity: boolean } { + if (input.needsBootstrap || input.meterValue == null) { + return { delta: null, parity: false } + } + const delta = input.d1Value - input.meterValue + return { delta, parity: delta === 0 } +} + +async function userExists(db: D1Database, stableUserId: string) { + const row = await db + .prepare(`SELECT 1 AS present FROM users WHERE stable_user_id = ?`) + .bind(stableUserId) + .first<{ present: number }>() + return row != null +} + +async function readD1DailyCounts(input: { + db: D1Database + stableUserId: string + day: string +}): Promise> { + const rows = await input.db + .prepare( + `SELECT resource, count + FROM entitlement_daily_counters + WHERE user_id = ? AND day = ?`, + ) + .bind(input.stableUserId, input.day) + .all<{ resource: string; count: number }>() + const counts = new Map() + for (const resource of dailyEntitlementResources) { + counts.set(resource, 0) + } + for (const row of rows.results ?? []) { + if ( + (dailyEntitlementResources as ReadonlyArray).includes( + row.resource, + ) + ) { + counts.set( + row.resource as DailyEntitlementResource, + Math.max(0, Number(row.count ?? 0)), + ) + } + } + return counts +} + +async function readDailyParity(input: { + db: D1Database + env: UserMeterEnv + stableUserId: string + day: string + generatedAt: string +}): Promise { + const d1Counts = await readD1DailyCounts(input) + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const resources: Array = [] + for (const resource of dailyEntitlementResources) { + const d1Count = d1Counts.get(resource) ?? 0 + const meterRead = await meter.read({ + resource, + day: input.day, + now: input.generatedAt, + }) + const needsBootstrap = meterRead.outcome === 'needs_bootstrap' + const meterCount = needsBootstrap ? null : meterRead.count + const { delta, parity } = countDeltaParity({ + d1Value: d1Count, + meterValue: meterCount, + needsBootstrap, + }) + resources.push({ + resource, + d1Count, + meterCount, + needsBootstrap, + delta, + parity, + }) + } + return { + day: input.day, + resources, + mismatchCount: resources.filter((row) => !row.parity).length, + } +} + +async function readStorageParity(input: { + db: D1Database + env: UserMeterEnv + stableUserId: string +}): Promise { + const d1Row = await input.db + .prepare( + `SELECT d1_storage_bytes AS bytes + FROM users + WHERE stable_user_id = ?`, + ) + .bind(input.stableUserId) + .first<{ bytes: number }>() + const d1Bytes = Math.max(0, Number(d1Row?.bytes ?? 0)) + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const meterRead = await meter.readStorageBytes() + const needsBootstrap = meterRead.outcome === 'needs_bootstrap' + const meterBytes = needsBootstrap ? null : meterRead.bytes + const { delta, parity } = countDeltaParity({ + d1Value: d1Bytes, + meterValue: meterBytes, + needsBootstrap, + }) + return { + d1Bytes, + meterBytes, + needsBootstrap, + delta, + parity, + } +} + +async function listAllMeterPackageServiceStates(input: { + env: UserMeterEnv + stableUserId: string + maxRows: number +}): Promise<{ + states: Array + truncated: boolean +}> { + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const states: Array = [] + let startAfter: string | null = null + let truncated = false + for (;;) { + const remaining = input.maxRows - states.length + if (remaining <= 0) { + truncated = true + break + } + const page = await meter.listPackageServiceStates({ + pageSize: Math.min(userMeterParityPageSize, remaining), + startAfter, + }) + states.push(...page.states) + if (!page.nextStartAfter) { + truncated = truncated || page.truncated + break + } + if (states.length >= input.maxRows) { + truncated = true + break + } + startAfter = page.nextStartAfter + } + return { states, truncated } +} + +async function listAllMeterWriteLeases(input: { + env: UserMeterEnv + stableUserId: string + maxRows: number +}): Promise<{ + leases: Array + truncated: boolean +}> { + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const leases: Array = [] + let startAfter: string | null = null + let truncated = false + for (;;) { + const remaining = input.maxRows - leases.length + if (remaining <= 0) { + truncated = true + break + } + const page = await meter.listWriteLeases({ + pageSize: Math.min(userMeterParityPageSize, remaining), + startAfter, + }) + leases.push(...page.leases) + if (!page.nextStartAfter) { + truncated = truncated || page.truncated + break + } + if (leases.length >= input.maxRows) { + truncated = true + break + } + startAfter = page.nextStartAfter + } + return { leases, truncated } +} + +function packageServiceKey(packageId: string, serviceName: string) { + return `${packageId}\0${serviceName}` +} + +function comparePackageServiceStates(input: { + d1Rows: ReadonlyArray + meterStates: ReadonlyArray +}): PackageServiceMismatchCategories { + const d1ByKey = new Map( + input.d1Rows.map((row) => [ + packageServiceKey(row.package_id, row.service_name), + row, + ]), + ) + const meterByKey = new Map( + input.meterStates.map((row) => [ + packageServiceKey(row.packageId, row.serviceName), + row, + ]), + ) + let d1Only = 0 + let meterOnly = 0 + let statusMismatch = 0 + let startedAtMismatch = 0 + let sourceUpdatedAtMismatch = 0 + for (const [key, d1Row] of d1ByKey) { + const meterRow = meterByKey.get(key) + if (!meterRow) { + d1Only += 1 + continue + } + if (d1Row.status !== meterRow.status) statusMismatch += 1 + const d1StartedAt = d1Row.started_at ?? null + const meterStartedAt = meterRow.startedAt ?? null + if (d1StartedAt !== meterStartedAt) startedAtMismatch += 1 + if (d1Row.updated_at !== meterRow.sourceUpdatedAt) { + sourceUpdatedAtMismatch += 1 + } + } + for (const key of meterByKey.keys()) { + if (!d1ByKey.has(key)) meterOnly += 1 + } + return { + d1Only, + meterOnly, + statusMismatch, + startedAtMismatch, + sourceUpdatedAtMismatch, + } +} + +async function readPackageServicesParity(input: { + db: D1Database + env: UserMeterEnv + stableUserId: string + now: Date +}): Promise { + const d1RowsResult = await input.db + .prepare( + `SELECT package_id, service_name, status, started_at, updated_at + FROM package_service_states + WHERE user_id = ? + ORDER BY package_id ASC, service_name ASC`, + ) + .bind(input.stableUserId) + .all() + const d1Rows = d1RowsResult.results ?? [] + const { states: meterStates, truncated } = + await listAllMeterPackageServiceStates({ + env: input.env, + stableUserId: input.stableUserId, + maxRows: userMeterParityInventoryMaxRows, + }) + const mismatchCategories = comparePackageServiceStates({ + d1Rows, + meterStates, + }) + const d1FreshRunningCount = await countRunningPackageServices({ + db: input.db, + userId: input.stableUserId, + now: input.now, + }) + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const meterRunning = await meter.countRunningPackageServices({ + staleAfterMs: packageServiceStateStaleMs, + now: input.now.toISOString(), + }) + const runningParity = d1FreshRunningCount === meterRunning.count + const inventoryMismatchCount = + mismatchCategories.d1Only + + mismatchCategories.meterOnly + + mismatchCategories.statusMismatch + + mismatchCategories.startedAtMismatch + + mismatchCategories.sourceUpdatedAtMismatch + return { + d1Count: d1Rows.length, + meterCount: meterStates.length, + truncated, + mismatchCategories, + running: { + d1FreshRunningCount, + meterRunningCount: meterRunning.count, + parity: runningParity, + }, + parity: inventoryMismatchCount === 0 && runningParity && !truncated, + } +} + +function compareDeletionLeaseTokens(input: { + d1Tokens: ReadonlyArray + meterLeases: ReadonlyArray +}): DeletionLeaseTokenMismatches { + const d1TokenSet = new Set(input.d1Tokens) + const doAuthorityTokens = new Set() + const legacyTokens = new Set() + for (const lease of input.meterLeases) { + if (lease.authority === 'do') doAuthorityTokens.add(lease.token) + else legacyTokens.add(lease.token) + } + let d1Only = 0 + for (const token of d1TokenSet) { + if (!doAuthorityTokens.has(token) && !legacyTokens.has(token)) { + d1Only += 1 + } + } + let doOnly = 0 + for (const token of doAuthorityTokens) { + if (!d1TokenSet.has(token)) doOnly += 1 + } + let legacyWithoutD1 = 0 + for (const token of legacyTokens) { + if (!d1TokenSet.has(token)) legacyWithoutD1 += 1 + } + return { d1Only, doOnly, legacyWithoutD1 } +} + +async function readDeletionParity(input: { + db: D1Database + env: UserMeterEnv + stableUserId: string +}): Promise { + const d1DeletingRow = await input.db + .prepare(`SELECT deleting_at FROM users WHERE stable_user_id = ?`) + .bind(input.stableUserId) + .first<{ deleting_at: string | null }>() + const d1DeletingAt = d1DeletingRow?.deleting_at ?? null + const d1LeaseRows = await input.db + .prepare( + `SELECT token + FROM account_write_leases + WHERE user_id = ? AND released_at IS NULL + ORDER BY acquired_at, token`, + ) + .bind(input.stableUserId) + .all() + const d1Tokens = (d1LeaseRows.results ?? []).map((row) => row.token) + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const meterDeletion = await meter.readDeletionState() + const { leases: meterLeases, truncated } = await listAllMeterWriteLeases({ + env: input.env, + stableUserId: input.stableUserId, + maxRows: userMeterParityInventoryMaxRows, + }) + let doAuthorityLeaseCount = 0 + let doLegacyLeaseCount = 0 + for (const lease of meterLeases) { + if (lease.authority === 'do') doAuthorityLeaseCount += 1 + else doLegacyLeaseCount += 1 + } + const tokenSetMismatches = compareDeletionLeaseTokens({ + d1Tokens, + meterLeases, + }) + const deletingAtParity = d1DeletingAt === meterDeletion.deletingAt + const d1ActiveLeaseCount = d1Tokens.length + const mirrorLeaseParity = + tokenSetMismatches.doOnly === 0 && + tokenSetMismatches.legacyWithoutD1 === 0 && + !truncated && + d1ActiveLeaseCount >= doAuthorityLeaseCount + return { + d1DeletingAt, + meterDeletingAt: meterDeletion.deletingAt, + deletingAtParity, + d1ActiveLeaseCount, + doAuthorityLeaseCount, + doLegacyLeaseCount, + tokenSetMismatches, + truncated, + mirrorLeaseParity, + } +} + +/** + * Read-only D1 ↔ UserMeter parity report for one account. Never bootstraps or + * writes parity state. Opening a cold UserMeter stub may still run Durable + * Object constructor schema maintenance and opportunistic stale daily-counter + * pruning; cold meter rows surface as `needsBootstrap`. + */ +export async function loadAdminUserMeterParityReport(input: { + db: D1Database + env: UserMeterEnv + stableUserId: string + now?: Date +}): Promise { + const stableUserId = normalizeStableUserId(input.stableUserId) + if (!isStableUserId(stableUserId)) return null + if (!(await userExists(input.db, stableUserId))) return null + + const now = input.now ?? new Date() + const generatedAt = now.toISOString() + const day = utcDayKey(now) + + const [daily, storage, packageServices, deletion] = await Promise.all([ + readDailyParity({ + db: input.db, + env: input.env, + stableUserId, + day, + generatedAt, + }), + readStorageParity({ + db: input.db, + env: input.env, + stableUserId, + }), + readPackageServicesParity({ + db: input.db, + env: input.env, + stableUserId, + now, + }), + readDeletionParity({ + db: input.db, + env: input.env, + stableUserId, + }), + ]) + + return { + generatedAt, + stableUserId, + daily, + storage, + packageServices, + deletion, + } +} diff --git a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts new file mode 100644 index 0000000000..ceccd65157 --- /dev/null +++ b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts @@ -0,0 +1,154 @@ +import { DatabaseSync } from 'node:sqlite' +import { expect, test, vi } from 'vitest' +import type * as AuditLog from '#worker/audit-log.ts' +import { createD1FromSqlite } from '#worker/test-support/create-d1-from-sqlite.ts' +import { createMcpCallerContext } from '#mcp/context.ts' +import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' + +const mockModule = vi.hoisted(() => ({ + logAuditEvent: vi.fn(async () => undefined), +})) + +vi.mock('#worker/audit-log.ts', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + logAuditEvent: (...args: Array) => + mockModule.logAuditEvent(...args), + } +}) + +const { adminDomain } = await import('./domain.ts') +const { adminUserMeterParityCapability } = + await import('./admin-user-meter-parity.ts') + +const stableUserId = testStableUserIdFromEmail('parity-cap@example.com') + +function createCapabilityTestDb() { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec(` + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + stable_user_id TEXT UNIQUE NOT NULL, + username TEXT NOT NULL, + email TEXT NOT NULL, + d1_storage_bytes INTEGER NOT NULL DEFAULT 0, + deleting_at TEXT, + created_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z', + updated_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z' + ); + CREATE TABLE entitlement_daily_counters ( + user_id TEXT NOT NULL, + resource TEXT NOT NULL, + day TEXT NOT NULL, + count INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, resource, day) + ); + CREATE TABLE package_service_states ( + user_id TEXT NOT NULL, + package_id TEXT NOT NULL, + service_name TEXT NOT NULL, + status TEXT NOT NULL, + started_at TEXT, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, package_id, service_name) + ); + CREATE TABLE account_write_leases ( + token TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + holder TEXT NOT NULL, + acquired_at TEXT NOT NULL, + released_at TEXT + ); + INSERT INTO users (stable_user_id, username, email, d1_storage_bytes) + VALUES ('${stableUserId}', 'parity-cap', 'parity-cap@example.com', 1); + INSERT INTO account_write_leases ( + token, user_id, holder, acquired_at, released_at + ) VALUES ( + 'lease-token-privacy', '${stableUserId}', 'holder-secret', + '2026-08-01T09:00:00.000Z', NULL + ); + `) + return createD1FromSqlite(sqlite) +} + +test('admin_user_meter_parity is registered admin-only read-only with parity keywords', () => { + const capability = adminDomain.capabilities.find( + (entry) => entry.name === 'admin_user_meter_parity', + ) + expect(capability).toMatchObject({ + requiredRole: 'admin', + readOnly: true, + idempotent: true, + destructive: false, + }) + expect(adminDomain.keywords).toEqual( + expect.arrayContaining(['user meter', 'parity', 'cutover']), + ) + expect(adminUserMeterParityCapability.name).toBe('admin_user_meter_parity') +}) + +test('admin_user_meter_parity returns null for missing users and omits lease secrets', async () => { + const db = createCapabilityTestDb() + const meter = createInMemoryUserMeterEnv() + const meterStub = userMeterRpc({ env: meter.env, userId: stableUserId }) + await meterStub.bootstrapDeletionState({ + leases: [ + { + token: 'lease-token-privacy', + holder: 'holder-secret', + acquiredAt: '2026-08-01T09:00:00.000Z', + authority: 'do', + }, + ], + }) + const ctx = { + env: { APP_DB: db, ...meter.env } as Env, + callerContext: createMcpCallerContext({ + baseUrl: 'https://heykody.dev', + user: { + userId: testStableUserIdFromEmail('admin@example.com'), + email: 'admin@example.com', + displayName: 'Admin', + roles: ['admin'], + }, + }), + } + + const missing = await adminUserMeterParityCapability.handler( + { + stable_user_id: testStableUserIdFromEmail('missing@example.com'), + }, + ctx, + ) + expect(missing).toEqual({ report: null }) + expect(mockModule.logAuditEvent).toHaveBeenCalledWith( + expect.objectContaining({ + action: 'admin_user_meter_parity', + result: 'success', + reason: `target_stable_user_id=${testStableUserIdFromEmail('missing@example.com')}`, + }), + ) + + const present = await adminUserMeterParityCapability.handler( + { stable_user_id: stableUserId }, + ctx, + ) + expect(present.report?.stableUserId).toBe(stableUserId) + expect(present.report?.deletion.mirrorLeaseParity).toBe(true) + expect(mockModule.logAuditEvent).toHaveBeenCalledWith( + expect.objectContaining({ + action: 'admin_user_meter_parity', + result: 'success', + reason: `target_stable_user_id=${stableUserId}`, + }), + ) + const serialized = JSON.stringify(present) + expect(serialized).not.toContain('lease-token-privacy') + expect(serialized).not.toContain('holder-secret') + expect(serialized).not.toContain('"token"') + expect(serialized).not.toContain('"holder"') +}) diff --git a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts new file mode 100644 index 0000000000..2d40113791 --- /dev/null +++ b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts @@ -0,0 +1,134 @@ +import { z } from 'zod' +import { loadAdminUserMeterParityReport } from '#worker/admin/user-meter-parity.ts' +import { dailyEntitlementResources } from '#worker/entitlements/user-meter-do.ts' +import { defineDomainCapability } from '#mcp/capabilities/define-domain-capability.ts' +import { capabilityDomainNames } from '#mcp/capabilities/domain-metadata.ts' +import { + adminCapabilityAccess, + auditAdminCapabilityInvocation, + stableUserIdSchema, +} from './admin-shared.ts' + +const dailyResourceSchema = z.enum(dailyEntitlementResources) + +const dailyResourceParitySchema = z.object({ + resource: dailyResourceSchema, + d1Count: z.number().int().nonnegative(), + meterCount: z.number().int().nonnegative().nullable(), + needsBootstrap: z.boolean(), + delta: z.number().int().nullable(), + parity: z.boolean(), +}) + +const storageParitySchema = z.object({ + d1Bytes: z.number().int().nonnegative(), + meterBytes: z.number().int().nonnegative().nullable(), + needsBootstrap: z.boolean(), + delta: z.number().int().nullable(), + parity: z.boolean(), +}) + +const packageServiceMismatchCategoriesSchema = z.object({ + d1Only: z.number().int().nonnegative(), + meterOnly: z.number().int().nonnegative(), + statusMismatch: z.number().int().nonnegative(), + startedAtMismatch: z.number().int().nonnegative(), + sourceUpdatedAtMismatch: z.number().int().nonnegative(), +}) + +const packageServicesParitySchema = z.object({ + d1Count: z.number().int().nonnegative(), + meterCount: z.number().int().nonnegative(), + truncated: z.boolean(), + mismatchCategories: packageServiceMismatchCategoriesSchema, + running: z.object({ + d1FreshRunningCount: z.number().int().nonnegative(), + meterRunningCount: z.number().int().nonnegative(), + parity: z.boolean(), + }), + parity: z.boolean(), +}) + +const deletionLeaseTokenMismatchesSchema = z.object({ + d1Only: z.number().int().nonnegative(), + doOnly: z.number().int().nonnegative(), + legacyWithoutD1: z.number().int().nonnegative(), +}) + +const deletionParitySchema = z.object({ + d1DeletingAt: z.string().nullable(), + meterDeletingAt: z.string().nullable(), + deletingAtParity: z.boolean(), + d1ActiveLeaseCount: z.number().int().nonnegative(), + doAuthorityLeaseCount: z.number().int().nonnegative(), + doLegacyLeaseCount: z.number().int().nonnegative(), + tokenSetMismatches: deletionLeaseTokenMismatchesSchema, + truncated: z.boolean(), + mirrorLeaseParity: z + .boolean() + .describe( + 'Temporary same-token D1 mirror readiness: doOnly and legacyWithoutD1 are zero, inventory is not truncated, and d1ActiveLeaseCount covers doAuthorityLeaseCount. d1Only is reported separately and does not fail this gate.', + ), +}) + +const reportSchema = z.object({ + generatedAt: z.string(), + stableUserId: stableUserIdSchema, + daily: z.object({ + day: z.string(), + resources: z.array(dailyResourceParitySchema), + mismatchCount: z.number().int().nonnegative(), + }), + storage: storageParitySchema, + packageServices: packageServicesParitySchema, + deletion: deletionParitySchema, +}) + +const inputSchema = z.object({ + stable_user_id: stableUserIdSchema, +}) + +const outputSchema = z.object({ + report: reportSchema.nullable(), +}) + +export const adminUserMeterParityCapability = defineDomainCapability( + capabilityDomainNames.admin, + { + ...adminCapabilityAccess, + name: 'admin_user_meter_parity', + description: + 'Read-only production verification report comparing one user D1 entitlement mirrors against UserMeter for daily counters, storage bytes, package-service liveness, and deletion write leases. Never bootstraps or writes parity state (DO constructor schema maintenance and stale daily pruning may still run); returns counts and parity only (no lease tokens/holders or email content). Admin-only.', + keywords: [ + 'admin', + 'user meter', + 'parity', + 'cutover', + 'daily counters', + 'storage bytes', + 'package services', + 'write leases', + 'deletion', + 'mirror', + 'bootstrap', + ], + inputSchema, + outputSchema, + async handler(args, ctx) { + return auditAdminCapabilityInvocation( + ctx, + 'admin_user_meter_parity', + async () => ({ + report: await loadAdminUserMeterParityReport({ + db: ctx.env.APP_DB, + env: ctx.env, + stableUserId: args.stable_user_id, + }), + }), + { + successReason: () => `target_stable_user_id=${args.stable_user_id}`, + }, + ) + }, + }, +) diff --git a/packages/worker/src/mcp/capabilities/admin/domain.ts b/packages/worker/src/mcp/capabilities/admin/domain.ts index 1cc69be789..f562e7d0a5 100644 --- a/packages/worker/src/mcp/capabilities/admin/domain.ts +++ b/packages/worker/src/mcp/capabilities/admin/domain.ts @@ -28,11 +28,12 @@ import { adminUserListCapability } from './admin-user-list.ts' import { adminUserUpdateCapability } from './admin-user-update.ts' import { adminAccountWriteLeaseListCapability } from './admin-account-write-lease-list.ts' import { adminAccountWriteLeaseRepairCapability } from './admin-account-write-lease-repair.ts' +import { adminUserMeterParityCapability } from './admin-user-meter-parity.ts' export const adminDomain = defineDomain({ name: capabilityDomainNames.admin, description: - 'Admin-only operator capabilities for account metadata, platform accounts, package scope grants, fleet package-codemod scan/dry-run/apply/revert over published package trees, feature flags, operator-owned system email, attributed platform feedback users explicitly submit for admin review, and metadata about activity on public community listings; never exposes private package source or unrelated user content such as secrets, memories, jobs, or user inbox email.', + 'Admin-only operator capabilities for account metadata, platform accounts, package scope grants, fleet package-codemod scan/dry-run/apply/revert over published package trees, feature flags, operator-owned system email, attributed platform feedback users explicitly submit for admin review, UserMeter↔D1 parity verification, and metadata about activity on public community listings; never exposes private package source or unrelated user content such as secrets, memories, jobs, or user inbox email.', keywords: [ 'admin', 'rbac', @@ -51,12 +52,16 @@ export const adminDomain = defineDomain({ 'package codemod', 'fleet', 'migration', + 'user meter', + 'parity', + 'cutover', ], capabilities: [ adminUserListCapability, adminUserGetCapability, adminUserCreateCapability, adminUserUpdateCapability, + adminUserMeterParityCapability, adminAccountWriteLeaseListCapability, adminAccountWriteLeaseRepairCapability, adminPlatformAccountCreateCapability, diff --git a/packages/worker/src/mcp/capabilities/registry.node.test.ts b/packages/worker/src/mcp/capabilities/registry.node.test.ts index 2633fcca93..ab53ac3ba4 100644 --- a/packages/worker/src/mcp/capabilities/registry.node.test.ts +++ b/packages/worker/src/mcp/capabilities/registry.node.test.ts @@ -110,10 +110,12 @@ test('getCapabilityRegistryForContext filters admin capabilities by current call }) expect(adminRegistry.capabilityMap.admin_user_list).toBeTruthy() + expect(adminRegistry.capabilityMap.admin_user_meter_parity).toBeTruthy() expect( adminRegistry.capabilityDomains.some((domain) => domain.name === 'admin'), ).toBe(true) expect(regularRegistry.capabilityMap.admin_user_list).toBeUndefined() + expect(regularRegistry.capabilityMap.admin_user_meter_parity).toBeUndefined() expect( regularRegistry.capabilityDomains.some((domain) => domain.name === 'admin'), ).toBe(false) From c44a518d8a2a02afa22e982011d478e7d0c4cd52 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 14:41:51 +0000 Subject: [PATCH 29/39] docs(entitlements): define production parity gates Co-authored-by: Kent C. Dodds --- .../contributing/architecture/entitlements.md | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 09cdd9a251..30117e2fc6 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -375,6 +375,40 @@ soak/parity review. waits until reporting-off-D1 work merges and mirror parity is verified in production. +### Admin UserMeter parity gates (`admin_user_meter_parity`) + +Production verification for mirror retirement and authority flips uses the +admin-only read-only capability `admin_user_meter_parity` (input: +`stable_user_id`). It compares production-shaped D1 rows for one account against +direct UserMeter RPCs and never bootstraps or writes parity state. Opening a +cold UserMeter stub may still run Durable Object constructor schema maintenance +and opportunistic stale daily-counter pruning. Cold meter rows surface as +`needsBootstrap` with `meterCount`/`meterBytes` null. + +Interpret the structured report as independent gates: + +| Gate | Pass condition | +| -------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Daily counters (current UTC day) | Each of the four daily resources has `parity: true` (`d1Count === meterCount`); aggregate `daily.mismatchCount === 0`. | +| Storage bytes | `storage.parity` — D1 `users.d1_storage_bytes` equals UserMeter `readStorageBytes` (not `needsBootstrap`). | +| Package services | `packageServices.parity` — inventory mismatch category counts are all zero (`d1Only` / `meterOnly` / `statusMismatch` / `startedAtMismatch` / `sourceUpdatedAtMismatch`), fresh-running counts match under the shared 24h stale window, and the meter page walk is not `truncated`. | +| Deletion tombstone | `deletion.deletingAtParity` — D1 `users.deleting_at` matches the meter tombstone. | +| Temporary D1 lease mirror | `deletion.mirrorLeaseParity` — `doOnly === 0`, `legacyWithoutD1 === 0`, inventory not truncated, and `d1ActiveLeaseCount >= doAuthorityLeaseCount` (same-token mirror coverage). `tokenSetMismatches.d1Only` is reported but does **not** fail this gate. | + +**D1-only leases:** email and other transition paths that omit `env` still take +exact D1 leases, so `d1Only > 0` is expected until that handoff. Mirror-removal +readiness therefore ignores `d1Only`. Operators separately confirm those rows +are known email/transition holders via `admin_account_write_lease_list` and +holder classification before retiring the temporary D1 mirror inventory. + +**Threshold:** treat unexplained mismatches as blocking for the corresponding +cutover (daily mirror retirement, storage authority flip, package-service +authority flip, or temporary D1 lease-mirror removal). Expected cold accounts +may report `needsBootstrap` until live traffic or an intentional bootstrap path +seeds the DO; that is a bootstrap gap, not a silent pass. Truncated inventories +fail closed (`parity` / `mirrorLeaseParity` false) so operators re-run or raise +the bounded page cap rather than approve a partial compare. + Module wiring: `consumeDailyEntitlement`, `refundDailyEntitlement`, and `readDailyEntitlementResourceUsage` require `env.USER_METER` and fail closed when the binding is missing. Storage-byte helpers (`readUserD1StorageBytes`, From a6180fc3b1a78feac33065d8fb69fd7150e4c61d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 14:57:31 +0000 Subject: [PATCH 30/39] fix(admin): bound parity page walks Co-authored-by: Kent C. Dodds --- .../src/admin/user-meter-parity.node.test.ts | 173 +++++++++++++++++- .../worker/src/admin/user-meter-parity.ts | 101 ++++++---- 2 files changed, 234 insertions(+), 40 deletions(-) diff --git a/packages/worker/src/admin/user-meter-parity.node.test.ts b/packages/worker/src/admin/user-meter-parity.node.test.ts index 1bd5481e8c..49b6b8a0e5 100644 --- a/packages/worker/src/admin/user-meter-parity.node.test.ts +++ b/packages/worker/src/admin/user-meter-parity.node.test.ts @@ -1,11 +1,18 @@ import { DatabaseSync } from 'node:sqlite' -import { expect, test } from 'vitest' +import { expect, test, vi } from 'vitest' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import { createD1FromSqlite } from '#worker/test-support/create-d1-from-sqlite.ts' import { testStableUserIdFromEmail } from '#worker/test-support/stable-user-id.ts' import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' -import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' -import { dailyEntitlementResources } from '#worker/entitlements/user-meter-do.ts' +import { + userMeterRpc, + type UserMeterEnv, +} from '#worker/entitlements/user-meter-client.ts' +import { + dailyEntitlementResources, + type UserMeterPackageServiceState, + type UserMeterWriteLeaseShadow, +} from '#worker/entitlements/user-meter-do.ts' import { loadAdminUserMeterParityReport } from './user-meter-parity.ts' const stableUserId = testStableUserIdFromEmail('parity@example.com') @@ -203,6 +210,58 @@ function assertNoLeaseSecrets(value: unknown) { expect(serialized).not.toContain('"holder"') } +type MeterListPageInput = { + pageSize?: number + startAfter?: string | null +} + +function withMeterListOverrides( + meter: ReturnType, + overrides: { + listPackageServiceStates?: (input: MeterListPageInput) => Promise<{ + states: Array + nextStartAfter: string | null + truncated: boolean + }> + listWriteLeases?: (input: MeterListPageInput) => Promise<{ + leases: Array + nextStartAfter: string | null + truncated: boolean + }> + }, +): UserMeterEnv { + const namespace = meter.env.USER_METER + if (!namespace) throw new Error('expected USER_METER binding') + const originalGet = namespace.get.bind(namespace) + return { + USER_METER: { + idFromName: namespace.idFromName.bind(namespace), + get: (id: DurableObjectId) => { + const stub = originalGet(id) as unknown as Record & { + listPackageServiceStates: (input: MeterListPageInput) => Promise<{ + states: Array + nextStartAfter: string | null + truncated: boolean + }> + listWriteLeases: (input: MeterListPageInput) => Promise<{ + leases: Array + nextStartAfter: string | null + truncated: boolean + }> + } + return { + ...stub, + listPackageServiceStates: + overrides.listPackageServiceStates ?? + stub.listPackageServiceStates.bind(stub), + listWriteLeases: + overrides.listWriteLeases ?? stub.listWriteLeases.bind(stub), + } + }, + }, + } as unknown as UserMeterEnv +} + test('loadAdminUserMeterParityReport reports full parity across daily/storage/services/deletion', async () => { const { sqlite, db } = createParityTestDb() const meter = createInMemoryUserMeterEnv() @@ -597,3 +656,111 @@ test('loadAdminUserMeterParityReport returns null for missing users', async () = }) expect(report).toBeNull() }) + +test('loadAdminUserMeterParityReport fails closed on repeated package-service cursor', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + const dailyCounts = { + email_sends_per_day: 0, + email_receives_per_day: 0, + execute_calls_per_day: 0, + outbound_fetches_per_day: 0, + } + seedD1ParityBaseline({ + sqlite, + stableUserId, + dailyCounts, + storageBytes: 0, + }) + await seedMeterParityBaseline({ + meter, + stableUserId, + dailyCounts, + storageBytes: 0, + }) + + const stuckCursor = '["pkg-a","web"]' + const listPackageServiceStates = vi.fn( + async (_input: MeterListPageInput) => ({ + states: [ + { + packageId: 'pkg-a', + serviceName: 'web', + status: 'idle' as const, + startedAt: null, + sourceUpdatedAt: '2026-08-01T11:00:00.000Z', + revision: 1, + updatedAt: now.toISOString(), + mirrorUpdatedAt: '1', + }, + ], + nextStartAfter: stuckCursor, + truncated: true, + }), + ) + const env = withMeterListOverrides(meter, { + listPackageServiceStates, + }) + + const report = await loadAdminUserMeterParityReport({ + db, + env, + stableUserId, + now, + }) + expect(listPackageServiceStates).toHaveBeenCalledTimes(2) + expect(listPackageServiceStates.mock.calls[0]?.[0]?.startAfter ?? null).toBe( + null, + ) + expect(listPackageServiceStates.mock.calls[1]?.[0]?.startAfter).toBe( + stuckCursor, + ) + expect(report?.packageServices).toMatchObject({ + truncated: true, + parity: false, + }) +}) + +test('loadAdminUserMeterParityReport fails closed on zero-item write-lease cursor', async () => { + const { sqlite, db } = createParityTestDb() + const meter = createInMemoryUserMeterEnv() + const dailyCounts = { + email_sends_per_day: 0, + email_receives_per_day: 0, + execute_calls_per_day: 0, + outbound_fetches_per_day: 0, + } + seedD1ParityBaseline({ + sqlite, + stableUserId, + dailyCounts, + storageBytes: 0, + }) + await seedMeterParityBaseline({ + meter, + stableUserId, + dailyCounts, + storageBytes: 0, + }) + + const listWriteLeases = vi.fn(async (_input: MeterListPageInput) => ({ + leases: [] as Array, + nextStartAfter: 'stuck-lease-cursor', + truncated: true, + })) + const env = withMeterListOverrides(meter, { + listWriteLeases, + }) + + const report = await loadAdminUserMeterParityReport({ + db, + env, + stableUserId, + now, + }) + expect(listWriteLeases).toHaveBeenCalledTimes(1) + expect(report?.deletion).toMatchObject({ + truncated: true, + mirrorLeaseParity: false, + }) +}) diff --git a/packages/worker/src/admin/user-meter-parity.ts b/packages/worker/src/admin/user-meter-parity.ts index 375d67448e..f5dd00b59a 100644 --- a/packages/worker/src/admin/user-meter-parity.ts +++ b/packages/worker/src/admin/user-meter-parity.ts @@ -235,39 +235,76 @@ async function readStorageParity(input: { } } -async function listAllMeterPackageServiceStates(input: { - env: UserMeterEnv - stableUserId: string +/** + * Bounded UserMeter list walk. Stops at maxRows, a missing cursor, a + * repeated/non-advancing cursor, or a zero-item page that still returns a + * cursor. Stalled/invalid progress marks `truncated` so callers fail closed. + */ +async function collectBoundedUserMeterPages(input: { maxRows: number -}): Promise<{ - states: Array - truncated: boolean -}> { - const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) - const states: Array = [] + fetchPage: (args: { + pageSize: number + startAfter: string | null + }) => Promise<{ + items: ReadonlyArray + nextStartAfter: string | null + truncated: boolean + }> +}): Promise<{ items: Array; truncated: boolean }> { + const items: Array = [] let startAfter: string | null = null let truncated = false for (;;) { - const remaining = input.maxRows - states.length + const remaining = input.maxRows - items.length if (remaining <= 0) { truncated = true break } - const page = await meter.listPackageServiceStates({ + const page = await input.fetchPage({ pageSize: Math.min(userMeterParityPageSize, remaining), startAfter, }) - states.push(...page.states) + items.push(...page.items) if (!page.nextStartAfter) { truncated = truncated || page.truncated break } - if (states.length >= input.maxRows) { + if (page.items.length === 0 || page.nextStartAfter === startAfter) { + truncated = true + break + } + if (items.length >= input.maxRows) { truncated = true break } startAfter = page.nextStartAfter } + return { items, truncated } +} + +async function listAllMeterPackageServiceStates(input: { + env: UserMeterEnv + stableUserId: string + maxRows: number +}): Promise<{ + states: Array + truncated: boolean +}> { + const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) + const { items: states, truncated } = await collectBoundedUserMeterPages({ + maxRows: input.maxRows, + fetchPage: async ({ pageSize, startAfter }) => { + const page = await meter.listPackageServiceStates({ + pageSize, + startAfter, + }) + return { + items: page.states, + nextStartAfter: page.nextStartAfter, + truncated: page.truncated, + } + }, + }) return { states, truncated } } @@ -280,30 +317,20 @@ async function listAllMeterWriteLeases(input: { truncated: boolean }> { const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) - const leases: Array = [] - let startAfter: string | null = null - let truncated = false - for (;;) { - const remaining = input.maxRows - leases.length - if (remaining <= 0) { - truncated = true - break - } - const page = await meter.listWriteLeases({ - pageSize: Math.min(userMeterParityPageSize, remaining), - startAfter, - }) - leases.push(...page.leases) - if (!page.nextStartAfter) { - truncated = truncated || page.truncated - break - } - if (leases.length >= input.maxRows) { - truncated = true - break - } - startAfter = page.nextStartAfter - } + const { items: leases, truncated } = await collectBoundedUserMeterPages({ + maxRows: input.maxRows, + fetchPage: async ({ pageSize, startAfter }) => { + const page = await meter.listWriteLeases({ + pageSize, + startAfter, + }) + return { + items: page.leases, + nextStartAfter: page.nextStartAfter, + truncated: page.truncated, + } + }, + }) return { leases, truncated } } From 177a4716d5438b6f9a5bc800e1a1ea591e52df17 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 15:41:44 +0000 Subject: [PATCH 31/39] refactor(entitlements): stop D1 daily counter mirror Co-authored-by: Kent C. Dodds --- packages/worker/src/account/data-targets.ts | 1 - .../src/admin/user-meter-parity.node.test.ts | 153 ++++++ .../worker/src/admin/user-meter-parity.ts | 59 ++- .../src/admin/user-usage-data.node.test.ts | 130 ++---- .../src/app/account-deletion.node.test.ts | 7 - .../src/app/account-retention-dispositions.ts | 1 - .../src/app/account-usage-data.node.test.ts | 73 +-- .../src/app/admin-insights-data.node.test.ts | 20 +- .../worker/src/app/admin-insights-data.ts | 37 +- .../app/handlers/account-email.node.test.ts | 105 ++--- .../worker/src/app/retention.node.test.ts | 37 +- packages/worker/src/app/retention.ts | 41 -- packages/worker/src/email/inbound-delivery.ts | 67 +-- .../inbound-entitlements.workers.test.ts | 24 - .../inbound-spam-controls.workers.test.ts | 13 +- .../worker/src/email/inbound.workers.test.ts | 27 +- packages/worker/src/email/outbound.ts | 3 +- .../entitlements/entitlements.node.test.ts | 218 ++++----- packages/worker/src/entitlements/service.ts | 206 ++------ .../worker/src/entitlements/test-schema.ts | 17 +- .../worker/src/entitlements/user-meter-do.ts | 2 +- .../entitlements/user-meter.workers.test.ts | 442 ++++-------------- .../admin/admin-capabilities.node.test.ts | 3 - .../admin-user-meter-parity.node.test.ts | 1 + .../admin/admin-user-meter-parity.ts | 10 +- .../email/email-usage-get.workers.test.ts | 39 +- .../worker/src/mcp/fetch-gateway.node.test.ts | 18 +- 27 files changed, 595 insertions(+), 1159 deletions(-) diff --git a/packages/worker/src/account/data-targets.ts b/packages/worker/src/account/data-targets.ts index eac40d291a..81fe11353f 100644 --- a/packages/worker/src/account/data-targets.ts +++ b/packages/worker/src/account/data-targets.ts @@ -219,7 +219,6 @@ export const accountUserDataTargets: ReadonlyArray = [ { kind: 'user_id', table: 'email_sender_identities' }, { kind: 'user_id', table: 'email_sender_rules' }, { kind: 'user_id', table: 'webhook_endpoints' }, - { kind: 'user_id', table: 'entitlement_daily_counters' }, { kind: 'user_columns', table: 'platform_feedback', diff --git a/packages/worker/src/admin/user-meter-parity.node.test.ts b/packages/worker/src/admin/user-meter-parity.node.test.ts index 49b6b8a0e5..47a1e7c792 100644 --- a/packages/worker/src/admin/user-meter-parity.node.test.ts +++ b/packages/worker/src/admin/user-meter-parity.node.test.ts @@ -330,6 +330,7 @@ test('loadAdminUserMeterParityReport reports full parity across daily/storage/se stableUserId, daily: { day, + mirrorRetired: false, mismatchCount: 0, }, storage: { @@ -764,3 +765,155 @@ test('loadAdminUserMeterParityReport fails closed on zero-item write-lease curso mirrorLeaseParity: false, }) }) + +test('loadAdminUserMeterParityReport reports mirrorRetired when entitlement_daily_counters is absent', async () => { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec(` + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + stable_user_id TEXT UNIQUE NOT NULL, + username TEXT NOT NULL, + email TEXT NOT NULL, + d1_storage_bytes INTEGER NOT NULL DEFAULT 0, + deleting_at TEXT, + created_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z', + updated_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z' + ); + CREATE TABLE package_service_states ( + user_id TEXT NOT NULL, + package_id TEXT NOT NULL, + service_name TEXT NOT NULL, + status TEXT NOT NULL CHECK (status IN ('running', 'idle', 'stopped', 'error')), + started_at TEXT, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, package_id, service_name) + ); + CREATE TABLE account_write_leases ( + token TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + holder TEXT NOT NULL, + acquired_at TEXT NOT NULL, + released_at TEXT + ); + `) + const db = createD1FromSqlite(sqlite) + const meter = createInMemoryUserMeterEnv() + insertUser(sqlite, { stableUserId, d1StorageBytes: 0 }) + const meterStub = userMeterRpc({ env: meter.env, userId: stableUserId }) + await meterStub.initialize({ + resource: 'email_sends_per_day', + day, + count: 4, + updatedAt: now.toISOString(), + }) + await meterStub.initializeStorageBytes({ + bytes: 0, + updatedAt: now.toISOString(), + }) + + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(report?.daily).toMatchObject({ + mirrorRetired: true, + mismatchCount: 0, + }) + expect( + report?.daily.resources.find( + (row) => row.resource === 'email_sends_per_day', + ), + ).toMatchObject({ + d1Count: null, + meterCount: 4, + delta: null, + parity: true, + }) +}) + +test('loadAdminUserMeterParityReport still compares D1 daily counts when mirror table exists', async () => { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec(` + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + stable_user_id TEXT UNIQUE NOT NULL, + username TEXT NOT NULL, + email TEXT NOT NULL, + password_hash TEXT, + d1_storage_bytes INTEGER NOT NULL DEFAULT 0, + deleting_at TEXT, + active_write_count INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z', + updated_at TEXT NOT NULL DEFAULT '2026-01-01T00:00:00.000Z' + ); + CREATE TABLE entitlement_daily_counters ( + user_id TEXT NOT NULL, + resource TEXT NOT NULL, + day TEXT NOT NULL, + count INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, resource, day) + ); + CREATE TABLE package_service_states ( + user_id TEXT NOT NULL, + package_id TEXT NOT NULL, + service_name TEXT NOT NULL, + status TEXT NOT NULL CHECK (status IN ('running', 'idle', 'stopped', 'error')), + started_at TEXT, + updated_at TEXT NOT NULL, + PRIMARY KEY (user_id, package_id, service_name) + ); + CREATE TABLE account_write_leases ( + token TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + holder TEXT NOT NULL, + acquired_at TEXT NOT NULL, + released_at TEXT + ); + `) + const db = createD1FromSqlite(sqlite) + const meter = createInMemoryUserMeterEnv() + insertUser(sqlite, { stableUserId, d1StorageBytes: 0 }) + for (const resource of dailyEntitlementResources) { + sqlite + .prepare( + `INSERT INTO entitlement_daily_counters ( + user_id, resource, day, count, updated_at + ) VALUES (?, ?, ?, ?, ?)`, + ) + .run(stableUserId, resource, day, 1, now.toISOString()) + } + const meterStub = userMeterRpc({ env: meter.env, userId: stableUserId }) + for (const resource of dailyEntitlementResources) { + await meterStub.initialize({ + resource, + day, + count: resource === 'email_sends_per_day' ? 9 : 1, + updatedAt: now.toISOString(), + }) + } + await meterStub.initializeStorageBytes({ + bytes: 0, + updatedAt: now.toISOString(), + }) + const report = await loadAdminUserMeterParityReport({ + db, + env: meter.env, + stableUserId, + now, + }) + expect(report?.daily.mirrorRetired).toBe(false) + expect(report?.daily.mismatchCount).toBe(1) + expect( + report?.daily.resources.find( + (row) => row.resource === 'email_sends_per_day', + ), + ).toMatchObject({ + d1Count: 1, + meterCount: 9, + delta: -8, + parity: false, + }) +}) diff --git a/packages/worker/src/admin/user-meter-parity.ts b/packages/worker/src/admin/user-meter-parity.ts index f5dd00b59a..6cf77d1c6a 100644 --- a/packages/worker/src/admin/user-meter-parity.ts +++ b/packages/worker/src/admin/user-meter-parity.ts @@ -24,10 +24,22 @@ export const userMeterParityPageSize = 500 type DailyResourceParity = { resource: DailyEntitlementResource - d1Count: number + /** + * D1 mirror count for the current UTC day. `null` when + * `daily.mirrorRetired` is true (table dropped / comparison retired). + */ + d1Count: number | null meterCount: number | null needsBootstrap: boolean + /** + * `d1Count - meterCount` when both sides are present; `null` when the + * mirror is retired or the meter still needs bootstrap. + */ delta: number | null + /** + * When the mirror is active: true iff counts match. When + * `mirrorRetired`, always true (no D1 comparison is claimed). + */ parity: boolean } @@ -91,6 +103,12 @@ export type AdminUserMeterParityReport = { stableUserId: string daily: { day: string + /** + * True when D1 `entitlement_daily_counters` is absent (retired). The + * daily gate then reports meter counts only and does not claim a D1 + * comparison (`d1Count`/`delta` are null; `parity` stays true). + */ + mirrorRetired: boolean resources: Array mismatchCount: number } @@ -131,6 +149,23 @@ async function userExists(db: D1Database, stableUserId: string) { return row != null } +/** + * Detect retirement without querying the missing table. Uses sqlite_master + * only (safe when the table has already been dropped). + */ +async function isEntitlementDailyCountersMirrorRetired( + db: D1Database, +): Promise { + const row = await db + .prepare( + `SELECT 1 AS present + FROM sqlite_master + WHERE type = 'table' AND name = 'entitlement_daily_counters'`, + ) + .first<{ present: number }>() + return row == null +} + async function readD1DailyCounts(input: { db: D1Database stableUserId: string @@ -170,11 +205,11 @@ async function readDailyParity(input: { day: string generatedAt: string }): Promise { - const d1Counts = await readD1DailyCounts(input) + const mirrorRetired = await isEntitlementDailyCountersMirrorRetired(input.db) + const d1Counts = mirrorRetired ? null : await readD1DailyCounts(input) const meter = userMeterRpc({ env: input.env, userId: input.stableUserId }) const resources: Array = [] for (const resource of dailyEntitlementResources) { - const d1Count = d1Counts.get(resource) ?? 0 const meterRead = await meter.read({ resource, day: input.day, @@ -182,6 +217,19 @@ async function readDailyParity(input: { }) const needsBootstrap = meterRead.outcome === 'needs_bootstrap' const meterCount = needsBootstrap ? null : meterRead.count + if (mirrorRetired) { + resources.push({ + resource, + d1Count: null, + meterCount, + needsBootstrap, + delta: null, + // Mirror comparison is retired; do not claim a D1 mismatch. + parity: true, + }) + continue + } + const d1Count = d1Counts?.get(resource) ?? 0 const { delta, parity } = countDeltaParity({ d1Value: d1Count, meterValue: meterCount, @@ -198,8 +246,11 @@ async function readDailyParity(input: { } return { day: input.day, + mirrorRetired, resources, - mismatchCount: resources.filter((row) => !row.parity).length, + mismatchCount: mirrorRetired + ? 0 + : resources.filter((row) => !row.parity).length, } } diff --git a/packages/worker/src/admin/user-usage-data.node.test.ts b/packages/worker/src/admin/user-usage-data.node.test.ts index 4cf563112b..c9b6248c03 100644 --- a/packages/worker/src/admin/user-usage-data.node.test.ts +++ b/packages/worker/src/admin/user-usage-data.node.test.ts @@ -32,13 +32,6 @@ type UsageRollupRow = { total_bytes: number } -type CounterRow = { - user_id: string - resource: string - day: string - count: number -} - type ResourceCount = Partial< Record< | 'saved_packages' @@ -59,12 +52,10 @@ function normalizeQuery(query: string) { function createAdminUserUsageTestDb(input: { users: Array usageRollups?: Array - dailyCounters?: Array resourceCounts?: Record }) { const users = input.users.map((user) => ({ ...user })) const usageRollups = input.usageRollups?.map((row) => ({ ...row })) ?? [] - const dailyCounters = input.dailyCounters?.map((row) => ({ ...row })) ?? [] const resourceCounts = input.resourceCounts ?? {} function countForQuery(normalizedQuery: string, userId: string) { @@ -106,15 +97,6 @@ function createAdminUserUsageTestDb(input: { return (users.find((user) => user.stable_user_id === params[0]) ?? null) as T | null } - if (normalizedQuery.includes('from entitlement_daily_counters')) { - const row = dailyCounters.find( - (counter) => - counter.user_id === params[0] && - counter.resource === params[1] && - counter.day === params[2], - ) - return (row ? { count: row.count } : null) as T | null - } const count = countForQuery(normalizedQuery, String(params[0])) if (count !== null) return { count } as T throw new Error(`Unsupported first query: ${query}`) @@ -243,14 +225,6 @@ test('loadAdminUserUsageData warns above eighty percent of plan limits', async ( event_count: 40, }), ], - dailyCounters: [ - { - user_id: usageUserId, - resource: 'email_sends_per_day', - day: '2026-07-05', - count: 170, - }, - ], resourceCounts: { [usageUserId]: { saved_packages: 85, @@ -262,8 +236,15 @@ test('loadAdminUserUsageData warns above eighty percent of plan limits', async ( }, }) + const env = withUserMeter({ APP_DB: db }) + await env.meter.seed({ + userId: usageUserId, + resource: 'email_sends_per_day', + day: '2026-07-05', + count: 170, + }) const data = await loadAdminUserUsageData( - withUserMeter({ APP_DB: db }) as Env, + env as Env, usageUserId, new Date('2026-07-05T12:00:00.000Z'), ) @@ -293,14 +274,6 @@ test('loadAdminUserUsageData rejects an invalid stored plan', async () => { stable_user_id: usageUserId, }, ], - dailyCounters: [ - { - user_id: usageUserId, - resource: 'email_receives_per_day', - day: '2026-07-05', - count: 190, - }, - ], resourceCounts: { [usageUserId]: { stored_email_messages: 12 }, }, @@ -439,43 +412,42 @@ test('loadAdminUserUsageData keeps current-month and month-over-month rollups on expect(getEventCount(data?.monthUsage[1]?.usage, 'execute')).toBe(12) }) -test('loadAdminUserUsageData reads daily counts from UserMeter (bootstrap then warm)', async () => { +test('loadAdminUserUsageData reads daily counts from UserMeter (seeded then warm)', async () => { const now = new Date('2026-07-05T12:00:00.000Z') const day = utcDayKey(now) const bootstrapEmail = 'bootstrap-drilldown@example.com' const bootstrapUserId = await createStableUserIdFromEmail(bootstrapEmail) - const bootstrapped = await loadAdminUserUsageData( - withUserMeter({ - APP_DB: createAdminUserUsageTestDb({ - users: [ - { - id: 3, - username: 'bootstrap', - email: bootstrapEmail, - plan: 'pro', - stable_user_id: bootstrapUserId, - }, - ], - dailyCounters: [ - { - user_id: bootstrapUserId, - resource: 'email_receives_per_day', - day, - count: 55, - }, - { - user_id: bootstrapUserId, - resource: 'outbound_fetches_per_day', - day, - count: 66, - }, - ], - resourceCounts: { - [bootstrapUserId]: { secrets: 3 }, + const bootstrapEnv = withUserMeter({ + APP_DB: createAdminUserUsageTestDb({ + users: [ + { + id: 3, + username: 'bootstrap', + email: bootstrapEmail, + plan: 'pro', + stable_user_id: bootstrapUserId, }, - }), - }) as Env, + ], + resourceCounts: { + [bootstrapUserId]: { secrets: 3 }, + }, + }), + }) + await bootstrapEnv.meter.seed({ + userId: bootstrapUserId, + resource: 'email_receives_per_day', + day, + count: 55, + }) + await bootstrapEnv.meter.seed({ + userId: bootstrapUserId, + resource: 'outbound_fetches_per_day', + day, + count: 66, + }) + const bootstrapped = await loadAdminUserUsageData( + bootstrapEnv as Env, bootstrapUserId, now, ) @@ -508,32 +480,6 @@ test('loadAdminUserUsageData reads daily counts from UserMeter (bootstrap then w stable_user_id: meterUserId, }, ], - dailyCounters: [ - { - user_id: meterUserId, - resource: 'email_sends_per_day', - day, - count: 11, - }, - { - user_id: meterUserId, - resource: 'email_receives_per_day', - day, - count: 22, - }, - { - user_id: meterUserId, - resource: 'execute_calls_per_day', - day, - count: 33, - }, - { - user_id: meterUserId, - resource: 'outbound_fetches_per_day', - day, - count: 44, - }, - ], resourceCounts: { [meterUserId]: { saved_packages: 6, diff --git a/packages/worker/src/app/account-deletion.node.test.ts b/packages/worker/src/app/account-deletion.node.test.ts index 01eb593fe9..170f7cfd1f 100644 --- a/packages/worker/src/app/account-deletion.node.test.ts +++ b/packages/worker/src/app/account-deletion.node.test.ts @@ -897,10 +897,6 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async email_attachments: [{ id: 'ea-1', message_id: 'em-1' }], email_delivery_events: [{ id: 'ed-1', user_id: userAaa }], email_sender_identities: [{ id: 'ei-1', user_id: userAaa }], - entitlement_daily_counters: [ - { user_id: userAaa, resource: 'email_sends_per_day', day: '2026-07-05' }, - { user_id: userBbb, resource: 'email_sends_per_day', day: '2026-07-05' }, - ], platform_feedback: [ { id: 'feedback-submitted-by-a', @@ -1306,9 +1302,6 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async 'email-raw:v1:user-aaa/em-1', 'email-raw:v1:user-aaa/em-2', ]) - expect(rows.entitlement_daily_counters).toEqual([ - { user_id: userBbb, resource: 'email_sends_per_day', day: '2026-07-05' }, - ]) expect(rows.platform_feedback).toEqual([ { id: 'feedback-reviewed-by-a', diff --git a/packages/worker/src/app/account-retention-dispositions.ts b/packages/worker/src/app/account-retention-dispositions.ts index 396c01cdcf..5178d8ace2 100644 --- a/packages/worker/src/app/account-retention-dispositions.ts +++ b/packages/worker/src/app/account-retention-dispositions.ts @@ -16,7 +16,6 @@ export const accountRetentionDispositions: ReadonlyArray) { return { ...env, ...meter.env, ...runLog.env, meter, runLog } } -type DailyCounterRow = { - user_id: string - resource: string - day: string - count: number -} - function createUsageTestDb(input: { userId: number email: string plan: string stripePlan?: string | null packageCount?: number - dailyCounters?: Array d1StorageBytes?: number }) { const stableUserId = testStableUserIdFromEmail(input.email) - const dailyCounters = input.dailyCounters?.map((row) => ({ ...row })) ?? [] const d1StorageBytes = input.d1StorageBytes ?? 0 return { stableUserId, @@ -53,15 +44,6 @@ function createUsageTestDb(input: { if (normalized.includes('from saved_packages')) { return { count: input.packageCount ?? 0 } as T } - if (normalized.includes('from entitlement_daily_counters')) { - const row = dailyCounters.find( - (counter) => - counter.user_id === params[0] && - counter.resource === params[1] && - counter.day === params[2], - ) - return (row ? { count: row.count } : { count: 0 }) as T - } if (normalized.includes('d1_storage_bytes')) { return { bytes: d1StorageBytes } as T } @@ -123,23 +105,22 @@ test('loadAccountUsageData returns plan rows and authoritative UserMeter daily c email: bootstrapEmail, plan: 'pro', packageCount: 1, - dailyCounters: [ - { - user_id: bootstrapUserId, - resource: 'email_sends_per_day', - day, - count: 17, - }, - { - user_id: bootstrapUserId, - resource: 'execute_calls_per_day', - day, - count: 91, - }, - ], + }) + const bootstrapEnv = withUsageEnv({ APP_DB: bootstrapDb }) + await bootstrapEnv.meter.seed({ + userId: bootstrapUserId, + resource: 'email_sends_per_day', + day, + count: 17, + }) + await bootstrapEnv.meter.seed({ + userId: bootstrapUserId, + resource: 'execute_calls_per_day', + day, + count: 91, }) const bootstrapped = await loadAccountUsageData({ - env: withUsageEnv({ APP_DB: bootstrapDb }) as Env, + env: bootstrapEnv as Env, userId: 8, now, }) @@ -154,32 +135,6 @@ test('loadAccountUsageData returns plan rows and authoritative UserMeter daily c email: meterEmail, plan: 'pro', packageCount: 4, - dailyCounters: [ - { - user_id: meterUserId, - resource: 'email_sends_per_day', - day, - count: 11, - }, - { - user_id: meterUserId, - resource: 'email_receives_per_day', - day, - count: 22, - }, - { - user_id: meterUserId, - resource: 'execute_calls_per_day', - day, - count: 33, - }, - { - user_id: meterUserId, - resource: 'outbound_fetches_per_day', - day, - count: 44, - }, - ], }) const warmEnv = withUsageEnv({ APP_DB: warmDb }) warmEnv.runLog.setActiveWorkflowCount(meterUserId, 3) diff --git a/packages/worker/src/app/admin-insights-data.node.test.ts b/packages/worker/src/app/admin-insights-data.node.test.ts index 6b3b0e849c..390e2b97a2 100644 --- a/packages/worker/src/app/admin-insights-data.node.test.ts +++ b/packages/worker/src/app/admin-insights-data.node.test.ts @@ -248,17 +248,6 @@ function createInsightsTestDb( ] as Array, } } - if (normalizedQuery.includes('from entitlement_daily_counters')) { - return { - results: [ - { - day: '2026-07-08', - resource: 'email_sends_per_day', - n: 3, - }, - ] as Array, - } - } if (normalizedQuery.includes('from email_delivery_events')) { return { results: [ @@ -322,12 +311,18 @@ function createInsightsTestDb( } test('loadAdminInsightsData assembles the dashboard payload', async () => { + consoleWarn.mockImplementation(() => {}) + consoleWarn.mockClear() const data = await loadAdminInsightsData( { APP_DB: createInsightsTestDb() } as Env, now, ) expect(data.ok).toBe(true) + expect(consoleWarn).toHaveBeenCalledWith( + 'admin-insights-email-quota-aggregate-unavailable', + { reason: 'entitlement-daily-counters-retired-local-dev' }, + ) expect(data.totals).toEqual({ users: 8, verifiedUsers: 5, @@ -353,7 +348,7 @@ test('loadAdminInsightsData assembles the dashboard payload', async () => { expect(data.emailByDay).toHaveLength(28) expect(data.emailByDay.at(-1)).toEqual({ day: '2026-07-08', - sends: 3, + sends: 0, receives: 0, }) expect(data.emailDeliveryByDay).toHaveLength(28) @@ -396,6 +391,7 @@ test('loadAdminInsightsData assembles the dashboard payload', async () => { }) test('activation latency excludes users who have no usable verification date', async () => { + consoleWarn.mockImplementation(() => {}) const data = await loadAdminInsightsData( { APP_DB: createInsightsTestDb({ activationLatencyRows: [] }) } as Env, now, diff --git a/packages/worker/src/app/admin-insights-data.ts b/packages/worker/src/app/admin-insights-data.ts index 9a2c3a7430..5babb4f167 100644 --- a/packages/worker/src/app/admin-insights-data.ts +++ b/packages/worker/src/app/admin-insights-data.ts @@ -281,29 +281,24 @@ async function loadEmailInsightsRows(input: { deliveryRows: Array }> { // Wrangler exposes a local Analytics Engine binding, but its SQL API cannot - // query the emulated dataset. Local development therefore keeps using the - // existing D1 tables, matching the feature-flag exposure fallback. + // query the emulated dataset. The D1 entitlement_daily_counters mirror is + // retired, so local/dev email quota aggregates degrade explicitly to empty + // (with a warning) rather than querying a removed table. Delivery-event + // aggregates still come from D1 email_delivery_events. if (!input.env.EMAIL_EVENTS || input.env.WRANGLER_IS_LOCAL_DEV === 'true') { - const [emailRows, deliveryRows] = await Promise.all([ - input.env.APP_DB.prepare( - `SELECT day, resource, SUM(count) AS n - FROM entitlement_daily_counters - WHERE day >= ? AND resource IN ('email_sends_per_day', 'email_receives_per_day') - GROUP BY day, resource`, - ) - .bind(input.dayCutoff) - .all(), - input.env.APP_DB.prepare( - `SELECT substr(created_at, 1, 10) AS day, event_type, COUNT(*) AS n - FROM email_delivery_events - WHERE provider = 'cloudflare-email' AND created_at >= ? - GROUP BY day, event_type`, - ) - .bind(input.dayCutoff) - .all(), - ]) + console.warn('admin-insights-email-quota-aggregate-unavailable', { + reason: 'entitlement-daily-counters-retired-local-dev', + }) + const deliveryRows = await input.env.APP_DB.prepare( + `SELECT substr(created_at, 1, 10) AS day, event_type, COUNT(*) AS n + FROM email_delivery_events + WHERE provider = 'cloudflare-email' AND created_at >= ? + GROUP BY day, event_type`, + ) + .bind(input.dayCutoff) + .all() return { - emailRows: emailRows.results ?? [], + emailRows: [], deliveryRows: deliveryRows.results ?? [], } } diff --git a/packages/worker/src/app/handlers/account-email.node.test.ts b/packages/worker/src/app/handlers/account-email.node.test.ts index 69293c4201..7110a388ff 100644 --- a/packages/worker/src/app/handlers/account-email.node.test.ts +++ b/packages/worker/src/app/handlers/account-email.node.test.ts @@ -4,6 +4,7 @@ import type * as EmailPlatformAddress from '#worker/email/platform-address.ts' import type * as EntitlementPlans from '#worker/entitlements/plans.ts' import type * as EntitlementService from '#worker/entitlements/service.ts' import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' const messageRow = { id: 'msg-1', @@ -251,41 +252,14 @@ function createListResult(rows: Array>) { function createEnv(input?: { meter?: ReturnType - dailyCounters?: Array<{ - user_id: string - resource: string - day: string - count: number - }> messageRows?: Array> messageTotal?: number }) { const meter = input?.meter ?? createInMemoryUserMeterEnv() - const dailyCounters = input?.dailyCounters?.map((row) => ({ ...row })) ?? [] const messageRows = input?.messageRows ?? [messageRow] const messageTotal = input?.messageTotal ?? messageRows.length mockModule.prepare.mockImplementation((query: string) => { const normalized = String(query).replace(/\s+/g, ' ').trim().toLowerCase() - if (normalized.includes('from entitlement_daily_counters')) { - return { - bind(...params: Array) { - return { - async first() { - const row = dailyCounters.find( - (counter) => - counter.user_id === params[0] && - counter.resource === params[1] && - counter.day === params[2], - ) - return (row ? { count: row.count } : { count: 0 }) as T - }, - async all() { - return { results: [] } - }, - } - }, - } - } if (normalized.includes('count(*)')) { return createCountResult(messageTotal) } @@ -321,20 +295,19 @@ test('email API lists messages with pagination, usage, and selected detail', asy const bootstrapMeter = createInMemoryUserMeterEnv() const env = createEnv({ meter: bootstrapMeter, - dailyCounters: [ - { - user_id: userId, - resource: 'email_sends_per_day', - day, - count: 2, - }, - { - user_id: userId, - resource: 'email_receives_per_day', - day, - count: 2, - }, - ], + }) + const meterStub = userMeterRpc({ env: bootstrapMeter.env, userId }) + await meterStub.initialize({ + resource: 'email_sends_per_day', + day, + count: 2, + updatedAt: new Date().toISOString(), + }) + await meterStub.initialize({ + resource: 'email_receives_per_day', + day, + count: 2, + updatedAt: new Date().toISOString(), }) const handler = createAccountEmailApiHandler(env) @@ -389,21 +362,25 @@ test('email API lists messages with pagination, usage, and selected detail', asy mockModule.prepare.mockClear() mockModule.getEmailMessageById.mockClear() + const selectionMeter = createInMemoryUserMeterEnv() const envWithSelection = createEnv({ - dailyCounters: [ - { - user_id: userId, - resource: 'email_sends_per_day', - day, - count: 2, - }, - { - user_id: userId, - resource: 'email_receives_per_day', - day, - count: 2, - }, - ], + meter: selectionMeter, + }) + const selectionStub = userMeterRpc({ + env: selectionMeter.env, + userId, + }) + await selectionStub.initialize({ + resource: 'email_sends_per_day', + day, + count: 2, + updatedAt: new Date().toISOString(), + }) + await selectionStub.initialize({ + resource: 'email_receives_per_day', + day, + count: 2, + updatedAt: new Date().toISOString(), }) const selectedResponse = await createAccountEmailApiHandler( envWithSelection, @@ -465,7 +442,7 @@ test('email API lists messages with pagination, usage, and selected detail', asy expect(unauthorizedResponse.status).toBe(401) }) -test('email API usage prefers initialized UserMeter daily counts over D1', async () => { +test('email API usage reads initialized UserMeter daily counts', async () => { using _frozenTime = useFrozenUtcTime('2026-07-31T15:00:00.000Z') const day = utcDayKey() const userId = 'stable-user-1' @@ -484,20 +461,6 @@ test('email API usage prefers initialized UserMeter daily counts over D1', async }) const env = createEnv({ meter, - dailyCounters: [ - { - user_id: userId, - resource: 'email_sends_per_day', - day, - count: 3, - }, - { - user_id: userId, - resource: 'email_receives_per_day', - day, - count: 5, - }, - ], }) const response = await createAccountEmailApiHandler(env).handler({ request: new Request('https://example.com/account/email.json'), @@ -513,7 +476,7 @@ test('email API usage prefers initialized UserMeter daily counts over D1', async receives_today: { count: 15, limit: 50 }, }), }) - expect(body.usage.sends_today.count).not.toBe(3) + expect(body.usage.sends_today.count).toBe(13) expect(body.usage.receives_today.count).not.toBe(5) }) diff --git a/packages/worker/src/app/retention.node.test.ts b/packages/worker/src/app/retention.node.test.ts index 95832df4cc..670863f3aa 100644 --- a/packages/worker/src/app/retention.node.test.ts +++ b/packages/worker/src/app/retention.node.test.ts @@ -6,14 +6,12 @@ import { auditEventRetentionDays, emailDeliveryEventRetentionDays, emailMessageRetentionDays, - entitlementDailyCounterRetentionDays, featureFlagExposureRetentionDays, getRetentionPolicyCoverage, memorySuppressionRetentionDays, platformFeedbackRetentionDays, pruneAuditEventsForRetention, pruneEmailDeliveryEventsForRetention, - pruneEntitlementDailyCountersForRetention, pruneFeatureFlagExposuresForRetention, pruneMemorySuppressionsForRetention, prunePlatformFeedbackForRetention, @@ -201,14 +199,6 @@ function createRetentionDb() { storage_key TEXT, created_at TEXT NOT NULL ); - CREATE TABLE entitlement_daily_counters ( - user_id TEXT NOT NULL, - resource TEXT NOT NULL, - day TEXT NOT NULL, - count INTEGER NOT NULL DEFAULT 0, - updated_at TEXT NOT NULL, - PRIMARY KEY (user_id, resource, day) - ); CREATE TABLE usage_rollups ( user_id TEXT NOT NULL, metric TEXT NOT NULL, @@ -865,22 +855,8 @@ test('retention prune reports selected separately from deleted when rows vanish ).toEqual({ selected: 2, deleted: 1 }) }) -test('entitlement counter and usage rollup retention respect boundaries', async () => { +test('usage rollup retention respects month boundaries', async () => { const { sqlite, db } = createRetentionDb() - const cutoffDay = daysAgo(entitlementDailyCounterRetentionDays).slice(0, 10) - for (const day of [ - daysAgo(entitlementDailyCounterRetentionDays + 1).slice(0, 10), - cutoffDay, - daysAgo(1).slice(0, 10), - ]) { - sqlite - .prepare( - `INSERT INTO entitlement_daily_counters ( - user_id, resource, day, count, updated_at - ) VALUES ('user-1', 'email_sends_per_day', ?, 1, ?)`, - ) - .run(day, now.toISOString()) - } // 24 months before 2026-07 keeps 2024-07 and later. for (const month of ['2024-06', '2024-07', '2026-06']) { sqlite @@ -892,22 +868,11 @@ test('entitlement counter and usage rollup retention respect boundaries', async .run(month, now.toISOString()) } - expect(await pruneEntitlementDailyCountersForRetention({ db, now })).toEqual({ - selected: 1, - deleted: 1, - }) expect(await pruneUsageRollupsForRetention({ db, now })).toEqual({ selected: 1, deleted: 1, }) - const days = sqlite - .prepare(`SELECT day FROM entitlement_daily_counters ORDER BY day`) - .all() as Array<{ day: string }> - expect(days.map((row) => row.day)).toEqual([ - cutoffDay, - daysAgo(1).slice(0, 10), - ]) const months = sqlite .prepare(`SELECT month FROM usage_rollups ORDER BY month`) .all() as Array<{ month: string }> diff --git a/packages/worker/src/app/retention.ts b/packages/worker/src/app/retention.ts index 423effa740..7db0133a1d 100644 --- a/packages/worker/src/app/retention.ts +++ b/packages/worker/src/app/retention.ts @@ -49,7 +49,6 @@ export const platformFeedbackRetentionDays = 365 export const publishedBundleArtifactRetentionDays = 30 export const emailDeliveryEventRetentionDays = 90 export const emailMessageRetentionDays = 365 -export const entitlementDailyCounterRetentionDays = 400 export const usageRollupRetentionMonths = 24 export const featureFlagExposureRetentionDays = 90 export const auditEventRetentionDays = 180 @@ -133,14 +132,6 @@ export const retentionPolicies: ReadonlyArray = [ description: 'Threads orphaned by email message retention are pruned for the affected users in the same run.', }, - { - table: 'entitlement_daily_counters', - scope: 'per-user', - retentionDays: entitlementDailyCounterRetentionDays, - batchSize: retentionDefaultBatchSize, - description: - 'Daily entitlement rate counters keep 400 days so year-over-year usage stays inspectable before rows age out.', - }, { table: 'usage_rollups', scope: 'per-user', @@ -210,7 +201,6 @@ export type RetentionPruneResult = { deletedAttachmentBlobs: number blobDeleteErrors: number } - entitlementDailyCounters: number usageRollups: number featureFlagExposureRollups: number auditEvents: number @@ -774,29 +764,6 @@ export async function pruneUserEmailMessagesForRetention(input: { return result } -export async function pruneEntitlementDailyCountersForRetention(input: { - db: D1Database - now?: Date - batchSize?: number -}) { - const cutoffDay = cutoffIso( - input.now ?? new Date(), - entitlementDailyCounterRetentionDays, - ).slice(0, 'YYYY-MM-DD'.length) - return selectAndDeleteByIds({ - db: input.db, - column: 'rowid', - bindings: [cutoffDay, input.batchSize ?? retentionDefaultBatchSize], - sql: `SELECT rowid - FROM entitlement_daily_counters - WHERE day < ? - ORDER BY day ASC, rowid ASC - LIMIT ?`, - table: 'entitlement_daily_counters', - idColumn: 'rowid', - }) -} - export async function pruneUsageRollupsForRetention(input: { db: D1Database now?: Date @@ -918,7 +885,6 @@ export async function pruneRetention(input: { deletedAttachmentBlobs: 0, blobDeleteErrors: 0, }, - entitlementDailyCounters: 0, usageRollups: 0, featureFlagExposureRollups: 0, auditEvents: 0, @@ -1011,13 +977,6 @@ export async function pruneRetention(input: { return batch.hasMore }, }, - countTask( - 'entitlement_daily_counters', - () => pruneEntitlementDailyCountersForRetention({ db, now }), - (count) => { - result.entitlementDailyCounters += count - }, - ), countTask( 'usage_rollups', () => pruneUsageRollupsForRetention({ db, now }), diff --git a/packages/worker/src/email/inbound-delivery.ts b/packages/worker/src/email/inbound-delivery.ts index f5499bf1fb..44726389c7 100644 --- a/packages/worker/src/email/inbound-delivery.ts +++ b/packages/worker/src/email/inbound-delivery.ts @@ -5,7 +5,6 @@ import { EntitlementLimitError, } from '#worker/entitlements/errors.ts' import { type PlanName } from '#worker/entitlements/plans.ts' -import { scheduleAbsoluteDailyEntitlementMirror } from '#worker/entitlements/service.ts' import { userMeterRpc, type UserMeterEnv, @@ -423,25 +422,11 @@ export async function pruneExpiredInboundDedupePointers(input: { ) } -async function readCounter(input: { +async function readSystemEmailDailyCounter(input: { db: D1Database - table: 'entitlement_daily_counters' | 'system_email_daily_counters' - userId?: string - localPart?: SystemEmailLocal + localPart: SystemEmailLocal day: string }) { - if (input.table === 'entitlement_daily_counters') { - const row = await input.db - .prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? - AND resource = 'email_receives_per_day' - AND day = ?`, - ) - .bind(input.userId, input.day) - .first<{ count: number }>() - return Number(row?.count ?? 0) - } const row = await input.db .prepare( `SELECT count FROM system_email_daily_counters @@ -452,7 +437,12 @@ async function readCounter(input: { return Number(row?.count ?? 0) } -/** Point-read today's user inbound receive count from UserMeter. */ +/** + * Point-read today's user inbound receive count from UserMeter. Cold meters + * initialize at zero; never touches the retired D1 daily counter table. + * + * `db` remains for call-site stability. + */ export async function readUserInboundReceiveCount(input: { db: D1Database env: UserMeterEnv @@ -460,6 +450,7 @@ export async function readUserInboundReceiveCount(input: { day: string now?: Date }) { + void input.db const now = input.now ?? new Date() const updatedAt = now.toISOString() const meter = userMeterRpc({ env: input.env, userId: input.userId }) @@ -469,16 +460,10 @@ export async function readUserInboundReceiveCount(input: { now: updatedAt, }) if (result.outcome === 'needs_bootstrap') { - const baseline = await readCounter({ - db: input.db, - table: 'entitlement_daily_counters', - userId: input.userId, - day: input.day, - }) await meter.initialize({ resource: 'email_receives_per_day', day: input.day, - count: baseline, + count: 0, updatedAt, }) result = await meter.read({ @@ -500,9 +485,8 @@ export async function readSystemInboundReceiveCount(input: { localPart: SystemEmailLocal day: string }) { - return await readCounter({ + return await readSystemEmailDailyCounter({ db: input.db, - table: 'system_email_daily_counters', localPart: input.localPart, day: input.day, }) @@ -521,7 +505,8 @@ async function resolveConcurrentDelivery(input: { /** * Charge one inbound receive via UserMeter, then persist the D1 delivery - * event. Delivery-id idempotency lives in the DO; D1 mirror is best-effort. + * event. Delivery-id idempotency lives in the DO. Never touches the retired + * D1 daily counter table. */ export async function chargeUserInboundDeliveryOnce(input: { db: D1Database @@ -530,8 +515,13 @@ export async function chargeUserInboundDeliveryOnce(input: { plan: PlanName limit: number now: Date + /** + * Retained for call-site stability. Daily counters no longer schedule D1 + * mirror work after mirror retirement. + */ waitUntil?: (promise: Promise) => void }): Promise { + void input.waitUntil const existing = await resolveConcurrentDelivery(input) if (existing) return existing @@ -548,16 +538,10 @@ export async function chargeUserInboundDeliveryOnce(input: { updatedAt, }) if (meterResult.outcome === 'needs_bootstrap') { - const baseline = await readCounter({ - db: input.db, - table: 'entitlement_daily_counters', - userId: input.delivery.userId, - day: input.delivery.quotaDay, - }) await meter.initialize({ resource: 'email_receives_per_day', day: input.delivery.quotaDay, - count: baseline, + count: 0, updatedAt, }) meterResult = await meter.consumeInboundDelivery({ @@ -585,16 +569,6 @@ export async function chargeUserInboundDeliveryOnce(input: { }) } - scheduleAbsoluteDailyEntitlementMirror({ - db: input.db, - userId: input.delivery.userId, - resource: meterResult.resource, - day: meterResult.day, - count: meterResult.count, - mirrorUpdatedAt: meterResult.mirrorUpdatedAt, - waitUntil: input.waitUntil, - }) - try { const insert = await input.db .prepare( @@ -635,9 +609,8 @@ export async function chargeSystemInboundDeliveryOnce(input: { }) { const existing = await resolveConcurrentDelivery(input) if (existing) return { delivery: existing, overLimit: false as const } - const current = await readCounter({ + const current = await readSystemEmailDailyCounter({ db: input.db, - table: 'system_email_daily_counters', localPart: input.localPart, day: input.delivery.quotaDay, }) diff --git a/packages/worker/src/email/inbound-entitlements.workers.test.ts b/packages/worker/src/email/inbound-entitlements.workers.test.ts index 53a9de2f46..fc9ca4746e 100644 --- a/packages/worker/src/email/inbound-entitlements.workers.test.ts +++ b/packages/worker/src/email/inbound-entitlements.workers.test.ts @@ -93,15 +93,6 @@ async function seedDailyReceiveCounter(userId: string, count: number) { updatedAt, ) }) - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, 'email_receives_per_day', ?, ?, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = excluded.count, - updated_at = excluded.updated_at`, - ) - .bind(userId, day, count, updatedAt) - .run() } async function readDailyReceiveCounter(userId: string) { @@ -888,19 +879,4 @@ test('inbound UserMeter replay keeps original claim day across UTC midnight', as expect( await meter.read({ resource: 'email_receives_per_day', day: dayN1 }), ).toMatchObject({ outcome: 'needs_bootstrap' }) - - const mirroredDayN = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, - ) - .bind(userId, dayN) - .first<{ count: number }>() - const mirroredDayN1 = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, - ) - .bind(userId, dayN1) - .first<{ count: number }>() - expect(Number(mirroredDayN?.count ?? 0)).toBe(1) - expect(mirroredDayN1).toBeNull() }, 30_000) diff --git a/packages/worker/src/email/inbound-spam-controls.workers.test.ts b/packages/worker/src/email/inbound-spam-controls.workers.test.ts index 09aade2a83..9d36ecb790 100644 --- a/packages/worker/src/email/inbound-spam-controls.workers.test.ts +++ b/packages/worker/src/email/inbound-spam-controls.workers.test.ts @@ -1,5 +1,6 @@ import { env } from 'cloudflare:workers' import { expect, test, vi } from 'vitest' +import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' import type * as PackageSubscriptionsModule from './package-subscriptions.ts' import { handleInboundEmail } from './inbound.ts' import { processInboundDeliveryEffects } from './inbound-effects.ts' @@ -68,13 +69,11 @@ async function seedVerifiedAccount(input: { } async function readUserDailyReceiveCount(userId: string) { - const row = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day' AND day = ?`, - ) - .bind(userId, new Date().toISOString().slice(0, 10)) - .first<{ count: number }>() - return Number(row?.count ?? 0) + const result = await userMeterRpc({ env, userId }).read({ + resource: 'email_receives_per_day', + day: new Date().toISOString().slice(0, 10), + }) + return result.outcome === 'ready' ? result.count : 0 } function dmarcFailAuthResults() { diff --git a/packages/worker/src/email/inbound.workers.test.ts b/packages/worker/src/email/inbound.workers.test.ts index 2505321ec9..e00a53a583 100644 --- a/packages/worker/src/email/inbound.workers.test.ts +++ b/packages/worker/src/email/inbound.workers.test.ts @@ -1,5 +1,6 @@ import { env } from 'cloudflare:workers' import { expect, test, vi } from 'vitest' +import { utcDayKey } from '@kody-internal/shared/date-keys.ts' import { userMeterRpc } from '#worker/entitlements/user-meter-client.ts' import { handleInboundEmail } from './inbound.ts' import { processInboundDeliveryEffects } from './inbound-effects.ts' @@ -576,13 +577,12 @@ test('inbound email handler rejects mail for unverified accounts', async () => { limit: 10, }) expect(messages).toEqual([]) - const counterRow = await env.APP_DB.prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = 'email_receives_per_day'`, - ) - .bind(userId) - .first<{ count: number }>() - expect(counterRow).toBeNull() + expect( + await userMeterRpc({ env, userId }).read({ + resource: 'email_receives_per_day', + day: utcDayKey(), + }), + ).toEqual({ outcome: 'needs_bootstrap' }) const events = await env.APP_DB.prepare( `SELECT event_type, detail_json FROM email_delivery_events WHERE user_id = ?`, ) @@ -1339,13 +1339,12 @@ test('pointer-only retry after midnight enforces the current quota day', async ( }) const receiveLimit = planLimits.free.maxEmailReceivesPerDay if (receiveLimit == null) throw new Error('Expected finite receive limit.') - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters ( - user_id, resource, day, count, updated_at - ) VALUES (?, 'email_receives_per_day', '2026-07-23', ?, ?)`, - ) - .bind(userId, receiveLimit, retryNow.toISOString()) - .run() + await userMeterRpc({ env, userId }).initialize({ + resource: 'email_receives_per_day', + day: '2026-07-23', + count: receiveLimit, + updatedAt: retryNow.toISOString(), + }) vi.setSystemTime(retryNow) const retry = createForwardableEmailMessage({ diff --git a/packages/worker/src/email/outbound.ts b/packages/worker/src/email/outbound.ts index ea135691fc..5b3c2d0684 100644 --- a/packages/worker/src/email/outbound.ts +++ b/packages/worker/src/email/outbound.ts @@ -588,8 +588,7 @@ export async function sendOutboundEmail( }) + attachmentBytesTotal, }) - // Atomic check-and-increment via UserMeter. No ExecutionContext here; - // the legacy D1 mirror uses the service's caught best-effort fallback. + // Atomic check-and-increment via UserMeter (sole daily authority). await consumeDailyEntitlement({ db: input.env.APP_DB, env: input.env, diff --git a/packages/worker/src/entitlements/entitlements.node.test.ts b/packages/worker/src/entitlements/entitlements.node.test.ts index f95a7c9a1a..b5c454dded 100644 --- a/packages/worker/src/entitlements/entitlements.node.test.ts +++ b/packages/worker/src/entitlements/entitlements.node.test.ts @@ -33,18 +33,8 @@ import { refundDailyEntitlement, } from './service.ts' import { utcDayKey } from '@kody-internal/shared/date-keys.ts' -import { - createInMemoryUserMeterEnv, - createWaitUntilDrain, -} from '#worker/test-support/user-meter.ts' - -type CounterRow = { - user_id: string - resource: string - day: string - count: number - updated_at?: string -} +import { createInMemoryUserMeterEnv } from '#worker/test-support/user-meter.ts' +import { userMeterRpc } from './user-meter-client.ts' function createEntitlementsTestDb( input: { @@ -71,12 +61,10 @@ function createEntitlementsTestDb( number > > - counters?: Array } = {}, ) { const users = input.users ?? [] const counts = input.counts ?? {} - const counters = input.counters ?? [] const queries: Array<{ sql: string; params: Array }> = [] const initialStorageBytes = Object.values(counts).reduce( (total, count) => total + (count ?? 0), @@ -180,15 +168,6 @@ function createEntitlementsTestDb( const bytes = storageBytesByUser.get(String(params[0])) return (bytes === undefined ? null : { bytes }) as T | null } - if (query.includes('FROM entitlement_daily_counters')) { - const row = counters.find( - (counter) => - counter.user_id === params[0] && - counter.resource === params[1] && - counter.day === params[2], - ) - return (row ? { count: row.count } : null) as T | null - } const count = countFor(query) if (count !== null) { return { count } as T @@ -210,66 +189,6 @@ function createEntitlementsTestDb( storageBytesByUser.set(userId, existing + Number(params[0])) return { meta: { changes: 1 } } } - if (query.includes('INSERT INTO entitlement_daily_counters')) { - const isAbsoluteMirror = - query.includes('count = excluded.count') && - query.includes('updated_at < excluded.updated_at') - const isConditionalConsume = query.includes('count + 1 <= ?') - const amount = isConditionalConsume ? 1 : Number(params[3]) - const mirrorUpdatedAt = String(params[4] ?? '') - const existing = counters.find( - (counter) => - counter.user_id === params[0] && - counter.resource === params[1] && - counter.day === params[2], - ) - if (existing) { - if ( - isConditionalConsume && - existing.count + 1 > Number(params[4]) - ) { - return { meta: { changes: 0 } } - } - if (isAbsoluteMirror) { - const existingUpdatedAt = existing.updated_at ?? '' - if ( - existingUpdatedAt !== '' && - !(existingUpdatedAt < mirrorUpdatedAt) - ) { - return { meta: { changes: 0 } } - } - existing.count = amount - existing.updated_at = mirrorUpdatedAt - } else { - existing.count += amount - } - } else { - counters.push({ - user_id: String(params[0]), - resource: String(params[1]), - day: String(params[2]), - count: amount, - ...(isAbsoluteMirror - ? { updated_at: mirrorUpdatedAt } - : {}), - }) - } - return { meta: { changes: 1 } } - } - if (query.includes('UPDATE entitlement_daily_counters')) { - // bind(updated_at, user_id, resource, day) - const existing = counters.find( - (counter) => - counter.user_id === params[1] && - counter.resource === params[2] && - counter.day === params[3], - ) - if (existing) { - existing.count = Math.max(0, existing.count - 1) - return { meta: { changes: 1 } } - } - return { meta: { changes: 0 } } - } throw new Error(`Unsupported run query: ${query}`) }, } @@ -278,7 +197,28 @@ function createEntitlementsTestDb( }, } as unknown as D1Database - return { db, counters, queries } + return { db, queries } +} + +async function readMeterDailyCount(input: { + env: ReturnType['env'] + userId: string + resource: + | 'email_sends_per_day' + | 'email_receives_per_day' + | 'execute_calls_per_day' + | 'outbound_fetches_per_day' + now: Date +}) { + const result = await userMeterRpc({ + env: input.env, + userId: input.userId, + }).read({ + resource: input.resource, + day: utcDayKey(input.now), + now: input.now.toISOString(), + }) + return result.outcome === 'ready' ? result.count : 0 } const plannedEmail = 'planned@example.com' @@ -770,11 +710,10 @@ test('persistent package services are gated as a 0/1 limit', async () => { test('plan user daily entitlements increment, enforce at limit, and reset on a new UTC day', async () => { const userId = await createStableUserIdFromEmail(plannedEmail) const now = new Date('2026-07-05T15:00:00.000Z') - const { db, counters } = createEntitlementsTestDb({ + const { db } = createEntitlementsTestDb({ users: [{ email: plannedEmail, plan: 'free', stable_user_id: userId }], }) const { env } = createInMemoryUserMeterEnv() - const mirror = createWaitUntilDrain() expect(utcDayKey(now)).toBe('2026-07-05') const limit = planLimits.free.maxEmailSendsPerDay @@ -787,19 +726,16 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n email: plannedEmail, resource: 'email_sends_per_day', now, - waitUntil: mirror.waitUntil, }) } - await mirror.drain() - expect(counters).toEqual([ - { - user_id: userId, + expect( + await readMeterDailyCount({ + env, + userId, resource: 'email_sends_per_day', - day: '2026-07-05', - count: limit, - updated_at: expect.stringMatching(/^r\/0+[0-9]+$/), - }, - ]) + now, + }), + ).toBe(limit) await expect( assertWithinEntitlement({ db, @@ -808,7 +744,7 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n resource: 'email_sends_per_day', now, }), - ).rejects.toBeInstanceOf(EntitlementLimitError) + ).rejects.toThrow(/must be read from UserMeter/) const denied = await consumeDailyEntitlement({ db, @@ -817,7 +753,6 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n email: plannedEmail, resource: 'email_sends_per_day', now, - waitUntil: mirror.waitUntil, }).then( () => null, (thrown: unknown) => thrown, @@ -832,7 +767,6 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n limit, current: limit, }) - expect(counters[0]?.count).toBe(limit) const nextDay = new Date('2026-07-06T00:00:01.000Z') await consumeDailyEntitlement({ @@ -842,17 +776,28 @@ test('plan user daily entitlements increment, enforce at limit, and reset on a n email: plannedEmail, resource: 'email_sends_per_day', now: nextDay, - waitUntil: mirror.waitUntil, }) - await mirror.drain() - expect(counters).toHaveLength(2) - expect(counters[1]?.count).toBe(1) + expect( + await readMeterDailyCount({ + env, + userId, + resource: 'email_sends_per_day', + now: nextDay, + }), + ).toBe(1) + expect( + await readMeterDailyCount({ + env, + userId, + resource: 'email_sends_per_day', + now, + }), + ).toBe(limit) }) test('refundDailyEntitlement decrements the user/day counter and floors at zero', async () => { - const { db, counters } = createEntitlementsTestDb() + const { db } = createEntitlementsTestDb() const { env } = createInMemoryUserMeterEnv() - const mirror = createWaitUntilDrain() const now = new Date('2026-07-05T15:00:00.000Z') for (let index = 0; index < 2; index += 1) { await consumeDailyEntitlement({ @@ -862,7 +807,6 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' email: null, resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) } for (let index = 0; index < 3; index += 1) { @@ -873,30 +817,30 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' email: null, resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) } - await mirror.drain() await refundDailyEntitlement({ db, env, userId: 'user-1', resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) - await mirror.drain() expect( - counters.find( - (row) => - row.user_id === 'user-1' && row.resource === 'email_receives_per_day', - )?.count, + await readMeterDailyCount({ + env, + userId: 'user-1', + resource: 'email_receives_per_day', + now, + }), ).toBe(1) expect( - counters.find( - (row) => - row.user_id === 'user-2' && row.resource === 'email_receives_per_day', - )?.count, + await readMeterDailyCount({ + env, + userId: 'user-2', + resource: 'email_receives_per_day', + now, + }), ).toBe(3) await refundDailyEntitlement({ @@ -905,7 +849,6 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' userId: 'user-1', resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) await refundDailyEntitlement({ db, @@ -913,21 +856,20 @@ test('refundDailyEntitlement decrements the user/day counter and floors at zero' userId: 'user-1', resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) - await mirror.drain() expect( - counters.find( - (row) => - row.user_id === 'user-1' && row.resource === 'email_receives_per_day', - )?.count, + await readMeterDailyCount({ + env, + userId: 'user-1', + resource: 'email_receives_per_day', + now, + }), ).toBe(0) }) test('missing-email lookups fail closed and honor free email caps', async () => { - const { db, counters } = createEntitlementsTestDb() + const { db } = createEntitlementsTestDb() const { env } = createInMemoryUserMeterEnv() - const mirror = createWaitUntilDrain() const sendLimit = planLimits.free.maxEmailSendsPerDay const now = new Date('2026-07-05T15:00:00.000Z') for (let index = 0; index < sendLimit; index += 1) { @@ -938,11 +880,16 @@ test('missing-email lookups fail closed and honor free email caps', async () => email: null, resource: 'email_sends_per_day', now, - waitUntil: mirror.waitUntil, }) } - await mirror.drain() - expect(counters[0]?.count).toBe(sendLimit) + expect( + await readMeterDailyCount({ + env, + userId: 'user-1', + resource: 'email_sends_per_day', + now, + }), + ).toBe(sendLimit) await expect( consumeDailyEntitlement({ db, @@ -951,7 +898,6 @@ test('missing-email lookups fail closed and honor free email caps', async () => email: null, resource: 'email_sends_per_day', now, - waitUntil: mirror.waitUntil, }), ).rejects.toBeInstanceOf(EntitlementLimitError) @@ -964,12 +910,15 @@ test('missing-email lookups fail closed and honor free email caps', async () => email: null, resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }) } - await mirror.drain() expect( - counters.find((row) => row.resource === 'email_receives_per_day')?.count, + await readMeterDailyCount({ + env, + userId: 'user-1', + resource: 'email_receives_per_day', + now, + }), ).toBe(receiveLimit) const denied = await consumeDailyEntitlement({ @@ -979,7 +928,6 @@ test('missing-email lookups fail closed and honor free email caps', async () => email: null, resource: 'email_receives_per_day', now, - waitUntil: mirror.waitUntil, }).then( () => null, (thrown: unknown) => thrown, diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index 5d3a5e6ae5..27e9591435 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -204,37 +204,6 @@ export async function findUserAccountByStableUserId( } } -/** - * Legacy D1 mirror write for a daily entitlement counter. Expand-phase - * enforcement is authoritative in UserMeter; this helper remains for tests, - * backfills, and the best-effort mirror scheduled after DO consume/refund. - */ -export async function incrementDailyEntitlementCounter(input: { - db: D1Database - userId: string - resource: EntitlementResource - amount?: number - now?: Date -}) { - const now = input.now ?? new Date() - await input.db - .prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = entitlement_daily_counters.count + excluded.count, - updated_at = excluded.updated_at`, - ) - .bind( - input.userId, - input.resource, - utcDayKey(now), - input.amount ?? 1, - now.toISOString(), - ) - .run() -} - function assertDailyEntitlementResource( resource: EntitlementResource, ): DailyEntitlementResource { @@ -246,119 +215,33 @@ function assertDailyEntitlementResource( return resource } -// Best-effort mirror scheduling: prefer `waitUntil` when available; otherwise -// catch so the promise cannot surface as unhandled. Never awaited by callers. -function scheduleDailyEntitlementMirror( - work: Promise, - waitUntil?: (promise: Promise) => void, -) { - const tracked = work.catch((error: unknown) => { - console.warn('entitlement-daily-mirror-failed', error) - }) - if (waitUntil) { - waitUntil(tracked) - return - } - void tracked -} - -// Absolute D1 mirror ordered by DO-minted `mirrorUpdatedAt` (revision-primary) -// so a late older write cannot overwrite newer state, including refunds. -async function mirrorDailyEntitlementAbsoluteCount(input: { - db: D1Database - userId: string - resource: DailyEntitlementResource - day: string - count: number - mirrorUpdatedAt: string -}) { - await input.db - .prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = excluded.count, - updated_at = excluded.updated_at - WHERE entitlement_daily_counters.updated_at < excluded.updated_at`, - ) - .bind( - input.userId, - input.resource, - input.day, - input.count, - input.mirrorUpdatedAt, - ) - .run() -} - -/** Best-effort absolute D1 mirror of UserMeter daily counter state. */ -export function scheduleAbsoluteDailyEntitlementMirror(input: { - db: D1Database - userId: string - resource: EntitlementResource - day: string - count: number - mirrorUpdatedAt: string - waitUntil?: (promise: Promise) => void -}): void { - const resource = assertDailyEntitlementResource(input.resource) - scheduleDailyEntitlementMirror( - mirrorDailyEntitlementAbsoluteCount({ - db: input.db, - userId: input.userId, - resource, - day: input.day, - count: input.count, - mirrorUpdatedAt: input.mirrorUpdatedAt, - }), - input.waitUntil, - ) -} - -async function ensureUserMeterCounterInitialized(input: { - db: D1Database +/** + * Seed a missing UserMeter `(resource, day)` at zero. `INSERT OR IGNORE` + * inside the DO keeps concurrent cold callers safe. + */ +async function ensureUserMeterCounterInitializedAtZero(input: { env: UserMeterEnv userId: string resource: DailyEntitlementResource day: string updatedAt: string - now: Date }) { - const baseline = await readDailyEntitlementCounter({ - db: input.db, - userId: input.userId, - resource: input.resource, - now: input.now, - }) const meter = userMeterRpc({ env: input.env, userId: input.userId }) await meter.initialize({ resource: input.resource, day: input.day, - count: baseline, + count: 0, updatedAt: input.updatedAt, }) } -async function readDailyEntitlementCounter(input: { - db: D1Database - userId: string - resource: EntitlementResource - now: Date -}) { - const row = await input.db - .prepare( - `SELECT count FROM entitlement_daily_counters - WHERE user_id = ? AND resource = ? AND day = ?`, - ) - .bind(input.userId, input.resource, utcDayKey(input.now)) - .first<{ count: number }>() - return Number(row?.count ?? 0) -} - /** * Point-read one daily entitlement counter from UserMeter. Cold meters - * bootstrap once from the legacy D1 point row, then re-read; warm meters - * return the DO count without touching D1. + * initialize the `(resource, day)` at zero, then re-read; warm meters return + * the DO count. Never touches D1 daily counter state. + * + * `db` remains on the signature for call-site stability; plan/account lookups + * on other paths still need APP_DB. */ export async function readDailyEntitlementResourceUsage(input: { db: D1Database @@ -367,6 +250,7 @@ export async function readDailyEntitlementResourceUsage(input: { resource: EntitlementResource now?: Date }): Promise { + void input.db const resource = assertDailyEntitlementResource(input.resource) const now = input.now ?? new Date() const day = utcDayKey(now) @@ -378,14 +262,12 @@ export async function readDailyEntitlementResourceUsage(input: { now: updatedAt, }) if (result.outcome === 'needs_bootstrap') { - await ensureUserMeterCounterInitialized({ - db: input.db, + await ensureUserMeterCounterInitializedAtZero({ env: input.env, userId: input.userId, resource, day, updatedAt, - now, }) result = await meter.read({ resource, @@ -914,12 +796,12 @@ export async function readEntitlementResourceUsage(input: { case 'email_receives_per_day': case 'execute_calls_per_day': case 'outbound_fetches_per_day': - return await readDailyEntitlementCounter({ - db, - userId, - resource, - now, - }) + // Authoritative daily counters live in UserMeter. Callers must use + // consumeDailyEntitlement / readDailyEntitlementResourceUsage / + // readCurrentEntitlementResourceUsage — the retired D1 mirror is gone. + throw new Error( + `${resource} usage must be read from UserMeter (use readDailyEntitlementResourceUsage or readCurrentEntitlementResourceUsage).`, + ) case 'stored_email_messages': return await countRows( db, @@ -1161,25 +1043,29 @@ export async function assertWithinEntitlement( export type ConsumeDailyEntitlementInput = { db: D1Database - /** Must expose `USER_METER` (authoritative); D1 is a best-effort mirror. */ + /** Must expose `USER_METER` (sole daily counter authority). */ env: UserMeterEnv userId: string email: string | null | undefined resource: EntitlementResource now?: Date - /** Prefer `ctx.waitUntil` for the legacy D1 mirror; never awaited here. */ + /** + * Retained for call-site stability. Daily counters no longer schedule D1 + * mirror work; unused after mirror retirement. + */ waitUntil?: (promise: Promise) => void } /** * Atomically consume one daily entitlement unit via UserMeter, throwing * EntitlementLimitError when the plan limit would be exceeded. Cold keys - * bootstrap once from legacy D1; warm path awaits only the DO RPC. Mirror - * writes are best-effort and cannot affect enforcement. + * initialize at zero (concurrent-safe); warm path awaits only the DO RPC. + * Never touches the retired D1 daily counter table. */ export async function consumeDailyEntitlement( input: ConsumeDailyEntitlementInput, ): Promise { + void input.waitUntil const resource = assertDailyEntitlementResource(input.resource) const now = input.now ?? new Date() const day = utcDayKey(now) @@ -1200,14 +1086,12 @@ export async function consumeDailyEntitlement( updatedAt, }) if (result.outcome === 'needs_bootstrap') { - await ensureUserMeterCounterInitialized({ - db: input.db, + await ensureUserMeterCounterInitializedAtZero({ env: input.env, userId: input.userId, resource, day, updatedAt, - now, }) result = await meter.consume({ resource, @@ -1230,17 +1114,6 @@ export async function consumeDailyEntitlement( upgradeHint: buildEntitlementUpgradeHint(resource), }) } - scheduleDailyEntitlementMirror( - mirrorDailyEntitlementAbsoluteCount({ - db: input.db, - userId: input.userId, - resource, - day, - count: result.count, - mirrorUpdatedAt: result.mirrorUpdatedAt, - }), - input.waitUntil, - ) } export type RefundDailyEntitlementInput = { @@ -1249,38 +1122,31 @@ export type RefundDailyEntitlementInput = { userId: string resource: EntitlementResource now?: Date + /** + * Retained for call-site stability. Daily counters no longer schedule D1 + * mirror work; unused after mirror retirement. + */ waitUntil?: (promise: Promise) => void } /** * Atomically refund one previously consumed daily entitlement unit in * UserMeter (floors at zero). Pass the same `now` (day key) as the matching - * consume. The legacy D1 mirror is best-effort and never awaited. + * consume. Never touches the retired D1 daily counter table. */ export async function refundDailyEntitlement( input: RefundDailyEntitlementInput, ): Promise { + void input.db + void input.waitUntil const resource = assertDailyEntitlementResource(input.resource) const now = input.now ?? new Date() const day = utcDayKey(now) const updatedAt = now.toISOString() const meter = userMeterRpc({ env: input.env, userId: input.userId }) - const result = await meter.refund({ + await meter.refund({ resource, day, updatedAt, }) - // revision 0 means the key was never initialized; nothing to mirror. - if (result.revision < 1) return - scheduleDailyEntitlementMirror( - mirrorDailyEntitlementAbsoluteCount({ - db: input.db, - userId: input.userId, - resource, - day, - count: result.count, - mirrorUpdatedAt: result.mirrorUpdatedAt, - }), - input.waitUntil, - ) } diff --git a/packages/worker/src/entitlements/test-schema.ts b/packages/worker/src/entitlements/test-schema.ts index a032c1a79a..938f01a5f3 100644 --- a/packages/worker/src/entitlements/test-schema.ts +++ b/packages/worker/src/entitlements/test-schema.ts @@ -4,8 +4,9 @@ import { ensureUsersTestSchema } from '#worker/users-test-schema.ts' /** * Non-destructive schema for entitlement primitives in workers-unit tests, * where the D1 database starts empty and each suite provisions the tables it - * needs. Mirrors migrations 0048 (entitlement_daily_counters) and 0066 (Stripe - * billing columns) on top of the shared `users` schema. + * needs. Mirrors Stripe billing columns (0066) and package-service state on + * top of the shared `users` schema. Code no longer reads or writes the D1 + * `entitlement_daily_counters` mirror; daily counters live only in UserMeter. * * Also provisions `user_storage_buckets` because entitlement suites that touch * StorageRunner writes register ownership through that table. @@ -20,18 +21,6 @@ export async function ensureEntitlementTestSchema(db: D1Database) { 'stripe_plan_refreshed_at', ], }) - await db - .prepare( - `CREATE TABLE IF NOT EXISTS entitlement_daily_counters ( - user_id TEXT NOT NULL, - resource TEXT NOT NULL, - day TEXT NOT NULL, - count INTEGER NOT NULL DEFAULT 0, - updated_at TEXT NOT NULL, - PRIMARY KEY (user_id, resource, day) -)`, - ) - .run() await db .prepare( `CREATE TABLE IF NOT EXISTS package_service_states ( diff --git a/packages/worker/src/entitlements/user-meter-do.ts b/packages/worker/src/entitlements/user-meter-do.ts index 4a3eff0eac..21317b1a19 100644 --- a/packages/worker/src/entitlements/user-meter-do.ts +++ b/packages/worker/src/entitlements/user-meter-do.ts @@ -829,7 +829,7 @@ class UserMeterBase extends DurableObject { } } - /** Cold-key seed from legacy D1; INSERT OR IGNORE is concurrency-safe. */ + /** Cold-key seed (zero after D1 mirror retirement); INSERT OR IGNORE is concurrency-safe. */ async initialize(input: { resource: string day: string diff --git a/packages/worker/src/entitlements/user-meter.workers.test.ts b/packages/worker/src/entitlements/user-meter.workers.test.ts index 823adf7f6d..ac79978f38 100644 --- a/packages/worker/src/entitlements/user-meter.workers.test.ts +++ b/packages/worker/src/entitlements/user-meter.workers.test.ts @@ -11,6 +11,7 @@ import { planLimits } from './plans.ts' import { assertWithinStorageBytesEntitlement, consumeDailyEntitlement, + readDailyEntitlementResourceUsage, readUserD1StorageBytes, refundDailyEntitlement, } from './service.ts' @@ -36,22 +37,6 @@ async function seedFreeUser(emailPrefix: string) { return { email, userId } } -async function readD1DailyCount(input: { - userId: string - resource: string - day: string -}) { - const row = await env.APP_DB.prepare( - `SELECT count, updated_at FROM entitlement_daily_counters - WHERE user_id = ? AND resource = ? AND day = ?`, - ) - .bind(input.userId, input.resource, input.day) - .first<{ count: number; updated_at: string }>() - return row - ? { count: Number(row.count), updatedAt: String(row.updated_at) } - : null -} - async function waitFor( predicate: () => boolean, timeoutMs = 5_000, @@ -66,60 +51,24 @@ async function waitFor( } } -const mirrorSql = `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT(user_id, resource, day) DO UPDATE SET - count = excluded.count, - updated_at = excluded.updated_at - WHERE entitlement_daily_counters.updated_at < excluded.updated_at` - -test('UserMeter bootstraps from legacy D1 once, then warms without APP_DB awaits', async () => { +test('cold daily consume initializes at zero without D1 prepare/run and first unit is 1', async () => { const now = new Date('2026-07-31T15:00:00.000Z') const day = utcDayKey(now) - const sendLimit = planLimits.free.maxEmailSendsPerDay - const user = await seedFreeUser('meter-bootstrap') - const drain = createWaitUntilDrain() + const user = await seedFreeUser('meter-cold-zero') const meter = userMeterRpc({ env, userId: user.userId }) - const legacyCount = sendLimit - 1 - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?)`, - ) - .bind( - user.userId, - 'email_sends_per_day', - day, - legacyCount, - '2026-07-31T12:00:00.000Z', - ) - .run() - - let d1PointReads = 0 + let dailyPrepareCalls = 0 using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { return ((query: string) => { - const statement = originalPrepare(query) - const isPointRead = - query.includes('SELECT count FROM entitlement_daily_counters') && - query.includes('WHERE user_id = ? AND resource = ? AND day = ?') - const originalBind = statement.bind.bind(statement) - return { - bind(...params: Array) { - const bound = originalBind(...params) - return { - first: async () => { - if (isPointRead) d1PointReads += 1 - return await bound.first() - }, - all: bound.all.bind(bound), - raw: bound.raw?.bind(bound), - run: bound.run.bind(bound), - } - }, - } + if (query.includes('entitlement_daily_counters')) dailyPrepareCalls += 1 + return originalPrepare(query) }) as D1Database['prepare'] }) + expect(await meter.read({ resource: 'email_sends_per_day', day })).toEqual({ + outcome: 'needs_bootstrap', + }) + await consumeDailyEntitlement({ db: env.APP_DB, env, @@ -127,285 +76,127 @@ test('UserMeter bootstraps from legacy D1 once, then warms without APP_DB awaits email: user.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) - expect(d1PointReads).toBe(1) expect( await meter.read({ resource: 'email_sends_per_day', day }), ).toMatchObject({ outcome: 'ready', - count: sendLimit, + count: 1, }) + expect(dailyPrepareCalls).toBe(0) +}, 30_000) - const denied = await consumeDailyEntitlement({ - db: env.APP_DB, - env, - userId: user.userId, - email: user.email, +test('warm daily consume/read never prepares entitlement_daily_counters', async () => { + const now = new Date('2026-07-31T15:00:00.000Z') + const day = utcDayKey(now) + const user = await seedFreeUser('meter-warm-no-d1') + const meter = userMeterRpc({ env, userId: user.userId }) + await meter.initialize({ resource: 'email_sends_per_day', - now, - waitUntil: drain.waitUntil, - }).then( - () => null, - (thrown: unknown) => thrown, - ) - expect(denied).toBeInstanceOf(EntitlementLimitError) - expect(denied).toMatchObject({ - details: { - resource: 'email_sends_per_day', - plan: 'free', - limit: sendLimit, - current: sendLimit, - }, + day, + count: 0, + updatedAt: now.toISOString(), }) - expect(d1PointReads).toBe(1) - await refundDailyEntitlement({ + let dailyPrepareCalls = 0 + using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + if (query.includes('entitlement_daily_counters')) dailyPrepareCalls += 1 + return originalPrepare(query) + }) as D1Database['prepare'] + }) + + await consumeDailyEntitlement({ db: env.APP_DB, env, userId: user.userId, + email: user.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, - }) - await drain.drain() - d1PointReads = 0 - let nonMirrorAppDbRuns = 0 - let releaseMirror!: () => void - const mirrorGate = new Promise((resolve) => { - releaseMirror = resolve - }) - let mirrorRunStarted = false - using _warmPatch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { - return ((query: string) => { - const statement = originalPrepare(query) - const isMirrorWrite = - query.includes('INSERT INTO entitlement_daily_counters') && - query.includes('updated_at < excluded.updated_at') - const isPointRead = - query.includes('SELECT count FROM entitlement_daily_counters') && - query.includes('WHERE user_id = ? AND resource = ? AND day = ?') - const originalBind = statement.bind.bind(statement) - return { - bind(...params: Array) { - const bound = originalBind(...params) - return { - first: async () => { - if (isPointRead) d1PointReads += 1 - return await bound.first() - }, - all: bound.all.bind(bound), - raw: bound.raw?.bind(bound), - run: async () => { - if (isMirrorWrite) { - mirrorRunStarted = true - await mirrorGate - return await bound.run() - } - nonMirrorAppDbRuns += 1 - return await bound.run() - }, - } - }, - } - }) as D1Database['prepare'] }) - try { - await expect( - consumeDailyEntitlement({ - db: env.APP_DB, - env, - userId: user.userId, - email: user.email, - resource: 'email_sends_per_day', - now, - waitUntil: drain.waitUntil, - }), - ).resolves.toBeUndefined() - expect(mirrorRunStarted).toBe(true) - expect(d1PointReads).toBe(0) - expect(nonMirrorAppDbRuns).toBe(0) - } finally { - releaseMirror() - } - await drain.drain() + await expect( + readDailyEntitlementResourceUsage({ + db: env.APP_DB, + env, + userId: user.userId, + resource: 'email_sends_per_day', + now, + }), + ).resolves.toBe(1) + expect(dailyPrepareCalls).toBe(0) }, 30_000) -test('UserMeter mirror tokens reject out-of-order absolute D1 writes', async () => { - const now = new Date('2026-07-31T16:00:00.000Z') - const day = utcDayKey(now) - const user = await seedFreeUser('meter-mirror-order') +test('next UTC day cold consume starts at zero independently', async () => { + const dayOne = new Date('2026-07-31T15:00:00.000Z') + const dayTwo = new Date('2026-08-01T01:00:00.000Z') + const user = await seedFreeUser('meter-next-day') const meter = userMeterRpc({ env, userId: user.userId }) - await meter.initialize({ - resource: 'email_sends_per_day', - day, - count: 0, - updatedAt: now.toISOString(), - }) - const first = await meter.consume({ - resource: 'email_sends_per_day', - day, - limit: 100, - updatedAt: now.toISOString(), - }) - expect(first).toMatchObject({ - outcome: 'ready', - consumed: true, - count: 1, - revision: 2, - }) - const second = await meter.consume({ - resource: 'email_sends_per_day', - day, - limit: 100, - updatedAt: now.toISOString(), - }) - expect(second).toMatchObject({ - outcome: 'ready', - consumed: true, - count: 2, - revision: 3, - }) - const refunded = await meter.refund({ + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, resource: 'email_sends_per_day', - day, - updatedAt: now.toISOString(), + now: dayOne, }) - expect(refunded).toMatchObject({ - outcome: 'ready', - count: 1, - revision: 4, - }) - if (first.outcome !== 'ready') { - throw new Error('Expected ready consume results.') - } + expect( + await meter.read({ + resource: 'email_sends_per_day', + day: utcDayKey(dayOne), + }), + ).toMatchObject({ outcome: 'ready', count: 1 }) - type PendingMirror = { - count: number - mirrorUpdatedAt: string - resolve: () => void - promise: Promise - } - const pending: Array = [] - function deferMirror(count: number, mirrorUpdatedAt: string) { - let resolve!: () => void - const promise = new Promise((res) => { - resolve = res - }) - pending.push({ count, mirrorUpdatedAt, resolve, promise }) - return promise - } + expect( + await meter.read({ + resource: 'email_sends_per_day', + day: utcDayKey(dayTwo), + now: dayTwo.toISOString(), + }), + ).toEqual({ outcome: 'needs_bootstrap' }) + let dailyPrepareCalls = 0 using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { return ((query: string) => { - const statement = originalPrepare(query) - const isMirrorWrite = - query.includes('INSERT INTO entitlement_daily_counters') && - query.includes('updated_at < excluded.updated_at') - const originalBind = statement.bind.bind(statement) - return { - bind(...params: Array) { - const bound = originalBind(...params) - return { - first: bound.first.bind(bound), - all: bound.all.bind(bound), - raw: bound.raw?.bind(bound), - run: async () => { - if (!isMirrorWrite) return await bound.run() - await deferMirror(Number(params[3]), String(params[4])) - return await bound.run() - }, - } - }, - } + if (query.includes('entitlement_daily_counters')) dailyPrepareCalls += 1 + return originalPrepare(query) }) as D1Database['prepare'] }) - const older = env.APP_DB.prepare(mirrorSql) - .bind( - user.userId, - 'email_sends_per_day', - day, - first.count, - first.mirrorUpdatedAt, - ) - .run() - const newer = env.APP_DB.prepare(mirrorSql) - .bind( - user.userId, - 'email_sends_per_day', - day, - refunded.count, - refunded.mirrorUpdatedAt, - ) - .run() - - await waitFor(() => pending.length === 2, 5_000, 'mirror gates') - expect(pending.map((entry) => entry.mirrorUpdatedAt)).toEqual([ - userMeterMirrorUpdatedAtToken(2), - userMeterMirrorUpdatedAtToken(4), - ]) - const [olderGate, newerGate] = pending - newerGate!.resolve() - await newer - expect( - await readD1DailyCount({ - userId: user.userId, - resource: 'email_sends_per_day', - day, - }), - ).toEqual({ - count: 1, - updatedAt: userMeterMirrorUpdatedAtToken(4), + await consumeDailyEntitlement({ + db: env.APP_DB, + env, + userId: user.userId, + email: user.email, + resource: 'email_sends_per_day', + now: dayTwo, }) - - olderGate!.resolve() - await older expect( - await readD1DailyCount({ - userId: user.userId, + await meter.read({ resource: 'email_sends_per_day', - day, + day: utcDayKey(dayTwo), }), - ).toEqual({ - count: 1, - updatedAt: userMeterMirrorUpdatedAtToken(4), - }) - expect( - await meter.read({ resource: 'email_sends_per_day', day }), - ).toMatchObject({ outcome: 'ready', count: 1, revision: 4 }) + ).toMatchObject({ outcome: 'ready', count: 1 }) + expect(dailyPrepareCalls).toBe(0) }, 30_000) -test('UserMeter daily entitlement consume/refund/read/export/purge workflow is per-user and mirror-safe', async () => { +test('UserMeter daily entitlement consume/refund/read/export/purge workflow is per-user without D1 daily table', async () => { const now = new Date('2026-07-31T15:00:00.000Z') const day = utcDayKey(now) const sendLimit = planLimits.free.maxEmailSendsPerDay const userA = await seedFreeUser('meter-a') const userB = await seedFreeUser('meter-b') - const drain = createWaitUntilDrain() const meterA = userMeterRpc({ env, userId: userA.userId }) const meterB = userMeterRpc({ env, userId: userB.userId }) expect(userMeterDurableObjectName(userA.userId)).toBe(userA.userId) - await meterA.initialize({ - resource: 'email_sends_per_day', - day, - count: 0, - updatedAt: now.toISOString(), - }) - await expect( - meterA.consume({ - resource: 'email_sends_per_day', - day, - limit: sendLimit, - updatedAt: now.toISOString(), - }), - ).resolves.toMatchObject({ - outcome: 'ready', - consumed: true, - count: 1, + let dailyPrepareCalls = 0 + using _patch = withPatchedDbPrepare(env.APP_DB, (originalPrepare) => { + return ((query: string) => { + if (query.includes('entitlement_daily_counters')) dailyPrepareCalls += 1 + return originalPrepare(query) + }) as D1Database['prepare'] }) await consumeDailyEntitlement({ @@ -415,14 +206,12 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userA.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) - await drain.drain() expect( await meterA.read({ resource: 'email_sends_per_day', day }), - ).toMatchObject({ outcome: 'ready', count: 2 }) + ).toMatchObject({ outcome: 'ready', count: 1 }) - for (let index = 2; index < sendLimit; index += 1) { + for (let index = 1; index < sendLimit; index += 1) { await consumeDailyEntitlement({ db: env.APP_DB, env, @@ -430,7 +219,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userA.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) } const concurrent = await Promise.all( @@ -442,7 +230,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userA.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }).then( () => null, (thrown: unknown) => thrown, @@ -454,17 +241,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p (result) => result instanceof EntitlementLimitError, ) expect(denials).toHaveLength(8) - for (const denial of denials) { - expect(denial).toMatchObject({ - details: { - code: 'entitlement_limit_exceeded', - resource: 'email_sends_per_day', - plan: 'free', - limit: sendLimit, - current: sendLimit, - }, - }) - } await consumeDailyEntitlement({ db: env.APP_DB, @@ -473,7 +249,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userB.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) expect( await meterB.read({ resource: 'email_sends_per_day', day }), @@ -488,7 +263,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p userId: userA.userId, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) expect( await meterA.read({ resource: 'email_sends_per_day', day }), @@ -500,7 +274,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p userId: userA.userId, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) } expect( @@ -514,7 +287,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userA.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }) await consumeDailyEntitlement({ db: env.APP_DB, @@ -523,7 +295,6 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p email: userA.email, resource: 'execute_calls_per_day', now, - waitUntil: drain.waitUntil, }) const exportedA = await meterA.exportCounters({}) expect(exportedA.counters).toEqual( @@ -574,53 +345,20 @@ test('UserMeter daily entitlement consume/refund/read/export/purge workflow is p await meterB.read({ resource: 'email_sends_per_day', day }), ).toMatchObject({ outcome: 'ready', count: 1 }) - await env.APP_DB.prepare( - `DELETE FROM entitlement_daily_counters - WHERE user_id = ? AND resource = ? AND day = ?`, - ) - .bind(userA.userId, 'email_sends_per_day', day) - .run() - - consoleWarn.mockImplementation(() => {}) - const failingMirrorDb = { - prepare(query: string) { - if ( - query.includes('INSERT INTO entitlement_daily_counters') && - query.includes('updated_at < excluded.updated_at') - ) { - return { - bind() { - return { - async run() { - throw new Error('mirror write failed') - }, - } - }, - } - } - return env.APP_DB.prepare(query) - }, - } as unknown as D1Database - await expect( consumeDailyEntitlement({ - db: failingMirrorDb, + db: env.APP_DB, env, userId: userA.userId, email: userA.email, resource: 'email_sends_per_day', now, - waitUntil: drain.waitUntil, }), ).resolves.toBeUndefined() - await drain.drain() expect( await meterA.read({ resource: 'email_sends_per_day', day }), ).toMatchObject({ outcome: 'ready', count: 1 }) - expect(consoleWarn).toHaveBeenCalledWith( - 'entitlement-daily-mirror-failed', - expect.any(Error), - ) + expect(dailyPrepareCalls).toBe(0) const stub = env.USER_METER.get( env.USER_METER.idFromName(userMeterDurableObjectName(userA.userId)), diff --git a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts index 7fa3dabbc9..c06cf8f4d6 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-capabilities.node.test.ts @@ -285,9 +285,6 @@ function createAdminCapabilityTestDb(input: { raw_mime_key: message.raw_mime_key ?? null, } as T } - if (normalizedQuery.includes('from entitlement_daily_counters')) { - return null - } const count = countForQuery(normalizedQuery) if (count !== null) return { count } as T return null diff --git a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts index ceccd65157..f14c29009c 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.node.test.ts @@ -138,6 +138,7 @@ test('admin_user_meter_parity returns null for missing users and omits lease sec ctx, ) expect(present.report?.stableUserId).toBe(stableUserId) + expect(present.report?.daily.mirrorRetired).toBe(false) expect(present.report?.deletion.mirrorLeaseParity).toBe(true) expect(mockModule.logAuditEvent).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts index 2d40113791..d20b563b1f 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-user-meter-parity.ts @@ -13,7 +13,7 @@ const dailyResourceSchema = z.enum(dailyEntitlementResources) const dailyResourceParitySchema = z.object({ resource: dailyResourceSchema, - d1Count: z.number().int().nonnegative(), + d1Count: z.number().int().nonnegative().nullable(), meterCount: z.number().int().nonnegative().nullable(), needsBootstrap: z.boolean(), delta: z.number().int().nullable(), @@ -76,6 +76,11 @@ const reportSchema = z.object({ stableUserId: stableUserIdSchema, daily: z.object({ day: z.string(), + mirrorRetired: z + .boolean() + .describe( + 'True when D1 entitlement_daily_counters is absent. Daily gate then reports meter counts only; d1Count/delta are null and parity stays true (no D1 comparison).', + ), resources: z.array(dailyResourceParitySchema), mismatchCount: z.number().int().nonnegative(), }), @@ -98,7 +103,7 @@ export const adminUserMeterParityCapability = defineDomainCapability( ...adminCapabilityAccess, name: 'admin_user_meter_parity', description: - 'Read-only production verification report comparing one user D1 entitlement mirrors against UserMeter for daily counters, storage bytes, package-service liveness, and deletion write leases. Never bootstraps or writes parity state (DO constructor schema maintenance and stale daily pruning may still run); returns counts and parity only (no lease tokens/holders or email content). Admin-only.', + 'Read-only production verification report for one user: UserMeter daily counters plus D1↔UserMeter parity for storage bytes, package-service liveness, and deletion write leases. While `entitlement_daily_counters` exists, daily rows compare D1 mirror counts (`mirrorRetired: false`); after the drop migration, daily rows report meter counts only (`mirrorRetired: true`). Never bootstraps or writes parity state (DO constructor schema maintenance and stale daily pruning may still run); returns counts and parity only (no lease tokens/holders or email content). Admin-only.', keywords: [ 'admin', 'user meter', @@ -111,6 +116,7 @@ export const adminUserMeterParityCapability = defineDomainCapability( 'deletion', 'mirror', 'bootstrap', + 'retired', ], inputSchema, outputSchema, diff --git a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts index 4b4b6a211b..c8a9fbe636 100644 --- a/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts +++ b/packages/worker/src/mcp/capabilities/email/email-usage-get.workers.test.ts @@ -29,18 +29,13 @@ async function seedDailyCounter(input: { count: number day: string }) { - await env.APP_DB.prepare( - `INSERT INTO entitlement_daily_counters (user_id, resource, day, count, updated_at) - VALUES (?, ?, ?, ?, ?)`, - ) - .bind( - input.userId, - input.resource, - input.day, - input.count, - new Date().toISOString(), - ) - .run() + const meter = userMeterRpc({ env, userId: input.userId }) + await meter.initialize({ + resource: input.resource, + day: input.day, + count: input.count, + updatedAt: new Date().toISOString(), + }) } async function seedStoredMessages(userId: string, count: number) { @@ -180,32 +175,18 @@ test( const meterEmail = `usage-meter-${crypto.randomUUID()}@example.com` const meterUserId = await createStableUserIdFromEmail(meterEmail) await seedUsageAccount({ email: meterEmail, plan: 'pro' }) + await seedStoredMessages(meterUserId, 2) await seedDailyCounter({ userId: meterUserId, resource: 'email_sends_per_day', - count: 3, + count: 13, day, }) await seedDailyCounter({ userId: meterUserId, resource: 'email_receives_per_day', - count: 5, - day, - }) - await seedStoredMessages(meterUserId, 2) - const updatedAt = new Date().toISOString() - const meter = userMeterRpc({ env, userId: meterUserId }) - await meter.initialize({ - resource: 'email_sends_per_day', - day, - count: 13, - updatedAt, - }) - await meter.initialize({ - resource: 'email_receives_per_day', - day, count: 15, - updatedAt, + day, }) const meterResult = await emailUsageGetCapability.handler( {}, diff --git a/packages/worker/src/mcp/fetch-gateway.node.test.ts b/packages/worker/src/mcp/fetch-gateway.node.test.ts index 189235e61e..237c06683c 100644 --- a/packages/worker/src/mcp/fetch-gateway.node.test.ts +++ b/packages/worker/src/mcp/fetch-gateway.node.test.ts @@ -24,16 +24,6 @@ const env = { bind(...params: Array) { return { async run() { - if ( - normalizedQuery.includes( - 'insert into entitlement_daily_counters', - ) && - params[0] !== 'user-123' - ) { - throw new Error( - 'Entitlement counter upsert must bind the acting userId.', - ) - } return { meta: { changes: 1 } } }, async first() { @@ -622,8 +612,8 @@ test('gateway fetch records outbound_fetch usage metering', async () => { }) expect(successResponse.status).toBe(200) expect(fetchStub).toHaveBeenCalledTimes(1) - // waitUntil: entitlement D1 mirror + usage metering - expect(waitUntil).toHaveBeenCalledTimes(2) + // waitUntil: usage metering only (daily D1 mirror retired) + expect(waitUntil).toHaveBeenCalledTimes(1) expect(recordUsageSpy).toHaveBeenCalledTimes(1) expect(recordUsageSpy).toHaveBeenCalledWith(env, { userId: 'user-123', @@ -747,8 +737,8 @@ test('gateway fetch records outbound_fetch usage metering', async () => { waitUntil, }), ).rejects.toThrow('network failed with waitUntil') - // waitUntil: entitlement D1 mirror (before fetch) + usage metering (error path) - expect(waitUntil).toHaveBeenCalledTimes(2) + // waitUntil: usage metering only (daily D1 mirror retired) + expect(waitUntil).toHaveBeenCalledTimes(1) expect(recordUsageSpy).toHaveBeenCalledTimes(1) expect(recordUsageSpy.mock.calls[0]?.[1]).toMatchObject({ entityId: 'api.example.com', From b955401ccc3e14b780910d8b330a01437dacd14f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 15:41:44 +0000 Subject: [PATCH 32/39] docs(entitlements): stage daily mirror retirement Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 34 +++-- .../contributing/architecture/entitlements.md | 123 ++++++++---------- 2 files changed, 73 insertions(+), 84 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index ab081153f0..6ca71ca496 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -388,8 +388,10 @@ layout is `index1 = userId`, `blob1 = event type`, `blob2 = delivery outcome`, `blob3 = source timestamp`, and `double1 = 1`. Admin queries return only platform-wide day/outcome counts and weight sampled rows by `_sample_interval`. When Analytics Engine SQL is unreachable, these two charts zero-fill while the -rest of the page renders. Local development uses the existing D1 counters and -delivery-event table because Wrangler's emulated dataset has no SQL API. +rest of the page renders. Local development cannot query Wrangler's emulated +Analytics Engine SQL API: email quota aggregates degrade to empty (with an +explicit warning) because code no longer reads `entitlement_daily_counters`, +while delivery-outcome aggregates still read D1 `email_delivery_events`. **Mailbox expand-phase parity events** reuse the same `EMAIL_EVENTS` dataset with a separate row shape defined in @@ -637,21 +639,16 @@ time-pruned. Deletion-fence legacy lease rows are bounded by the D1 snapshot replace on `markDeleting` rather than time retention; DO-authority rows clear on release/repair/purge. -**Expand-phase D1 mirrors (daily counters only):** enforcement and point reads -are authoritative in UserMeter for daily counters. D1 -`entitlement_daily_counters` is **not** dropped — it remains a best-effort -mirror for existing readers and reporting. After each DO consume/refund/inbound -claim, the entitlements service schedules a non-awaited absolute mirror write -keyed by `(user_id, resource, day)` with a revision-ordered `updated_at` token -(`r/` + zero-padded revision from `userMeterMirrorUpdatedAtToken`) so late -writes cannot overwrite newer state. See +**D1 daily mirror writes stopped:** enforcement and point reads for daily +counters are solely in UserMeter. Code no longer reads or writes +`entitlement_daily_counters`; the physical table stays quiescent until a +follow-up migration-only deploy drops it. See [Entitlements](./entitlements.md#usermeter-expand-phase). **Daily cold bootstrap:** a missing `(resource, day)` row returns -`needs_bootstrap`. The service performs one legacy D1 point read on -`entitlement_daily_counters`, then `initialize()` seeds the DO row with -`INSERT OR IGNORE` (concurrent callers cannot double-apply the baseline). Warm -daily paths never read D1 for enforcement. +`needs_bootstrap`. The service calls `initialize({ count: 0 })` with +`INSERT OR IGNORE` (concurrent callers stay safe). Warm daily paths never read +D1 for enforcement. Account deletion calls `UserMeter.purge()` (one RPC per user, no D1 id scan; `deleteAll` clears counters, claims, storage-byte shadow, package-service @@ -1664,10 +1661,11 @@ Current retention policies: After Mailbox cut-over, the same 365-day window and blob-before-row ordering are self-enforced by the Mailbox DO alarm; `system:email` stays on the D1 system-email retention job. -- `entitlement_daily_counters`: expand-phase **mirror** of UserMeter daily - counters (authoritative state lives in the per-user `UserMeter` DO). Rows keep - 400 days by `day` key until mirror retirement is verified after - reporting-off-D1 merges; the table is not dropped in this phase. +- `entitlement_daily_counters`: **quiescent pending drop** — mirror writes, + bootstrap reads, account-deletion inventory, and retention pruning all stopped + in the code deploy; existing rows age in place until a follow-up migration + drops the table. Daily counter retention lives in the per-user `UserMeter` DO + (`userMeterDailyCounterRetentionDays`). - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; raw Analytics Engine usage events follow platform retention. - `feature_flag_exposure_rollups`: local-dev/test flag exposure rollups keep 90 diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 30117e2fc6..df727c933f 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -147,25 +147,20 @@ reserve bytes in D1 or UserMeter. `consumeInboundDelivery` RPCs check the plan limit and increment inside the DO. The Durable Object request model serializes mutations per user; counter updates use optimistic concurrency on monotonic `revision` so concurrent consumes cannot -overshoot. Missing `(resource, day)` rows return `needs_bootstrap` rather than -silently starting at zero on a warm account. - -**Cold bootstrap:** on `needs_bootstrap`, the service performs one legacy D1 -point read on `entitlement_daily_counters`, then `UserMeter.initialize()` seeds -the row with `INSERT OR IGNORE` (concurrent callers cannot double-apply the -baseline). The warm enforcement path awaits only the DO RPC — never APP_DB. - -**Non-awaited D1 mirror:** after each successful consume, refund, or inbound -delivery claim, the service schedules a best-effort absolute mirror write to -`entitlement_daily_counters` via `waitUntil` when available (otherwise a caught -void promise). Mirror ordering uses the DO-minted `mirrorUpdatedAt` token -(`r/` + zero-padded revision) in the existing `updated_at` TEXT column so late -writes cannot overwrite newer state, including refunds that lower `count`. -Mirror failures are logged and never affect enforcement. The D1 table is **not** -dropped in this phase — it remains for existing readers and reporting. +overshoot. Missing `(resource, day)` rows return `needs_bootstrap`; the service +then initializes that key at zero via `UserMeter.initialize()` +(`INSERT OR IGNORE`, concurrent-safe) before retrying. Warm enforcement awaits +only the DO RPC and never touches D1 daily counter state. + +**D1 daily mirror writes stopped:** consume, refund, inbound charge/read, +point-read surfaces, account export/deletion inventory, and retention no longer +read or write `entitlement_daily_counters`. The physical table stays in +production until a follow-up migration-only deploy drops it (migrations apply +before Workers); existing rows are quiescent historical mirror state. Analytics +Engine remains the production reporting path for email send/receive aggregates. **Point-read surfaces** call `readDailyEntitlementResourceUsage` (UserMeter with -the same cold-bootstrap path): +the same cold zero-init path): - Account usage UI — `packages/worker/src/app/account-usage-data.ts` - Account email usage panel — `packages/worker/src/app/account-email-data.ts` @@ -173,23 +168,15 @@ the same cold-bootstrap path): - Admin per-user usage drill-down — `packages/worker/src/admin/user-usage-data.ts` -Non-daily resources and contexts without `USER_METER` still use -`readEntitlementResourceUsage` against D1. - -During a rolling deployment, requests already running on the previous Worker -version may still increment D1 after a new-version request bootstraps its DO -row. Cloudflare activation bounds that overlap to in-flight requests, but -operators should treat mirror parity during the deploy window as approximate; -post-deploy requests have one authority in UserMeter. +Non-daily resources still use `readEntitlementResourceUsage` against D1. +`readEntitlementResourceUsage` for daily resources throws and directs callers to +the UserMeter helpers above. **Inbound retry idempotency:** inbound receive quota uses `UserMeter.consumeInboundDelivery`, which atomically claims `delivery_id` and consumes one `email_receives_per_day` unit inside a SQLite transaction. Retries return the accepted counter without incrementing (`replayed: true`). -Cross-UTC-day retries use the original claim's resource/day. The legacy D1 -mirror is scheduled from the email path via -`scheduleAbsoluteDailyEntitlementMirror` so the email subsystem does not -duplicate mirror SQL. +Cross-UTC-day retries use the original claim's resource/day. ### D1 payload storage bytes — UserMeter shadow (expand phase slice 3) @@ -371,29 +358,37 @@ same-token rollout mirrors; email keeps its D1 lease path. Package-service and storage authority flips remain separate high-risk contract follow-ups after soak/parity review. -**Daily-counter mirror retirement:** dropping D1 `entitlement_daily_counters` -waits until reporting-off-D1 work merges and mirror parity is verified in -production. +**Daily-counter mirror retirement (two-deploy):** this code deploy stops all D1 +mirror/bootstrap/inventory/retention use while leaving the physical +`entitlement_daily_counters` table in place. A follow-up migration-only PR drops +the table after this deploy is healthy. Production `admin_user_meter_parity` +scans across 38 users showed zero daily mismatches with Analytics Engine +reporting active — the deploy rationale for stopping mirror writes before the +drop. While the table exists, parity still compares `d1Count === meterCount` +(`mirrorRetired: false`); after the drop migration, `mirrorRetired: true` +reports meter counts only. ### Admin UserMeter parity gates (`admin_user_meter_parity`) -Production verification for mirror retirement and authority flips uses the -admin-only read-only capability `admin_user_meter_parity` (input: +Production verification for mirror retirement and remaining authority flips uses +the admin-only read-only capability `admin_user_meter_parity` (input: `stable_user_id`). It compares production-shaped D1 rows for one account against -direct UserMeter RPCs and never bootstraps or writes parity state. Opening a -cold UserMeter stub may still run Durable Object constructor schema maintenance -and opportunistic stale daily-counter pruning. Cold meter rows surface as +direct UserMeter RPCs and never bootstraps or writes parity state. Daily +comparison retires automatically once the drop migration removes +`entitlement_daily_counters` (`daily.mirrorRetired: true`). Opening a cold +UserMeter stub may still run Durable Object constructor schema maintenance and +opportunistic stale daily-counter pruning. Cold meter rows surface as `needsBootstrap` with `meterCount`/`meterBytes` null. Interpret the structured report as independent gates: -| Gate | Pass condition | -| -------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Daily counters (current UTC day) | Each of the four daily resources has `parity: true` (`d1Count === meterCount`); aggregate `daily.mismatchCount === 0`. | -| Storage bytes | `storage.parity` — D1 `users.d1_storage_bytes` equals UserMeter `readStorageBytes` (not `needsBootstrap`). | -| Package services | `packageServices.parity` — inventory mismatch category counts are all zero (`d1Only` / `meterOnly` / `statusMismatch` / `startedAtMismatch` / `sourceUpdatedAtMismatch`), fresh-running counts match under the shared 24h stale window, and the meter page walk is not `truncated`. | -| Deletion tombstone | `deletion.deletingAtParity` — D1 `users.deleting_at` matches the meter tombstone. | -| Temporary D1 lease mirror | `deletion.mirrorLeaseParity` — `doOnly === 0`, `legacyWithoutD1 === 0`, inventory not truncated, and `d1ActiveLeaseCount >= doAuthorityLeaseCount` (same-token mirror coverage). `tokenSetMismatches.d1Only` is reported but does **not** fail this gate. | +| Gate | Pass condition | +| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Daily counters (current UTC day) | While `daily.mirrorRetired` is false (table present): each of the four daily resources has `parity: true` (`d1Count === meterCount`); aggregate `daily.mismatchCount === 0`. After the drop migration (`mirrorRetired: true`): report meter counts only; `d1Count`/`delta` are null, each resource has `parity: true`, and `mismatchCount === 0` (no D1 comparison). | +| Storage bytes | `storage.parity` — D1 `users.d1_storage_bytes` equals UserMeter `readStorageBytes` (not `needsBootstrap`). | +| Package services | `packageServices.parity` — inventory mismatch category counts are all zero (`d1Only` / `meterOnly` / `statusMismatch` / `startedAtMismatch` / `sourceUpdatedAtMismatch`), fresh-running counts match under the shared 24h stale window, and the meter page walk is not `truncated`. | +| Deletion tombstone | `deletion.deletingAtParity` — D1 `users.deleting_at` matches the meter tombstone. | +| Temporary D1 lease mirror | `deletion.mirrorLeaseParity` — `doOnly === 0`, `legacyWithoutD1 === 0`, inventory not truncated, and `d1ActiveLeaseCount >= doAuthorityLeaseCount` (same-token mirror coverage). `tokenSetMismatches.d1Only` is reported but does **not** fail this gate. | **D1-only leases:** email and other transition paths that omit `env` still take exact D1 leases, so `d1Only > 0` is expected until that handoff. Mirror-removal @@ -402,12 +397,12 @@ are known email/transition holders via `admin_account_write_lease_list` and holder classification before retiring the temporary D1 mirror inventory. **Threshold:** treat unexplained mismatches as blocking for the corresponding -cutover (daily mirror retirement, storage authority flip, package-service -authority flip, or temporary D1 lease-mirror removal). Expected cold accounts -may report `needsBootstrap` until live traffic or an intentional bootstrap path -seeds the DO; that is a bootstrap gap, not a silent pass. Truncated inventories -fail closed (`parity` / `mirrorLeaseParity` false) so operators re-run or raise -the bounded page cap rather than approve a partial compare. +cutover (daily mirror retirement before the drop migration, storage authority +flip, package-service authority flip, or temporary D1 lease-mirror removal). +Expected cold accounts may report `needsBootstrap` until live traffic seeds the +DO; that is a bootstrap gap, not a silent pass. Truncated inventories fail +closed (`parity` / `mirrorLeaseParity` false) so operators re-run or raise the +bounded page cap rather than approve a partial compare. Module wiring: `consumeDailyEntitlement`, `refundDailyEntitlement`, and `readDailyEntitlementResourceUsage` require `env.USER_METER` and fail closed @@ -568,14 +563,13 @@ Rules: the limit abuse-resistant for permanent rejects (parse failures, entitlement/quota rejects). - **Cold bootstrap:** missing `(resource, day)` rows trigger one legacy D1 point - read and a single `UserMeter.initialize()` before retrying the consume. + **Cold bootstrap:** missing `(resource, day)` rows trigger + `UserMeter.initialize({ count: 0 })` (`INSERT OR IGNORE`) before retrying the + consume. Concurrent cold callers cannot double-apply a non-zero baseline. - **D1 mirror (expand phase):** after each successful consume/refund/inbound - claim, a best-effort, non-awaited absolute mirror write updates - `entitlement_daily_counters` with revision-ordered `updated_at` tokens. The - table is not dropped — it remains for existing readers and reporting until - reporting-off-D1 retirement is verified. + **D1 mirror writes stopped:** consume/refund/inbound charge/read paths never + touch `entitlement_daily_counters`; the physical table remains quiescent until + a follow-up migration drops it. A delivery claim remains charged when later storage fails. Cloudflare Email Routing retries replay that same `delivery_id` through @@ -583,9 +577,6 @@ Rules: across a UTC-day boundary. The retained claim is the idempotency boundary; production inbound handling does not call `refundDailyEntitlement`. - `incrementDailyEntitlementCounter` remains for raw D1 counter writes (tests, - backfills, and legacy paths). - - **Boolean allowances** (persistent package services) are modeled as limit `0` (not allowed) vs `1` (allowed) so the numeric contract stays uniform. - **Per-unit size limits** (`email_message_bytes`) compare one candidate value @@ -795,9 +786,9 @@ in [`../environment-variables.md`](../environment-variables.md). `0066-stripe-billing.sql`; owned by `packages/worker/src/billing/`, read by `getUserPlan` via `resolveEffectivePlan`. `stripe_plan` stays nullable because it is Stripe-derived; `max` is manual-only. -- `entitlement_daily_counters` — expand-phase **mirror** of UserMeter daily - counters (authoritative state in the per-user `UserMeter` DO), created by - migration `0048-user-plans-and-entitlement-counters.sql`; included in the - account-deletion cascade (`packages/worker/src/app/account-deletion.ts`). - Table retirement waits until reporting-off-D1 merges and mirror parity is - verified. +- `entitlement_daily_counters` — **quiescent pending drop**. Created by + migration `0048-user-plans-and-entitlement-counters.sql`; mirror writes, + bootstrap reads, account-deletion inventory, and retention pruning all stopped + in the code deploy. Daily counters are authoritative in the per-user + `UserMeter` DO; account export/deletion use `UserMeter` RPCs. A follow-up + migration-only PR drops the physical table after this deploy is healthy. From 164b275a60529f821fc2f9892b55216762e724da Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 15:44:08 +0000 Subject: [PATCH 33/39] fix(entitlements): preserve pre-drop schema guards Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 7 ++-- .../contributing/architecture/entitlements.md | 39 ++++++++++--------- packages/worker/src/account/data-targets.ts | 1 + .../src/app/account-retention-dispositions.ts | 6 +++ 4 files changed, 32 insertions(+), 21 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 6ca71ca496..1788af1bab 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -1662,9 +1662,10 @@ Current retention policies: are self-enforced by the Mailbox DO alarm; `system:email` stays on the D1 system-email retention job. - `entitlement_daily_counters`: **quiescent pending drop** — mirror writes, - bootstrap reads, account-deletion inventory, and retention pruning all stopped - in the code deploy; existing rows age in place until a follow-up migration - drops the table. Daily counter retention lives in the per-user `UserMeter` DO + bootstrap reads, and scheduled retention pruning stop in the code deploy. + Account deletion keeps removing rows while the physical table exists; a + follow-up code deploy removes that inventory target before the later drop + migration. Daily counter retention lives in the per-user `UserMeter` DO (`userMeterDailyCounterRetentionDays`). - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; raw Analytics Engine usage events follow platform retention. diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index df727c933f..3fd3b0c82d 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -153,11 +153,12 @@ then initializes that key at zero via `UserMeter.initialize()` only the DO RPC and never touches D1 daily counter state. **D1 daily mirror writes stopped:** consume, refund, inbound charge/read, -point-read surfaces, account export/deletion inventory, and retention no longer -read or write `entitlement_daily_counters`. The physical table stays in -production until a follow-up migration-only deploy drops it (migrations apply -before Workers); existing rows are quiescent historical mirror state. Analytics -Engine remains the production reporting path for email send/receive aggregates. +point-read surfaces, account export, and retention no longer read or write +`entitlement_daily_counters`. Account deletion keeps removing rows while the +physical table exists. A follow-up code deploy removes that target before the +later migration-only drop (migrations apply before Workers); existing rows are +otherwise quiescent historical mirror state. Analytics Engine remains the +production reporting path for email send/receive aggregates. **Point-read surfaces** call `readDailyEntitlementResourceUsage` (UserMeter with the same cold zero-init path): @@ -358,15 +359,16 @@ same-token rollout mirrors; email keeps its D1 lease path. Package-service and storage authority flips remain separate high-risk contract follow-ups after soak/parity review. -**Daily-counter mirror retirement (two-deploy):** this code deploy stops all D1 -mirror/bootstrap/inventory/retention use while leaving the physical -`entitlement_daily_counters` table in place. A follow-up migration-only PR drops -the table after this deploy is healthy. Production `admin_user_meter_parity` -scans across 38 users showed zero daily mismatches with Analytics Engine -reporting active — the deploy rationale for stopping mirror writes before the -drop. While the table exists, parity still compares `d1Count === meterCount` -(`mirrorRetired: false`); after the drop migration, `mirrorRetired: true` -reports meter counts only. +**Daily-counter mirror retirement (three-deploy):** this code deploy stops all +D1 mirror/bootstrap/retention use while leaving the physical +`entitlement_daily_counters` table and account-deletion target in place. A +follow-up code deploy removes that target after this Worker is healthy; the +third, migration-only deploy drops the table. Production +`admin_user_meter_parity` scans across 38 users showed zero daily mismatches +with Analytics Engine reporting active — the deploy rationale for stopping +mirror writes before the drop. While the table exists, parity still compares +`d1Count === meterCount` (`mirrorRetired: false`); after the drop migration, +`mirrorRetired: true` reports meter counts only. ### Admin UserMeter parity gates (`admin_user_meter_parity`) @@ -788,7 +790,8 @@ in [`../environment-variables.md`](../environment-variables.md). it is Stripe-derived; `max` is manual-only. - `entitlement_daily_counters` — **quiescent pending drop**. Created by migration `0048-user-plans-and-entitlement-counters.sql`; mirror writes, - bootstrap reads, account-deletion inventory, and retention pruning all stopped - in the code deploy. Daily counters are authoritative in the per-user - `UserMeter` DO; account export/deletion use `UserMeter` RPCs. A follow-up - migration-only PR drops the physical table after this deploy is healthy. + bootstrap reads, and retention pruning stop in the code deploy. Account + deletion keeps removing user rows while the table exists. Daily counters are + authoritative in the per-user `UserMeter` DO; account export uses UserMeter + RPCs. A follow-up code deploy removes the deletion target before a later + migration-only PR drops the table. diff --git a/packages/worker/src/account/data-targets.ts b/packages/worker/src/account/data-targets.ts index 81fe11353f..eac40d291a 100644 --- a/packages/worker/src/account/data-targets.ts +++ b/packages/worker/src/account/data-targets.ts @@ -219,6 +219,7 @@ export const accountUserDataTargets: ReadonlyArray = [ { kind: 'user_id', table: 'email_sender_identities' }, { kind: 'user_id', table: 'email_sender_rules' }, { kind: 'user_id', table: 'webhook_endpoints' }, + { kind: 'user_id', table: 'entitlement_daily_counters' }, { kind: 'user_columns', table: 'platform_feedback', diff --git a/packages/worker/src/app/account-retention-dispositions.ts b/packages/worker/src/app/account-retention-dispositions.ts index 5178d8ace2..d4f23694d0 100644 --- a/packages/worker/src/app/account-retention-dispositions.ts +++ b/packages/worker/src/app/account-retention-dispositions.ts @@ -32,6 +32,12 @@ export const accountRetentionDispositions: ReadonlyArray Date: Sat, 1 Aug 2026 16:02:03 +0000 Subject: [PATCH 34/39] fix(entitlements): clarify mirror-stop fallbacks Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 20 ++++++----- .../src/app/admin-insights-data.node.test.ts | 34 +++++++++++++++++-- .../worker/src/app/admin-insights-data.ts | 9 ++++- 3 files changed, 51 insertions(+), 12 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 1788af1bab..9792bb8e48 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -390,8 +390,8 @@ platform-wide day/outcome counts and weight sampled rows by `_sample_interval`. When Analytics Engine SQL is unreachable, these two charts zero-fill while the rest of the page renders. Local development cannot query Wrangler's emulated Analytics Engine SQL API: email quota aggregates degrade to empty (with an -explicit warning) because code no longer reads `entitlement_daily_counters`, -while delivery-outcome aggregates still read D1 `email_delivery_events`. +explicit warning) rather than reading the retired D1 mirror, while +delivery-outcome aggregates still read D1 `email_delivery_events`. **Mailbox expand-phase parity events** reuse the same `EMAIL_EVENTS` dataset with a separate row shape defined in @@ -639,9 +639,10 @@ time-pruned. Deletion-fence legacy lease rows are bounded by the D1 snapshot replace on `markDeleting` rather than time retention; DO-authority rows clear on release/repair/purge. -**D1 daily mirror writes stopped:** enforcement and point reads for daily -counters are solely in UserMeter. Code no longer reads or writes -`entitlement_daily_counters`; the physical table stays quiescent until a +**D1 daily mirror writes stopped:** enforcement, point reads, bootstrap, and +mirror paths no longer read or write `entitlement_daily_counters`. The admin +parity report (`admin_user_meter_parity`) temporarily reads the table while it +exists for migration verification. The physical table stays quiescent until a follow-up migration-only deploy drops it. See [Entitlements](./entitlements.md#usermeter-expand-phase). @@ -1662,10 +1663,11 @@ Current retention policies: are self-enforced by the Mailbox DO alarm; `system:email` stays on the D1 system-email retention job. - `entitlement_daily_counters`: **quiescent pending drop** — mirror writes, - bootstrap reads, and scheduled retention pruning stop in the code deploy. - Account deletion keeps removing rows while the physical table exists; a - follow-up code deploy removes that inventory target before the later drop - migration. Daily counter retention lives in the per-user `UserMeter` DO + bootstrap reads, and scheduled retention pruning stop in the code deploy. The + admin parity report (`admin_user_meter_parity`) temporarily reads the table + while it exists; account deletion keeps removing rows while the physical table + exists; a follow-up code deploy removes that inventory target before the later + drop migration. Daily counter retention lives in the per-user `UserMeter` DO (`userMeterDailyCounterRetentionDays`). - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; raw Analytics Engine usage events follow platform retention. diff --git a/packages/worker/src/app/admin-insights-data.node.test.ts b/packages/worker/src/app/admin-insights-data.node.test.ts index 390e2b97a2..256c879265 100644 --- a/packages/worker/src/app/admin-insights-data.node.test.ts +++ b/packages/worker/src/app/admin-insights-data.node.test.ts @@ -314,7 +314,11 @@ test('loadAdminInsightsData assembles the dashboard payload', async () => { consoleWarn.mockImplementation(() => {}) consoleWarn.mockClear() const data = await loadAdminInsightsData( - { APP_DB: createInsightsTestDb() } as Env, + { + APP_DB: createInsightsTestDb(), + EMAIL_EVENTS: {} as AnalyticsEngineDataset, + WRANGLER_IS_LOCAL_DEV: 'true', + } as Env, now, ) @@ -323,6 +327,10 @@ test('loadAdminInsightsData assembles the dashboard payload', async () => { 'admin-insights-email-quota-aggregate-unavailable', { reason: 'entitlement-daily-counters-retired-local-dev' }, ) + expect(consoleWarn).not.toHaveBeenCalledWith( + 'admin-insights-email-quota-aggregate-unavailable', + { reason: 'missing-email-events-binding' }, + ) expect(data.totals).toEqual({ users: 8, verifiedUsers: 5, @@ -390,10 +398,32 @@ test('loadAdminInsightsData assembles the dashboard payload', async () => { expect(data.activation.medianHoursToActivation).toBe(36) }) +test('loadAdminInsightsData warns when EMAIL_EVENTS binding is missing', async () => { + consoleWarn.mockImplementation(() => {}) + consoleWarn.mockClear() + const data = await loadAdminInsightsData( + { APP_DB: createInsightsTestDb() } as Env, + now, + ) + + expect(data.ok).toBe(true) + expect(consoleWarn).toHaveBeenCalledWith( + 'admin-insights-email-quota-aggregate-unavailable', + { reason: 'missing-email-events-binding' }, + ) + expect(consoleWarn).not.toHaveBeenCalledWith( + 'admin-insights-email-quota-aggregate-unavailable', + { reason: 'entitlement-daily-counters-retired-local-dev' }, + ) +}) + test('activation latency excludes users who have no usable verification date', async () => { consoleWarn.mockImplementation(() => {}) const data = await loadAdminInsightsData( - { APP_DB: createInsightsTestDb({ activationLatencyRows: [] }) } as Env, + { + APP_DB: createInsightsTestDb({ activationLatencyRows: [] }), + WRANGLER_IS_LOCAL_DEV: 'true', + } as Env, now, ) // Two users reached the milestone, but neither has timing we can measure. diff --git a/packages/worker/src/app/admin-insights-data.ts b/packages/worker/src/app/admin-insights-data.ts index 5babb4f167..c979787e64 100644 --- a/packages/worker/src/app/admin-insights-data.ts +++ b/packages/worker/src/app/admin-insights-data.ts @@ -285,10 +285,17 @@ async function loadEmailInsightsRows(input: { // retired, so local/dev email quota aggregates degrade explicitly to empty // (with a warning) rather than querying a removed table. Delivery-event // aggregates still come from D1 email_delivery_events. - if (!input.env.EMAIL_EVENTS || input.env.WRANGLER_IS_LOCAL_DEV === 'true') { + if (input.env.WRANGLER_IS_LOCAL_DEV === 'true') { console.warn('admin-insights-email-quota-aggregate-unavailable', { reason: 'entitlement-daily-counters-retired-local-dev', }) + } + if (!input.env.EMAIL_EVENTS) { + console.warn('admin-insights-email-quota-aggregate-unavailable', { + reason: 'missing-email-events-binding', + }) + } + if (input.env.WRANGLER_IS_LOCAL_DEV === 'true' || !input.env.EMAIL_EVENTS) { const deliveryRows = await input.env.APP_DB.prepare( `SELECT substr(created_at, 1, 10) AS day, event_type, COUNT(*) AS n FROM email_delivery_events From 436c27231f0ce634f9a5369fcef552fea8a7ffab Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 16:09:07 +0000 Subject: [PATCH 35/39] docs(entitlements): clarify pre-drop inventory reads Co-authored-by: Kent C. Dodds --- docs/contributing/architecture/entitlements.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 3fd3b0c82d..80dda96eeb 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -153,12 +153,13 @@ then initializes that key at zero via `UserMeter.initialize()` only the DO RPC and never touches D1 daily counter state. **D1 daily mirror writes stopped:** consume, refund, inbound charge/read, -point-read surfaces, account export, and retention no longer read or write -`entitlement_daily_counters`. Account deletion keeps removing rows while the -physical table exists. A follow-up code deploy removes that target before the -later migration-only drop (migrations apply before Workers); existing rows are -otherwise quiescent historical mirror state. Analytics Engine remains the -production reporting path for email send/receive aggregates. +point-read surfaces, and retention no longer read or write +`entitlement_daily_counters`. Generic account export and deletion keep reading +or removing user rows while the physical table exists. A follow-up code deploy +removes those inventory targets before the later migration-only drop (migrations +apply before Workers); existing rows are otherwise quiescent historical mirror +state. Analytics Engine remains the production reporting path for email +send/receive aggregates. **Point-read surfaces** call `readDailyEntitlementResourceUsage` (UserMeter with the same cold zero-init path): From fc3bf7a58d99c2fad6d78122c0f15c1118b2c32e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 16:37:02 +0000 Subject: [PATCH 36/39] refactor(account): detach pending daily counter table Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 55 +++++---- .../contributing/architecture/entitlements.md | 45 ++++---- .../src/account/data-targets.node.test.ts | 104 ++++++++++++++++++ packages/worker/src/account/data-targets.ts | 43 +++++++- .../worker/src/account/export.node.test.ts | 5 + ...ccount-retention-dispositions.node.test.ts | 9 ++ .../src/app/account-retention-dispositions.ts | 2 +- 7 files changed, 214 insertions(+), 49 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index c569a4fe78..54ebff1ad7 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -71,11 +71,13 @@ Deletion must cover these user-owned surfaces: - **D1:** every live table with `user_id` / `*_user_id` ownership columns, plus transitive children (`secret_entries`, `value_entries`, `email_attachments`) - and listing children for community-owned listings. The guardrail test in - `packages/worker/src/app/account-deletion.node.test.ts` applies the live - migrations to SQLite and fails if a user-owned schema column is not - represented in the deletion target list, or if the deletion target list - references a stale column. + and listing children for community-owned listings. Physical tables pending a + later drop migration are registered in `accountUserDataPendingDropTargets` + (schema coverage only; runtime deletion never queries them). The guardrail + test in `packages/worker/src/app/account-deletion.node.test.ts` applies the + live migrations to SQLite and fails if a user-owned schema column is not + represented in the deletion target list or pending-drop registry, or if those + lists reference a stale column. - **Durable Objects:** `JobManager`, `StorageRunner`, `RepoSession`, `RemoteConnectorSession`, `PackageRealtimeSession`, `PackageServiceInstance`, `McpClientHub`, `RunLog`, `UserMeter`, `StripePlanRefresh`, and `Mailbox` are @@ -114,22 +116,25 @@ Deletion must cover these user-owned surfaces: Account export is implemented in `packages/worker/src/account/export.ts`. It mirrors the deletion inventory so portability and account migration cover the same user-owned storage surfaces. The D1 table list and shared kind→SQL match -builders live in `account-data-targets.ts` (`accountUserDataTargets`, -`buildUserScopedTargetMatch`); export redaction columns also live there. -Out-of-band surfaces (Durable Objects, KV schemes, R2, Vectorize, Artifacts) are -declared in `account-user-owned-surfaces.ts` and consumed by both deletion and -export. Growth-table retention dispositions are linked in +builders live in `account/data-targets.ts` (`accountUserDataTargets`, +`buildUserScopedTargetMatch`); export redaction columns and +`accountUserDataPendingDropTargets` also live there. Out-of-band surfaces +(Durable Objects, KV schemes, R2, Vectorize, Artifacts) are declared in +`account-user-owned-surfaces.ts` and consumed by both deletion and export. +Growth-table retention dispositions are linked in `account-retention-dispositions.ts`. `packages/worker/src/account/export.node.test.ts` applies the live migrations to SQLite and fails if a `user_id` / `*_user_id` column is not covered by the -export list. The hard invariant is the same as every storage path: callers pass -the authenticated user's stable MCP `userId`, and every query or Durable Object -lookup is scoped to that id. +export list or pending-drop registry. The hard invariant is the same as every +storage path: callers pass the authenticated user's stable MCP `userId`, and +every query or Durable Object lookup is scoped to that id. System email rows owned by `system:email` are intentionally absent from account exports for the same reason they are absent from deletion: they belong to the -operator inbox surface, not to the exporting user. The export manifest lists -this under `excludedD1Surfaces` so the omission is explicit. +operator inbox surface, not to the exporting user. Pending-drop physical tables +(such as quiescent `entitlement_daily_counters`) are also absent from export +queries. The export manifest lists both under `excludedD1Surfaces` so the +omission is explicit. Platform-feedback submissions are included in the submitting user's own D1 export section. An export never includes submissions owned by other users, @@ -639,11 +644,12 @@ time-pruned. Deletion-fence legacy lease rows are bounded by the D1 snapshot replace on `markDeleting` rather than time retention; DO-authority rows clear on release/repair/purge. -**D1 daily mirror writes stopped:** enforcement, point reads, bootstrap, and -mirror paths no longer read or write `entitlement_daily_counters`. The admin -parity report (`admin_user_meter_parity`) temporarily reads the table while it -exists for migration verification. The physical table stays quiescent until a -follow-up migration-only deploy drops it. See +**D1 daily mirror writes stopped:** enforcement, point reads, bootstrap, mirror, +and account export/deletion inventory paths no longer read or write +`entitlement_daily_counters`. The admin parity report +(`admin_user_meter_parity`) temporarily reads the table while it exists for +migration verification. The physical table stays registered as a pending-drop +coverage exemption until a later migration-only deploy drops it. See [Entitlements](./entitlements.md#usermeter-expand-phase). **Daily cold bootstrap:** a missing `(resource, day)` row returns @@ -1667,10 +1673,11 @@ Current retention policies: are self-enforced by the Mailbox DO alarm; `system:email` stays on the D1 system-email retention job. - `entitlement_daily_counters`: **quiescent pending drop** — mirror writes, - bootstrap reads, and scheduled retention pruning stop in the code deploy. The - admin parity report (`admin_user_meter_parity`) temporarily reads the table - while it exists; account deletion keeps removing rows while the physical table - exists; a follow-up code deploy removes that inventory target before the later + bootstrap reads, scheduled retention pruning, and account export/deletion + inventory stopped across the stage 1/2 code deploys. The admin parity report + (`admin_user_meter_parity`) temporarily reads the table while it exists. + Runtime account inventory no longer queries the table; it remains registered + in `accountUserDataPendingDropTargets` for schema coverage until the later drop migration. Daily counter retention lives in the per-user `UserMeter` DO (`userMeterDailyCounterRetentionDays`). - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 80dda96eeb..1d03451b13 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -153,12 +153,14 @@ then initializes that key at zero via `UserMeter.initialize()` only the DO RPC and never touches D1 daily counter state. **D1 daily mirror writes stopped:** consume, refund, inbound charge/read, -point-read surfaces, and retention no longer read or write -`entitlement_daily_counters`. Generic account export and deletion keep reading -or removing user rows while the physical table exists. A follow-up code deploy -removes those inventory targets before the later migration-only drop (migrations -apply before Workers); existing rows are otherwise quiescent historical mirror -state. Analytics Engine remains the production reporting path for email +point-read surfaces, retention, and generic account export/deletion no longer +read or write `entitlement_daily_counters`. The physical table remains in the +migrated schema as a pending-drop target (`accountUserDataPendingDropTargets`) +so schema coverage still recognizes `entitlement_daily_counters.user_id` without +runtime inventory querying it; export lists the omission under +`excludedD1Surfaces` (no raw rows). A later migration-only deploy drops the +table (migrations apply before Workers); existing rows are quiescent historical +mirror state. Analytics Engine remains the production reporting path for email send/receive aggregates. **Point-read surfaces** call `readDailyEntitlementResourceUsage` (UserMeter with @@ -360,16 +362,17 @@ same-token rollout mirrors; email keeps its D1 lease path. Package-service and storage authority flips remain separate high-risk contract follow-ups after soak/parity review. -**Daily-counter mirror retirement (three-deploy):** this code deploy stops all -D1 mirror/bootstrap/retention use while leaving the physical -`entitlement_daily_counters` table and account-deletion target in place. A -follow-up code deploy removes that target after this Worker is healthy; the -third, migration-only deploy drops the table. Production -`admin_user_meter_parity` scans across 38 users showed zero daily mismatches -with Analytics Engine reporting active — the deploy rationale for stopping -mirror writes before the drop. While the table exists, parity still compares -`d1Count === meterCount` (`mirrorRetired: false`); after the drop migration, -`mirrorRetired: true` reports meter counts only. +**Daily-counter mirror retirement (three-deploy):** stage 1 (Worker #1133) +stopped all D1 mirror/bootstrap/retention use while leaving the physical +`entitlement_daily_counters` table and account inventory target in place. Stage +2 removes the runtime account export/deletion inventory target and registers the +table in `accountUserDataPendingDropTargets` for schema coverage (export +documents the omission; no raw rows). The third, migration-only deploy drops the +table. Production `admin_user_meter_parity` scans across 38 users showed zero +daily mismatches with Analytics Engine reporting active — the deploy rationale +for stopping mirror writes before the drop. While the table exists, parity still +compares `d1Count === meterCount` (`mirrorRetired: false`); after the drop +migration, `mirrorRetired: true` reports meter counts only. ### Admin UserMeter parity gates (`admin_user_meter_parity`) @@ -791,8 +794,8 @@ in [`../environment-variables.md`](../environment-variables.md). it is Stripe-derived; `max` is manual-only. - `entitlement_daily_counters` — **quiescent pending drop**. Created by migration `0048-user-plans-and-entitlement-counters.sql`; mirror writes, - bootstrap reads, and retention pruning stop in the code deploy. Account - deletion keeps removing user rows while the table exists. Daily counters are - authoritative in the per-user `UserMeter` DO; account export uses UserMeter - RPCs. A follow-up code deploy removes the deletion target before a later - migration-only PR drops the table. + bootstrap reads, retention pruning, and account export/deletion inventory + stopped across the stage 1/2 code deploys. The table remains in + `accountUserDataPendingDropTargets` for schema coverage until a later + migration-only PR drops it. Daily counters are authoritative in the per-user + `UserMeter` DO; account export uses UserMeter RPCs (no raw D1 mirror rows). diff --git a/packages/worker/src/account/data-targets.node.test.ts b/packages/worker/src/account/data-targets.node.test.ts index 15e642cba3..336f7cf9a9 100644 --- a/packages/worker/src/account/data-targets.node.test.ts +++ b/packages/worker/src/account/data-targets.node.test.ts @@ -1,14 +1,30 @@ +import { quoteSqlIdentifier } from '@kody-internal/shared/sql-literals.ts' +import { readdirSync, readFileSync } from 'node:fs' +import { DatabaseSync } from 'node:sqlite' import { expect, test } from 'vitest' import { accountExportForeignUserIdColumnsByTable, accountExportRedactedColumnsByTable, accountExportRedactedForeignUserId, + accountUserDataPendingDropTargets, accountUserDataTargets, buildUserScopedDeleteOrUpdateSql, buildUserScopedTargetMatch, + getAccountD1UserColumnCoverage, + getAccountExportExcludedD1Surfaces, + isExcludedFromAccountExport, type UserScopedDataTarget, } from './data-targets.ts' +function applyMigrations(db: DatabaseSync) { + const migrationsDir = new URL('../../migrations/', import.meta.url) + for (const fileName of readdirSync(migrationsDir) + .filter((file) => file.endsWith('.sql')) + .sort()) { + db.exec(readFileSync(new URL(fileName, migrationsDir), 'utf8')) + } +} + function matchFor(target: UserScopedDataTarget) { return buildUserScopedTargetMatch({ target, @@ -197,3 +213,91 @@ test('every accountUserDataTargets kind has a shared match builder and export gu ), ).toBe(true) }) + +test('pending-drop entitlement_daily_counters stays out of runtime deletion/export SQL while schema coverage holds', () => { + const pendingDrop = accountUserDataPendingDropTargets.find( + (target) => target.table === 'entitlement_daily_counters', + ) + expect(pendingDrop).toEqual( + expect.objectContaining({ + table: 'entitlement_daily_counters', + column: 'user_id', + surface: 'entitlement_daily_counters', + reason: expect.stringContaining('pending a later drop migration'), + }), + ) + expect(pendingDrop?.reason).toContain('No raw rows exported') + + expect( + accountUserDataTargets.some( + (target) => + 'table' in target && target.table === 'entitlement_daily_counters', + ), + ).toBe(false) + + const deletionStatements = accountUserDataTargets.map((target) => { + const match = matchFor(target) + return buildUserScopedDeleteOrUpdateSql(match).sql + }) + expect(deletionStatements.join('\n')).not.toMatch( + /entitlement_daily_counters/u, + ) + + const exportStatements = accountUserDataTargets + .filter((target) => !isExcludedFromAccountExport(target)) + .map((target) => { + const match = matchFor(target) + return `SELECT * FROM ${match.table} WHERE ${match.qualifiedWhereSql}` + }) + expect(exportStatements.join('\n')).not.toMatch(/entitlement_daily_counters/u) + + expect(getAccountExportExcludedD1Surfaces()).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + name: 'entitlement_daily_counters', + reason: expect.stringContaining('No raw rows exported'), + }), + ]), + ) + + const db = new DatabaseSync(':memory:') + applyMigrations(db) + const tableExists = db + .prepare( + `SELECT 1 AS present + FROM sqlite_schema + WHERE type = 'table' AND name = 'entitlement_daily_counters'`, + ) + .get() as { present: number } | undefined + expect(tableExists?.present).toBe(1) + + const liveUserColumns = new Set() + const tables = db + .prepare( + `SELECT name + FROM sqlite_schema + WHERE type = 'table' AND name NOT LIKE 'sqlite_%' + ORDER BY name`, + ) + .all() as Array<{ name: string }> + for (const table of tables) { + const columns = db + .prepare(`PRAGMA table_info(${quoteSqlIdentifier(table.name)})`) + .all() as Array<{ name: string }> + for (const column of columns) { + if (column.name === 'user_id' || column.name.endsWith('_user_id')) { + liveUserColumns.add(`${table.name}.${column.name}`) + } + } + } + const coveredColumns = getAccountD1UserColumnCoverage() + expect(coveredColumns.has('entitlement_daily_counters.user_id')).toBe(true) + const missing = [...liveUserColumns].filter( + (column) => !coveredColumns.has(column), + ) + const stale = [...coveredColumns].filter( + (column) => !liveUserColumns.has(column), + ) + expect(missing).toEqual([]) + expect(stale).toEqual([]) +}) diff --git a/packages/worker/src/account/data-targets.ts b/packages/worker/src/account/data-targets.ts index eac40d291a..133eecd72f 100644 --- a/packages/worker/src/account/data-targets.ts +++ b/packages/worker/src/account/data-targets.ts @@ -72,6 +72,33 @@ export const accountUserDataExcludedOwnerIds = [ }, ] as const +/** + * Physical D1 tables that still exist in the migrated schema (pending a later + * drop migration) but are intentionally absent from runtime account + * deletion/export inventory. Schema coverage guardrails treat these + * `user_id` / `*_user_id` columns as covered so the live table does not look + * like a missing inventory target; runtime deletion and export never query + * them. Export documents each omission under `excludedD1Surfaces` (no raw + * rows). + */ +export type AccountUserDataPendingDropTarget = { + table: string + column: string + surface: string + reason: string +} + +export const accountUserDataPendingDropTargets: ReadonlyArray = + [ + { + table: 'entitlement_daily_counters', + column: 'user_id', + surface: 'entitlement_daily_counters', + reason: + 'Quiescent daily-counter D1 mirror pending a later drop migration after mirror-stop Worker #1133. Authoritative daily counters live in UserMeter; account export and deletion must not query this table before the drop. No raw rows exported.', + }, + ] + /** Targets that account export should skip (deletion still covers them). */ export function isExcludedFromAccountExport( target: UserScopedDataTarget, @@ -81,7 +108,8 @@ export function isExcludedFromAccountExport( /** * Documented D1 surfaces omitted from account export. Includes owner-id - * exclusions plus targets marked includeInExport: false with surface/reason. + * exclusions, pending-drop physical tables, plus targets marked + * includeInExport: false with surface/reason. */ export function getAccountExportExcludedD1Surfaces(): Array<{ name: string @@ -92,6 +120,10 @@ export function getAccountExportExcludedD1Surfaces(): Array<{ name: exclusion.surface, reason: exclusion.reason, })), + ...accountUserDataPendingDropTargets.map((target) => ({ + name: target.surface, + reason: target.reason, + })), ] for (const target of accountUserDataTargets) { if (!isExcludedFromAccountExport(target)) continue @@ -116,7 +148,10 @@ export function getAccountExportExcludedD1Surfaces(): Array<{ * update. * Rows owned by accountUserDataExcludedOwnerIds are operator/platform data, * not user data; tests assert those owner ids stay deliberately excluded from - * user account operations. + * user account operations. Physical tables pending a later drop migration are + * listed in accountUserDataPendingDropTargets instead of here so runtime + * deletion/export never query them while schema coverage still recognizes + * their user columns. * * Order matters for deletion: child tables come before parent tables so the * cascade is self-contained even on engines / configs where foreign-key @@ -219,7 +254,6 @@ export const accountUserDataTargets: ReadonlyArray = [ { kind: 'user_id', table: 'email_sender_identities' }, { kind: 'user_id', table: 'email_sender_rules' }, { kind: 'user_id', table: 'webhook_endpoints' }, - { kind: 'user_id', table: 'entitlement_daily_counters' }, { kind: 'user_columns', table: 'platform_feedback', @@ -424,6 +458,9 @@ export function getAccountD1UserColumnCoverage() { } } } + for (const target of accountUserDataPendingDropTargets) { + covered.add(`${target.table}.${target.column}`) + } return covered } diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 4c757e37dc..6c5e19e4d1 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -156,8 +156,13 @@ test('account export documents and excludes operator-owned system email rows', a name: 'system_email_inboxes', reason: expect.stringContaining('Operator-owned inbound mail'), }), + expect.objectContaining({ + name: 'entitlement_daily_counters', + reason: expect.stringContaining('pending a later drop migration'), + }), ]), ) + expect(accountExport.d1).not.toHaveProperty('entitlement_daily_counters') }) test('account export includes submitted feedback but excludes reviewer-only relationships', async () => { diff --git a/packages/worker/src/app/account-retention-dispositions.node.test.ts b/packages/worker/src/app/account-retention-dispositions.node.test.ts index c8dae3edea..1e063711c7 100644 --- a/packages/worker/src/app/account-retention-dispositions.node.test.ts +++ b/packages/worker/src/app/account-retention-dispositions.node.test.ts @@ -43,6 +43,15 @@ test('retention dispositions stay aligned with scheduled policies and documented disposition.kind === 'alternate_cleanup', ), ).toBe(true) + expect( + nonScheduled.some( + (disposition) => + disposition.table === 'entitlement_daily_counters' && + disposition.kind === 'alternate_cleanup' && + disposition.reason.includes('pending a later drop migration') && + disposition.reason.includes('no scheduled retention'), + ), + ).toBe(true) expect( nonScheduled.some( (disposition) => diff --git a/packages/worker/src/app/account-retention-dispositions.ts b/packages/worker/src/app/account-retention-dispositions.ts index d4f23694d0..292bb312af 100644 --- a/packages/worker/src/app/account-retention-dispositions.ts +++ b/packages/worker/src/app/account-retention-dispositions.ts @@ -36,7 +36,7 @@ export const accountRetentionDispositions: ReadonlyArray Date: Sat, 1 Aug 2026 16:50:38 +0000 Subject: [PATCH 37/39] docs(storage): point pending-drop guardrail correctly Co-authored-by: Kent C. Dodds --- docs/contributing/architecture/data-storage.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 54ebff1ad7..7a7068a5ce 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -74,9 +74,9 @@ Deletion must cover these user-owned surfaces: and listing children for community-owned listings. Physical tables pending a later drop migration are registered in `accountUserDataPendingDropTargets` (schema coverage only; runtime deletion never queries them). The guardrail - test in `packages/worker/src/app/account-deletion.node.test.ts` applies the - live migrations to SQLite and fails if a user-owned schema column is not - represented in the deletion target list or pending-drop registry, or if those + test in `packages/worker/src/account/data-targets.node.test.ts` applies the + live migrations to SQLite and fails if a user-owned schema column lacks schema + coverage in the runtime target list or pending-drop registry, or if those lists reference a stale column. - **Durable Objects:** `JobManager`, `StorageRunner`, `RepoSession`, `RemoteConnectorSession`, `PackageRealtimeSession`, `PackageServiceInstance`, From 659033c521de6358636df430c5433e79b807d88f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 17:15:14 +0000 Subject: [PATCH 38/39] refactor(entitlements): drop retired daily counter table Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 34 ++++--- .../contributing/architecture/entitlements.md | 86 ++++++++--------- .../0126-drop-entitlement-daily-counters.sql | 8 ++ .../src/account/data-targets.node.test.ts | 39 +++----- packages/worker/src/account/data-targets.ts | 13 +-- .../worker/src/account/export.node.test.ts | 9 +- ...ccount-retention-dispositions.node.test.ts | 11 +-- .../src/app/account-retention-dispositions.ts | 6 -- ...ment-daily-counters-migration.node.test.ts | 96 +++++++++++++++++++ tools/migration-ledger.json | 4 + 10 files changed, 193 insertions(+), 113 deletions(-) create mode 100644 packages/worker/migrations/0126-drop-entitlement-daily-counters.sql create mode 100644 packages/worker/src/entitlements/drop-entitlement-daily-counters-migration.node.test.ts diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 7a7068a5ce..75128bfaf2 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -131,10 +131,10 @@ every query or Durable Object lookup is scoped to that id. System email rows owned by `system:email` are intentionally absent from account exports for the same reason they are absent from deletion: they belong to the -operator inbox surface, not to the exporting user. Pending-drop physical tables -(such as quiescent `entitlement_daily_counters`) are also absent from export -queries. The export manifest lists both under `excludedD1Surfaces` so the -omission is explicit. +operator inbox surface, not to the exporting user. The export manifest lists +that omission under `excludedD1Surfaces` so it is explicit. The retired +`entitlement_daily_counters` D1 mirror is absent from the final schema +(migration `0126`) and therefore from export inventory and `excludedD1Surfaces`. Platform-feedback submissions are included in the submitting user's own D1 export section. An export never includes submissions owned by other users, @@ -644,12 +644,12 @@ time-pruned. Deletion-fence legacy lease rows are bounded by the D1 snapshot replace on `markDeleting` rather than time retention; DO-authority rows clear on release/repair/purge. -**D1 daily mirror writes stopped:** enforcement, point reads, bootstrap, mirror, -and account export/deletion inventory paths no longer read or write -`entitlement_daily_counters`. The admin parity report -(`admin_user_meter_parity`) temporarily reads the table while it exists for -migration verification. The physical table stays registered as a pending-drop -coverage exemption until a later migration-only deploy drops it. See +**D1 daily mirror retired:** enforcement, point reads, bootstrap, mirror, and +account export/deletion inventory paths never read or write +`entitlement_daily_counters`. The three-deploy retirement is complete (Workers +#1133 / #1134, then migration `0126-drop-entitlement-daily-counters.sql`). The +final live schema has no table or day index; `admin_user_meter_parity` reports +`daily.mirrorRetired: true` (meter counts only). See [Entitlements](./entitlements.md#usermeter-expand-phase). **Daily cold bootstrap:** a missing `(resource, day)` row returns @@ -1672,14 +1672,12 @@ Current retention policies: After Mailbox cut-over, the same 365-day window and blob-before-row ordering are self-enforced by the Mailbox DO alarm; `system:email` stays on the D1 system-email retention job. -- `entitlement_daily_counters`: **quiescent pending drop** — mirror writes, - bootstrap reads, scheduled retention pruning, and account export/deletion - inventory stopped across the stage 1/2 code deploys. The admin parity report - (`admin_user_meter_parity`) temporarily reads the table while it exists. - Runtime account inventory no longer queries the table; it remains registered - in `accountUserDataPendingDropTargets` for schema coverage until the later - drop migration. Daily counter retention lives in the per-user `UserMeter` DO - (`userMeterDailyCounterRetentionDays`). +- `entitlement_daily_counters`: **retired** — dropped by migration + `0126-drop-entitlement-daily-counters.sql` after stages 1/2 stopped mirror + writes and detached runtime inventory. No scheduled retention disposition or + pending-drop coverage remains. Daily counter retention lives in the per-user + `UserMeter` DO (`userMeterDailyCounterRetentionDays`); + `admin_user_meter_parity` reports `daily.mirrorRetired: true`. - `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; raw Analytics Engine usage events follow platform retention. - `feature_flag_exposure_rollups`: local-dev/test flag exposure rollups keep 90 diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 1d03451b13..8b40d21beb 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -152,16 +152,14 @@ then initializes that key at zero via `UserMeter.initialize()` (`INSERT OR IGNORE`, concurrent-safe) before retrying. Warm enforcement awaits only the DO RPC and never touches D1 daily counter state. -**D1 daily mirror writes stopped:** consume, refund, inbound charge/read, -point-read surfaces, retention, and generic account export/deletion no longer -read or write `entitlement_daily_counters`. The physical table remains in the -migrated schema as a pending-drop target (`accountUserDataPendingDropTargets`) -so schema coverage still recognizes `entitlement_daily_counters.user_id` without -runtime inventory querying it; export lists the omission under -`excludedD1Surfaces` (no raw rows). A later migration-only deploy drops the -table (migrations apply before Workers); existing rows are quiescent historical -mirror state. Analytics Engine remains the production reporting path for email -send/receive aggregates. +**D1 daily mirror retired:** consume, refund, inbound charge/read, point-read +surfaces, retention, and account export/deletion never read or write +`entitlement_daily_counters`. The three-deploy retirement is complete: Worker +#1133 stopped mirror writes, #1134 detached runtime inventory, and migration +`0126-drop-entitlement-daily-counters.sql` dropped the table and day index. +`admin_user_meter_parity` reports `daily.mirrorRetired: true` (meter counts +only; `d1Count`/`delta` null). Analytics Engine remains the production reporting +path for email send/receive aggregates. **Point-read surfaces** call `readDailyEntitlementResourceUsage` (UserMeter with the same cold zero-init path): @@ -362,39 +360,39 @@ same-token rollout mirrors; email keeps its D1 lease path. Package-service and storage authority flips remain separate high-risk contract follow-ups after soak/parity review. -**Daily-counter mirror retirement (three-deploy):** stage 1 (Worker #1133) -stopped all D1 mirror/bootstrap/retention use while leaving the physical -`entitlement_daily_counters` table and account inventory target in place. Stage -2 removes the runtime account export/deletion inventory target and registers the -table in `accountUserDataPendingDropTargets` for schema coverage (export -documents the omission; no raw rows). The third, migration-only deploy drops the -table. Production `admin_user_meter_parity` scans across 38 users showed zero -daily mismatches with Analytics Engine reporting active — the deploy rationale -for stopping mirror writes before the drop. While the table exists, parity still -compares `d1Count === meterCount` (`mirrorRetired: false`); after the drop -migration, `mirrorRetired: true` reports meter counts only. +**Daily-counter mirror retirement (three-deploy, complete):** stage 1 (Worker +#1133) stopped all D1 mirror/bootstrap/retention use while leaving the physical +`entitlement_daily_counters` table in place. Stage 2 (#1134) detached the +runtime account export/deletion inventory target and kept a pending-drop schema +coverage exemption while the table still existed. Stage 3 (migration +`0126-drop-entitlement-daily-counters.sql`) drops the table and +`idx_entitlement_daily_counters_day`. Production `admin_user_meter_parity` scans +across 38 users showed zero daily mismatches with Analytics Engine reporting +active — the deploy rationale for stopping mirror writes before the drop. Final +live schema has no `entitlement_daily_counters` table; admin parity reports +`daily.mirrorRetired: true` (meter counts only). ### Admin UserMeter parity gates (`admin_user_meter_parity`) Production verification for mirror retirement and remaining authority flips uses the admin-only read-only capability `admin_user_meter_parity` (input: `stable_user_id`). It compares production-shaped D1 rows for one account against -direct UserMeter RPCs and never bootstraps or writes parity state. Daily -comparison retires automatically once the drop migration removes -`entitlement_daily_counters` (`daily.mirrorRetired: true`). Opening a cold -UserMeter stub may still run Durable Object constructor schema maintenance and -opportunistic stale daily-counter pruning. Cold meter rows surface as -`needsBootstrap` with `meterCount`/`meterBytes` null. +direct UserMeter RPCs and never bootstraps or writes parity state. After +migration `0126`, daily comparison is retired (`daily.mirrorRetired: true`): the +report returns meter counts only with `d1Count`/`delta` null and each resource +`parity: true`. Opening a cold UserMeter stub may still run Durable Object +constructor schema maintenance and opportunistic stale daily-counter pruning. +Cold meter rows surface as `needsBootstrap` with `meterCount`/`meterBytes` null. Interpret the structured report as independent gates: -| Gate | Pass condition | -| -------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Daily counters (current UTC day) | While `daily.mirrorRetired` is false (table present): each of the four daily resources has `parity: true` (`d1Count === meterCount`); aggregate `daily.mismatchCount === 0`. After the drop migration (`mirrorRetired: true`): report meter counts only; `d1Count`/`delta` are null, each resource has `parity: true`, and `mismatchCount === 0` (no D1 comparison). | -| Storage bytes | `storage.parity` — D1 `users.d1_storage_bytes` equals UserMeter `readStorageBytes` (not `needsBootstrap`). | -| Package services | `packageServices.parity` — inventory mismatch category counts are all zero (`d1Only` / `meterOnly` / `statusMismatch` / `startedAtMismatch` / `sourceUpdatedAtMismatch`), fresh-running counts match under the shared 24h stale window, and the meter page walk is not `truncated`. | -| Deletion tombstone | `deletion.deletingAtParity` — D1 `users.deleting_at` matches the meter tombstone. | -| Temporary D1 lease mirror | `deletion.mirrorLeaseParity` — `doOnly === 0`, `legacyWithoutD1 === 0`, inventory not truncated, and `d1ActiveLeaseCount >= doAuthorityLeaseCount` (same-token mirror coverage). `tokenSetMismatches.d1Only` is reported but does **not** fail this gate. | +| Gate | Pass condition | +| -------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Daily counters (current UTC day) | Final post-drop semantics (`daily.mirrorRetired: true`): report meter counts only; `d1Count`/`delta` are null, each of the four daily resources has `parity: true`, and `mismatchCount === 0` (no D1 comparison). Pre-drop Workers that still see the table compare `d1Count === meterCount` with `mirrorRetired: false`. | +| Storage bytes | `storage.parity` — D1 `users.d1_storage_bytes` equals UserMeter `readStorageBytes` (not `needsBootstrap`). | +| Package services | `packageServices.parity` — inventory mismatch category counts are all zero (`d1Only` / `meterOnly` / `statusMismatch` / `startedAtMismatch` / `sourceUpdatedAtMismatch`), fresh-running counts match under the shared 24h stale window, and the meter page walk is not `truncated`. | +| Deletion tombstone | `deletion.deletingAtParity` — D1 `users.deleting_at` matches the meter tombstone. | +| Temporary D1 lease mirror | `deletion.mirrorLeaseParity` — `doOnly === 0`, `legacyWithoutD1 === 0`, inventory not truncated, and `d1ActiveLeaseCount >= doAuthorityLeaseCount` (same-token mirror coverage). `tokenSetMismatches.d1Only` is reported but does **not** fail this gate. | **D1-only leases:** email and other transition paths that omit `env` still take exact D1 leases, so `d1Only > 0` is expected until that handoff. Mirror-removal @@ -573,9 +571,8 @@ Rules: `UserMeter.initialize({ count: 0 })` (`INSERT OR IGNORE`) before retrying the consume. Concurrent cold callers cannot double-apply a non-zero baseline. - **D1 mirror writes stopped:** consume/refund/inbound charge/read paths never - touch `entitlement_daily_counters`; the physical table remains quiescent until - a follow-up migration drops it. + **D1 mirror retired:** consume/refund/inbound charge/read paths never touch + `entitlement_daily_counters`; migration `0126` dropped the table. A delivery claim remains charged when later storage fails. Cloudflare Email Routing retries replay that same `delivery_id` through @@ -792,10 +789,11 @@ in [`../environment-variables.md`](../environment-variables.md). `0066-stripe-billing.sql`; owned by `packages/worker/src/billing/`, read by `getUserPlan` via `resolveEffectivePlan`. `stripe_plan` stays nullable because it is Stripe-derived; `max` is manual-only. -- `entitlement_daily_counters` — **quiescent pending drop**. Created by - migration `0048-user-plans-and-entitlement-counters.sql`; mirror writes, - bootstrap reads, retention pruning, and account export/deletion inventory - stopped across the stage 1/2 code deploys. The table remains in - `accountUserDataPendingDropTargets` for schema coverage until a later - migration-only PR drops it. Daily counters are authoritative in the per-user - `UserMeter` DO; account export uses UserMeter RPCs (no raw D1 mirror rows). +- `entitlement_daily_counters` — **retired**. Created by migration + `0048-user-plans-and-entitlement-counters.sql`, indexed by + `0055-retention-indexes.sql`, and dropped by + `0126-drop-entitlement-daily-counters.sql` after stages 1/2 stopped mirror + writes and detached runtime inventory. Final live schema has no table or day + index; `admin_user_meter_parity` reports `daily.mirrorRetired: true`. Daily + counters are authoritative in the per-user `UserMeter` DO; account export uses + UserMeter RPCs. diff --git a/packages/worker/migrations/0126-drop-entitlement-daily-counters.sql b/packages/worker/migrations/0126-drop-entitlement-daily-counters.sql new file mode 100644 index 0000000000..96dbc6e1a5 --- /dev/null +++ b/packages/worker/migrations/0126-drop-entitlement-daily-counters.sql @@ -0,0 +1,8 @@ +-- Stage 3 (migration-only): drop the quiescent entitlement_daily_counters D1 +-- mirror after Worker #1133 stopped mirror writes and #1134 detached runtime +-- account export/deletion inventory. Authoritative daily counters live in the +-- per-user UserMeter Durable Object; admin_user_meter_parity reports +-- daily.mirrorRetired: true once this table is absent. + +DROP INDEX IF EXISTS idx_entitlement_daily_counters_day; +DROP TABLE IF EXISTS entitlement_daily_counters; diff --git a/packages/worker/src/account/data-targets.node.test.ts b/packages/worker/src/account/data-targets.node.test.ts index 336f7cf9a9..92807c0c27 100644 --- a/packages/worker/src/account/data-targets.node.test.ts +++ b/packages/worker/src/account/data-targets.node.test.ts @@ -214,26 +214,23 @@ test('every accountUserDataTargets kind has a shared match builder and export gu ).toBe(true) }) -test('pending-drop entitlement_daily_counters stays out of runtime deletion/export SQL while schema coverage holds', () => { - const pendingDrop = accountUserDataPendingDropTargets.find( - (target) => target.table === 'entitlement_daily_counters', - ) - expect(pendingDrop).toEqual( - expect.objectContaining({ - table: 'entitlement_daily_counters', - column: 'user_id', - surface: 'entitlement_daily_counters', - reason: expect.stringContaining('pending a later drop migration'), - }), - ) - expect(pendingDrop?.reason).toContain('No raw rows exported') - +test('final schema drops entitlement_daily_counters without stale inventory coverage', () => { + expect( + accountUserDataPendingDropTargets.some( + (target) => target.table === 'entitlement_daily_counters', + ), + ).toBe(false) expect( accountUserDataTargets.some( (target) => 'table' in target && target.table === 'entitlement_daily_counters', ), ).toBe(false) + expect(getAccountExportExcludedD1Surfaces()).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ name: 'entitlement_daily_counters' }), + ]), + ) const deletionStatements = accountUserDataTargets.map((target) => { const match = matchFor(target) @@ -251,15 +248,6 @@ test('pending-drop entitlement_daily_counters stays out of runtime deletion/expo }) expect(exportStatements.join('\n')).not.toMatch(/entitlement_daily_counters/u) - expect(getAccountExportExcludedD1Surfaces()).toEqual( - expect.arrayContaining([ - expect.objectContaining({ - name: 'entitlement_daily_counters', - reason: expect.stringContaining('No raw rows exported'), - }), - ]), - ) - const db = new DatabaseSync(':memory:') applyMigrations(db) const tableExists = db @@ -269,7 +257,7 @@ test('pending-drop entitlement_daily_counters stays out of runtime deletion/expo WHERE type = 'table' AND name = 'entitlement_daily_counters'`, ) .get() as { present: number } | undefined - expect(tableExists?.present).toBe(1) + expect(tableExists).toBeUndefined() const liveUserColumns = new Set() const tables = db @@ -291,7 +279,8 @@ test('pending-drop entitlement_daily_counters stays out of runtime deletion/expo } } const coveredColumns = getAccountD1UserColumnCoverage() - expect(coveredColumns.has('entitlement_daily_counters.user_id')).toBe(true) + expect(coveredColumns.has('entitlement_daily_counters.user_id')).toBe(false) + expect(liveUserColumns.has('entitlement_daily_counters.user_id')).toBe(false) const missing = [...liveUserColumns].filter( (column) => !coveredColumns.has(column), ) diff --git a/packages/worker/src/account/data-targets.ts b/packages/worker/src/account/data-targets.ts index 133eecd72f..3b918eabc8 100644 --- a/packages/worker/src/account/data-targets.ts +++ b/packages/worker/src/account/data-targets.ts @@ -79,7 +79,8 @@ export const accountUserDataExcludedOwnerIds = [ * `user_id` / `*_user_id` columns as covered so the live table does not look * like a missing inventory target; runtime deletion and export never query * them. Export documents each omission under `excludedD1Surfaces` (no raw - * rows). + * rows). Empty after `entitlement_daily_counters` was dropped by migration + * `0126`. */ export type AccountUserDataPendingDropTarget = { table: string @@ -89,15 +90,7 @@ export type AccountUserDataPendingDropTarget = { } export const accountUserDataPendingDropTargets: ReadonlyArray = - [ - { - table: 'entitlement_daily_counters', - column: 'user_id', - surface: 'entitlement_daily_counters', - reason: - 'Quiescent daily-counter D1 mirror pending a later drop migration after mirror-stop Worker #1133. Authoritative daily counters live in UserMeter; account export and deletion must not query this table before the drop. No raw rows exported.', - }, - ] + [] /** Targets that account export should skip (deletion still covers them). */ export function isExcludedFromAccountExport( diff --git a/packages/worker/src/account/export.node.test.ts b/packages/worker/src/account/export.node.test.ts index 6c5e19e4d1..f6d4e1c7b6 100644 --- a/packages/worker/src/account/export.node.test.ts +++ b/packages/worker/src/account/export.node.test.ts @@ -156,10 +156,11 @@ test('account export documents and excludes operator-owned system email rows', a name: 'system_email_inboxes', reason: expect.stringContaining('Operator-owned inbound mail'), }), - expect.objectContaining({ - name: 'entitlement_daily_counters', - reason: expect.stringContaining('pending a later drop migration'), - }), + ]), + ) + expect(accountExport.manifest.excludedD1Surfaces).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ name: 'entitlement_daily_counters' }), ]), ) expect(accountExport.d1).not.toHaveProperty('entitlement_daily_counters') diff --git a/packages/worker/src/app/account-retention-dispositions.node.test.ts b/packages/worker/src/app/account-retention-dispositions.node.test.ts index 1e063711c7..7747a7c78a 100644 --- a/packages/worker/src/app/account-retention-dispositions.node.test.ts +++ b/packages/worker/src/app/account-retention-dispositions.node.test.ts @@ -45,13 +45,12 @@ test('retention dispositions stay aligned with scheduled policies and documented ).toBe(true) expect( nonScheduled.some( - (disposition) => - disposition.table === 'entitlement_daily_counters' && - disposition.kind === 'alternate_cleanup' && - disposition.reason.includes('pending a later drop migration') && - disposition.reason.includes('no scheduled retention'), + (disposition) => disposition.table === 'entitlement_daily_counters', ), - ).toBe(true) + ).toBe(false) + expect(getRetentionPolicyCoverage().has('entitlement_daily_counters')).toBe( + false, + ) expect( nonScheduled.some( (disposition) => diff --git a/packages/worker/src/app/account-retention-dispositions.ts b/packages/worker/src/app/account-retention-dispositions.ts index 292bb312af..5178d8ace2 100644 --- a/packages/worker/src/app/account-retention-dispositions.ts +++ b/packages/worker/src/app/account-retention-dispositions.ts @@ -32,12 +32,6 @@ export const accountRetentionDispositions: ReadonlyArray file.endsWith('.sql')) + .sort() +} + +function applyMigration(db: DatabaseSync, fileName: string) { + db.exec(readFileSync(new URL(fileName, migrationsDirectory), 'utf8')) +} + +function applyMigrationsThrough( + db: DatabaseSync, + inclusiveUpperBound: string, + exclusiveLowerBound = '', +) { + for (const fileName of migrationFileNames().filter( + (file) => file > exclusiveLowerBound && file <= inclusiveUpperBound, + )) { + applyMigration(db, fileName) + } +} + +function applyAllMigrations(db: DatabaseSync) { + db.exec('PRAGMA foreign_keys = ON') + for (const fileName of migrationFileNames()) { + applyMigration(db, fileName) + } +} + +function schemaObjectPresent( + db: DatabaseSync, + type: 'table' | 'index', + name: string, +) { + return ( + db + .prepare( + `SELECT 1 AS present + FROM sqlite_schema + WHERE type = ? AND name = ?`, + ) + .get(type, name) as { present: number } | undefined + )?.present +} + +test('0048 creates entitlement_daily_counters, 0055 indexes it, and 0126 drops both', () => { + const db = new DatabaseSync(':memory:') + db.exec('PRAGMA foreign_keys = ON') + + applyMigrationsThrough(db, createEntitlementDailyCountersMigration) + expect(schemaObjectPresent(db, 'table', 'entitlement_daily_counters')).toBe(1) + expect( + schemaObjectPresent(db, 'index', 'idx_entitlement_daily_counters_day'), + ).toBeUndefined() + + applyMigrationsThrough( + db, + retentionIndexesMigration, + createEntitlementDailyCountersMigration, + ) + expect(schemaObjectPresent(db, 'table', 'entitlement_daily_counters')).toBe(1) + expect( + schemaObjectPresent(db, 'index', 'idx_entitlement_daily_counters_day'), + ).toBe(1) + + applyMigration(db, dropEntitlementDailyCountersMigration) + expect( + schemaObjectPresent(db, 'table', 'entitlement_daily_counters'), + ).toBeUndefined() + expect( + schemaObjectPresent(db, 'index', 'idx_entitlement_daily_counters_day'), + ).toBeUndefined() +}) + +test('applying all migrations leaves no entitlement_daily_counters table or index', () => { + const db = new DatabaseSync(':memory:') + applyAllMigrations(db) + expect( + schemaObjectPresent(db, 'table', 'entitlement_daily_counters'), + ).toBeUndefined() + expect( + schemaObjectPresent(db, 'index', 'idx_entitlement_daily_counters_day'), + ).toBeUndefined() +}) diff --git a/tools/migration-ledger.json b/tools/migration-ledger.json index 070dc3f59f..396dc6b4b8 100644 --- a/tools/migration-ledger.json +++ b/tools/migration-ledger.json @@ -515,6 +515,10 @@ { "filename": "0125-mailbox-parity-state.sql", "sha256": "6d2d54b8d8a2cc7399c4f0dbb8069cfa5d6686958d8fef1a3b874c472a564754" + }, + { + "filename": "0126-drop-entitlement-daily-counters.sql", + "sha256": "1fd413a77395d0dd183533863fe381ee5c006e75c8cf4224a4f00c1a57496e26" } ] } From 65b2828108a54a2d8c03f280eb8a9b6253d970a1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 1 Aug 2026 17:28:40 +0000 Subject: [PATCH 39/39] docs(entitlements): format retirement references Co-authored-by: Kent C. Dodds --- docs/contributing/architecture/data-storage.md | 6 +++--- docs/contributing/architecture/entitlements.md | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 75128bfaf2..a8edec4bbd 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -647,9 +647,9 @@ release/repair/purge. **D1 daily mirror retired:** enforcement, point reads, bootstrap, mirror, and account export/deletion inventory paths never read or write `entitlement_daily_counters`. The three-deploy retirement is complete (Workers -#1133 / #1134, then migration `0126-drop-entitlement-daily-counters.sql`). The -final live schema has no table or day index; `admin_user_meter_parity` reports -`daily.mirrorRetired: true` (meter counts only). See +`#1133` / `#1134`, then migration `0126-drop-entitlement-daily-counters.sql`). +The final live schema has no table or day index; `admin_user_meter_parity` +reports `daily.mirrorRetired: true` (meter counts only). See [Entitlements](./entitlements.md#usermeter-expand-phase). **Daily cold bootstrap:** a missing `(resource, day)` row returns diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index 8b40d21beb..c0e5bdb1aa 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -155,7 +155,7 @@ only the DO RPC and never touches D1 daily counter state. **D1 daily mirror retired:** consume, refund, inbound charge/read, point-read surfaces, retention, and account export/deletion never read or write `entitlement_daily_counters`. The three-deploy retirement is complete: Worker -#1133 stopped mirror writes, #1134 detached runtime inventory, and migration +`#1133` stopped mirror writes, `#1134` detached runtime inventory, and migration `0126-drop-entitlement-daily-counters.sql` dropped the table and day index. `admin_user_meter_parity` reports `daily.mirrorRetired: true` (meter counts only; `d1Count`/`delta` null). Analytics Engine remains the production reporting @@ -361,10 +361,10 @@ storage authority flips remain separate high-risk contract follow-ups after soak/parity review. **Daily-counter mirror retirement (three-deploy, complete):** stage 1 (Worker -#1133) stopped all D1 mirror/bootstrap/retention use while leaving the physical -`entitlement_daily_counters` table in place. Stage 2 (#1134) detached the -runtime account export/deletion inventory target and kept a pending-drop schema -coverage exemption while the table still existed. Stage 3 (migration +`#1133`) stopped all D1 mirror/bootstrap/retention use while leaving the +physical `entitlement_daily_counters` table in place. Stage 2 (`#1134`) detached +the runtime account export/deletion inventory target and kept a pending-drop +schema coverage exemption while the table still existed. Stage 3 (migration `0126-drop-entitlement-daily-counters.sql`) drops the table and `idx_entitlement_daily_counters_day`. Production `admin_user_meter_parity` scans across 38 users showed zero daily mismatches with Analytics Engine reporting