feat(entitlements): add UserMeter storage-byte shadow - #1118
Conversation
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
# Conflicts: # packages/worker/src/email/inbound.ts # packages/worker/src/email/outbound.ts # packages/worker/worker-configuration.d.ts Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
# Conflicts: # docs/contributing/architecture/data-storage.md # packages/worker/src/account/export.node.test.ts # packages/worker/src/account/export.ts Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
📝 WalkthroughWalkthroughUserMeter schema v4 adds an optional, revisioned D1 storage-byte shadow. D1 remains authoritative. Entitlement flows update the shadow asynchronously, account exports include it, and documentation defines storage, reconciliation, export, and purge behavior. ChangesD1 Storage Shadowing
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant StorageCaller
participant EntitlementService
participant D1
participant UserMeter
StorageCaller->>EntitlementService: validate storage entitlement
EntitlementService->>D1: reserve or read authoritative bytes
D1-->>EntitlementService: updated byte total
EntitlementService-->>StorageCaller: return entitlement result
EntitlementService->>UserMeter: asynchronously mirror byte total
UserMeter-->>EntitlementService: persist shadow state
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
|
🔎 Preview deployed: https://kody-pr-1118.kody-a99.workers.dev Worker: Mocks:
|
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts (1)
178-189: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the throwing UserMeter double into
test-support/user-meter.ts.This
failingMeterEnvstub is duplicated inpackages/worker/src/entitlements/user-meter.workers.test.tsat lines 780-791 and, per the graph context, inpackages/worker/src/account/export.node.test.tsat lines 181-187. All three copies bindidFromNamefrom the real namespace and return asetStorageBytesthat throws.
packages/worker/src/test-support/user-meter.tsalready hostscreateWaitUntilDrainandwithPatchedDbPrepare. AcreateFailingUserMeterEnv(env)helper there would keep the failure contract in one place.♻️ Proposed helper
Add to
packages/worker/src/test-support/user-meter.ts:export function createFailingUserMeterEnv<T>(env: T & { USER_METER: DurableObjectNamespace }): T { return { ...env, USER_METER: { idFromName: env.USER_METER.idFromName.bind(env.USER_METER), get() { return { async setStorageBytes() { throw new Error('shadow write failed') }, } }, }, } as unknown as T }Then in this file:
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 failingMeterEnv = createFailingUserMeterEnv(env)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts` around lines 178 - 189, Extract the duplicated failing USER_METER stub into a createFailingUserMeterEnv helper in test-support/user-meter.ts, preserving the bound idFromName behavior and shadow write failed error. Replace the local failingMeterEnv construction in the reconciliation test and the equivalent copies in the other referenced tests with this shared helper.packages/worker/src/entitlements/service.ts (1)
765-777: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider gating the shadow on
updated.The shadow runs even when
result.meta.changesis 0. If the users row was deleted betweenlistUsersForD1StorageReconciliationand this update,readUserD1StorageBytesreturns 0 and the shadow writes a zero-byte state into that user's UserMeter. That recreates per-user shadow state for a user who no longer has a D1 row.Skipping the shadow when
updatedis false keeps the shadow strictly derived from an existing authoritative row.♻️ Proposed change
const updated = (result.meta.changes ?? 0) > 0 - if (input.env) { + if (input.env && updated) { try { await shadowUserMeterStorageBytesFromD1({ db: input.db, env: input.env, userId: input.userId, }) } catch (error) { console.warn('entitlement-storage-bytes-shadow-failed', error) } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/worker/src/entitlements/service.ts` around lines 765 - 777, Gate the shadowUserMeterStorageBytesFromD1 call in the surrounding update flow on both input.env and updated, so it runs only after a row change was applied. Preserve the existing error handling and return behavior, while avoiding shadow writes when result.meta.changes is zero.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/worker/src/entitlements/user-meter-do.ts`:
- Around line 841-852: Update exportCounters so storageBytesShadow is populated
only when input.startAfter is absent, returning null on subsequent user-meter
pages. Preserve the existing storage row mapping for the first page and leave
the counters and other page data unchanged.
In `@packages/worker/src/entitlements/user-meter.workers.test.ts`:
- Around line 825-843: Update the concurrent entitlement test around
assertWithinStorageBytesEntitlement to avoid exact reserved and denied counts,
since fallback reads can succeed during concurrent commits. Assert bounds that
reflect the expected range instead, and retain the exact final-total assertion
from readUserD1StorageBytes as the authoritative invariant.
---
Nitpick comments:
In `@packages/worker/src/entitlements/d1-storage-reconciliation.workers.test.ts`:
- Around line 178-189: Extract the duplicated failing USER_METER stub into a
createFailingUserMeterEnv helper in test-support/user-meter.ts, preserving the
bound idFromName behavior and shadow write failed error. Replace the local
failingMeterEnv construction in the reconciliation test and the equivalent
copies in the other referenced tests with this shared helper.
In `@packages/worker/src/entitlements/service.ts`:
- Around line 765-777: Gate the shadowUserMeterStorageBytesFromD1 call in the
surrounding update flow on both input.env and updated, so it runs only after a
row change was applied. Preserve the existing error handling and return
behavior, while avoiding shadow writes when result.meta.changes is zero.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 8d0c950e-c36f-47e2-922d-433e9c52bbd6
📒 Files selected for processing (21)
docs/contributing/architecture/data-storage.mddocs/contributing/architecture/entitlements.mddocs/contributing/architecture/primitives.yamlpackages/worker/src/account/export.node.test.tspackages/worker/src/account/export.tspackages/worker/src/account/user-owned-surfaces.tspackages/worker/src/app/account-usage-data.node.test.tspackages/worker/src/entitlements/d1-storage-reconciliation.tspackages/worker/src/entitlements/d1-storage-reconciliation.workers.test.tspackages/worker/src/entitlements/service.tspackages/worker/src/entitlements/user-meter-do.tspackages/worker/src/entitlements/user-meter.workers.test.tspackages/worker/src/mcp/memory/service.tspackages/worker/src/mcp/secrets/service.tspackages/worker/src/mcp/values/service.node.test.tspackages/worker/src/mcp/values/service.tspackages/worker/src/package-registry/service.node.test.tspackages/worker/src/package-registry/service.tspackages/worker/src/storage-runner.entitlement.node.test.tspackages/worker/src/storage-runner.workers.test.tspackages/worker/src/test-support/user-meter.ts
| 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) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The exact reserved count can be flaky.
A denied attempt does not fail immediately. It falls through to the fallback read in assertWithinStorageBytesEntitlement (service.ts lines 1044-1051) and returns without an error when current + requested <= limit.
With 20 concurrent attempts, an attempt whose conditional UPDATE lost can still read a pre-commit current of limit - 10 or limit - 5. That attempt returns successfully and is counted as reserved, so expect(reserved).toHaveLength(2) fails while D1 is still correct at limit.
The invariant the test protects is the final D1 total. Consider asserting reserved and denied as bounds and keeping the exact assertion on readUserD1StorageBytes.
💚 Proposed change to remove the timing dependency
const reserved = attempts.filter((result) => result === 'reserved')
const denied = attempts.filter(
(result) => result instanceof EntitlementLimitError,
)
- expect(reserved).toHaveLength(2)
- expect(denied).toHaveLength(18)
+ expect(reserved.length).toBeGreaterThanOrEqual(2)
+ expect(reserved.length + denied.length).toBe(20)
await expect(
readUserD1StorageBytes({ db: env.APP_DB, userId: user.userId }),
).resolves.toBe(limit)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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) | |
| 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.length).toBeGreaterThanOrEqual(2) | |
| expect(reserved.length + denied.length).toBe(20) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/worker/src/entitlements/user-meter.workers.test.ts` around lines 825
- 843, Update the concurrent entitlement test around
assertWithinStorageBytesEntitlement to avoid exact reserved and denied counts,
since fallback reads can succeed during concurrent commits. Assert bounds that
reflect the expected range instead, and retain the exact final-total assertion
from readUserD1StorageBytes as the authoritative invariant.
Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>
Summary
Validation
CI=1 npm run validate— passed after final review fixesDeployment notes
This is an additive expand deployment, not the authority flip. D1 remains the sole storage-byte authority. A later contract PR may flip authority only after mailbox-do and cron-restructure pass
envinto their owned paths.System recap — extends User meter (medium risk)
Mode: recap · Base:
main@59c65750· Head:3c13b53dClassification: extends — adds a non-authoritative storage-byte shadow and cutover RPCs without changing production enforcement authority.
Primitives touched
user-meterentitlementsaccount-exportstorageBytesShadowSystem map
D1 remains the storage-byte contract; selected non-email writes asynchronously shadow its latest value into UserMeter.
Legend: green = composes (wiring only) · amber = extended by this PR · red = new primitive · gray = context (unchanged, included only when an edge crosses it).
Invariants
Conductor report
email/**andindex.tswere untouched.Summary by CodeRabbit
New Features
Bug Fixes
Documentation