Repository navigation
Cut sequential D1/KV round trips from the execute hot path - #1035
Merged
Merged
Conversation
- Reentrant account write leases: nested withAccountWriteLease calls for the same user reuse the outer lease (MCP call -> package invocation previously paid ~5 extra D1 statements per invoke) - Short-TTL (60s) per-binding caches for quota plan resolution and the fetch gateway's account reverse-resolution; the atomic daily counter write stays uncached - packages.invokeChecked resolves the package once: the check phase skips the full-source export projection it discards and hands its saved package + manifest + bundle artifact to the invoke phase - Reuse the already-loaded entity source row instead of re-querying when the artifact's sourceId matches - Thread the MCP Durable Object's waitUntil through the execute tool so terminal run-record writes and error-subscription dispatch no longer serialize the response
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThe changes add async-local account lease tracking, per-D1 entitlement caching, reusable package-invocation preloads, user-scoped source resolution, and Durable Object ChangesAccount write lease reuse
Entitlement lookup caching
Preloaded package execution and background work
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant invokeChecked
participant checkPackageInvokeForRuntimeWithPreloads
participant invokePackageExportForExecuteRuntime
participant invokeSavedPackageModule
invokeChecked->>checkPackageInvokeForRuntimeWithPreloads: check and preload package artifacts
checkPackageInvokeForRuntimeWithPreloads-->>invokeChecked: result and preloads
invokeChecked->>invokePackageExportForExecuteRuntime: invoke with preloads
invokePackageExportForExecuteRuntime->>invokeSavedPackageModule: pass preloaded module artifact
invokeSavedPackageModule-->>invokeChecked: execute package export
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 |
kody-bot
marked this pull request as ready for review
July 30, 2026 01:55
Contributor
|
🔎 Preview deployed: https://kody-pr-1035.kody-a99.workers.dev Worker: Mocks:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/account/deletion-state.ts`:
- Around line 94-112: Update withAccountWriteLease and
heldAccountWriteLeaseStorage so the AsyncLocalStorage state reflects lease
lifetime, preventing child or detached async work created by write from treating
a released lease as active. Track and invalidate each user’s active lease state
(or equivalent nested-call guard) when the outer acquisition completes, while
preserving synchronous nested writes during the active lease; add a regression
test where a child spawned inside write calls withAccountWriteLease only after
the outer call resolves.
In `@packages/worker/src/package-invocations/idempotent-module-invocation.ts`:
- Around line 296-303: Update the repoSource resolution around
getEntitySourceById so the artifact source lookup is scoped to
input.actor.userId. Pass the user ID to the existing helper or extend
getEntitySourceById to require and apply that user constraint, while preserving
the current sourceRow path and source-ID matching behavior.
🪄 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: 69f9b29c-fbcf-4c5e-83a7-277cd0bd4155
📒 Files selected for processing (12)
packages/worker/src/account/deletion-state.node.test.tspackages/worker/src/account/deletion-state.tspackages/worker/src/entitlements/entitlements.node.test.tspackages/worker/src/entitlements/service.tspackages/worker/src/mcp/fetch-gateway.tspackages/worker/src/mcp/index.tspackages/worker/src/mcp/mcp-registration-agent.tspackages/worker/src/mcp/tools/execute.tspackages/worker/src/package-invocations/http-invoke.tspackages/worker/src/package-invocations/idempotent-module-invocation.tspackages/worker/src/package-invocations/invoke-check.tspackages/worker/src/package-invocations/runtime-tool-factories.ts
Comment on lines
+94
to
+112
| const heldAccountWriteLeaseStorage = new AsyncLocalStorage< | ||
| ReadonlySet<string> | ||
| >() | ||
|
|
||
| export async function withAccountWriteLease<T>(input: { | ||
| db: D1Database | ||
| stableUserId: string | ||
| holder?: string | ||
| write: () => Promise<T> | ||
| }) { | ||
| const heldLeases = heldAccountWriteLeaseStorage.getStore() | ||
| if (heldLeases?.has(input.stableUserId)) { | ||
| return await input.write() | ||
| } | ||
| const nextHeldLeases = new Set(heldLeases) | ||
| nextHeldLeases.add(input.stableUserId) | ||
| return await heldAccountWriteLeaseStorage.run(nextHeldLeases, async () => | ||
| acquireAccountWriteLeaseAndWrite(input), | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect lease callers for detached work that can inherit AsyncLocalStorage context.
ast-grep outline packages/worker/src/account/deletion-state.ts --items all
rg -n -C3 --type=ts '\bwithAccountWriteLease\s*\(' packages/worker/src
rg -n -C3 --type=ts '\bvoid\s+|\bwaitUntil\s*\(' \
packages/worker/src/package-invocations packages/worker/src/mcpRepository: kentcdodds/kody
Length of output: 41280
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== deletion-state outline/sections =="
wc -l packages/worker/src/account/deletion-state.ts
sed -n '1,280p' packages/worker/src/account/deletion-state.ts
echo "== related test sections =="
wc -l packages/worker/src/account/deletion-state.node.test.ts
sed -n '1,240p' packages/worker/src/account/deletion-state.node.test.ts
echo "== package invocation common/invoker waitUntil context candidates =="
sed -n '75,140p' packages/worker/src/package-invocations/common.ts
rg -n -C4 --type=ts 'waitForCompletion|waitUntilFrom|waitUntil|ctx=|executionContext|executor|run' packages/worker/src/package-invocations packages/worker/src/mcp/src packages/worker/src/mcp 2>/dev/null || trueRepository: kentcdodds/kody
Length of output: 50372
Prevent detached async work from reusing a released lease.
withAccountWriteLease puts the user in heldAccountWriteLeaseStorage, and that AsyncLocalStorage context is captured by child promises spawned in write. If a void/waitUntil callback invokes this function after the outer finally removes the user, the nested call believes the lease is still held and can run a write while the lease row was released. Track lease lifetime/active count per user or explicitly invalidate nested calls in detached work; add a regression test for a child created inside the outer callback that calls withAccountWriteLease only after the outer call resolves.
🤖 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/account/deletion-state.ts` around lines 94 - 112, Update
withAccountWriteLease and heldAccountWriteLeaseStorage so the AsyncLocalStorage
state reflects lease lifetime, preventing child or detached async work created
by write from treating a released lease as active. Track and invalidate each
user’s active lease state (or equivalent nested-call guard) when the outer
acquisition completes, while preserving synchronous nested writes during the
active lease; add a regression test where a child spawned inside write calls
withAccountWriteLease only after the outer call resolves.
…lookup by user - Lease reentrancy now tracks per-user frames that deactivate at release, so detached waitUntil work inheriting the async context re-acquires its own lease instead of reusing a released one (with regression test) - Package invocation repo-source fallback uses the user-scoped getEntitySourceByIdForUser lookup
This was referenced Jul 30, 2026
Merged
This was referenced Aug 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Execute performance: root-cause analysis and hot-path fixes
What production data shows
Queried
kody_usage_events(Analytics Engine) forexecutesandbox durations:Two distinct findings:
kody.*capability dispatch is fine warm (~40–105ms) but the first dispatch in a cold isolate costs ~1.5s, every sandboxfetchpays a D1 plan read + counter write in the gateway (Add bad-actor abuse controls: suspension, outbound-email pause, compute quotas, delivery insights #911, Jul 24), and one warmpackages.invokeCheckedof a trivial export costs ~950ms–1.9s — it is a long chain of sequential D1/KV/DO round trips. Agent workflows chain package invocations (github, discord), so tool calls stack multiple seconds of platform overhead on top of real API time.Measured warm
packages.invokeCheckedbreakdown (before this PR): check phase resolves the package + manifest + artifact and loads the full package source for a projection that invokeChecked discards; the invoke phase then re-resolves the package, re-loads manifest + artifact, acquires a second account write lease (~5 D1 statements — the MCP call already holds one), inserts/updates idempotency rows, and awaits run-record Durable Object writes before responding.Changes
account/deletion-state.ts): nestedwithAccountWriteLeasecalls for the same user inside the same async context reuse the outer lease. The MCP request boundary, app requests, and job runs all take a lease, so every nested package invocation previously paid a full extra acquire/release (~5 D1 statements, 4 of them writes). Deletion remains blocked for exactly the outer lease's span.entitlements/service.ts):getCachedUserPlanandfindCachedUserAccountByStableUserId, keyed by theD1Databasebinding, used byconsumeDailyEntitlement(every execute call and every sandbox outbound fetch) and the fetch gateway's email reverse-resolution. The atomic daily counter upsert stays uncached; only plan-limit resolution tolerates ≤60s staleness. Failures are never cached.packages.invokeCheckedresolves its package once (package-invocations/*): the check phase skips the full-source export projection it discards, and hands its saved-package row + manifest + prepared bundle artifact to the invoke phase, which no longer re-resolves or re-loads them. Barepackages.checkbehavior is unchanged.getEntitySourceByIdre-query when the artifact'ssourceIdmatches the row already loaded with the manifest.waitUntilthrough the execute tool (mcp/index.ts,tools/execute.ts): terminal run-record writes, nested-invocation run records, and error-subscription dispatch now ridectx.waitUntilinstead of serializing the response (the options and never-block contract already existed; the public execute tool just never passed it).Testing
npm run validatepasses (1578 unit tests across 463 files, 20 Playwright e2e, MCP e2e, lint/format/typecheck/structure checks).Follow-ups (not this PR)
System recap — extends existing primitives (medium risk)
Mode: recap · Base:
main@7ebea332· Head:fa6d87e4Classification: extends — no new primitives; changes the account write-lease contract (reentrant per async context), entitlement plan resolution (60s cached for quota checks), and the package-invocation check/invoke contract (single load).
Primitives touched
mcp-serverwaitUntil; fetch gateway uses cached account resolveentitlementsconsumeDailyEntitlementresolves plan limits through a 60s per-binding cache; counter writes unchangedpackage-invocations(unmatched paths)System map
The MCP execute tool's nested package invocations now reuse the request's account write lease and the check phase's package loads instead of re-reading D1/KV.
Legend: green = composes (wiring only) · amber = extended by this PR · red = new primitive · gray = context (unchanged, included only when an edge crosses it).
Invariants
Per-user isolation is preserved: lease reentrancy is keyed by
stableUserIdwithin one async context; plan caches are keyed by user id + email perD1Databasebinding; preloaded package artifacts come from the same user-scoped check call.Summary by CodeRabbit
waitUntil.