From fa6d87e4ef67c4607546b51419aa9be6afd50b5b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 30 Jul 2026 01:50:50 +0000 Subject: [PATCH 1/2] Cut sequential D1/KV round trips from the execute hot path - 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 --- .../src/account/deletion-state.node.test.ts | 108 ++++++++++++ packages/worker/src/account/deletion-state.ts | 32 ++++ .../entitlements/entitlements.node.test.ts | 106 ++++++++++++ packages/worker/src/entitlements/service.ts | 100 ++++++++++- packages/worker/src/mcp/fetch-gateway.ts | 4 +- packages/worker/src/mcp/index.ts | 3 + .../worker/src/mcp/mcp-registration-agent.ts | 6 + packages/worker/src/mcp/tools/execute.ts | 6 + .../src/package-invocations/http-invoke.ts | 31 ++-- .../idempotent-module-invocation.ts | 31 ++-- .../src/package-invocations/invoke-check.ts | 161 +++++++++++++----- .../runtime-tool-factories.ts | 28 ++- 12 files changed, 546 insertions(+), 70 deletions(-) diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index ac749aefc8..c9f2358c1d 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -86,3 +86,111 @@ test('non-expiring lease blocks deletion until audited repair fences callback', reason: 'Inspected worker crash and confirmed process termination.', }) }) + +function createLeaseTestDb() { + const sqlite = new DatabaseSync(':memory:') + sqlite.exec(` + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + stable_user_id TEXT UNIQUE, + deleting_at TEXT, + active_write_count INTEGER NOT NULL DEFAULT 0, + updated_at TEXT + ); + 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 (id, stable_user_id) VALUES (1, 'user-a'); + INSERT INTO users (id, stable_user_id) VALUES (2, 'user-b'); + `) + return { sqlite, db: createD1FromSqlite(sqlite) } +} + +function countLeaseRows(sqlite: DatabaseSync) { + const row = sqlite + .prepare(`SELECT COUNT(*) AS count FROM account_write_leases`) + .get() as { count: number } + return Number(row.count) +} + +test('nested same-user lease reuses the outer lease instead of re-acquiring', async () => { + const { sqlite, db } = createLeaseTestDb() + const result = await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:outer', + async write() { + expect(countLeaseRows(sqlite)).toBe(1) + return await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:nested', + async write() { + // Still only the outer lease: the nested call must not pay + // another acquire/release round trip. + expect(countLeaseRows(sqlite)).toBe(1) + const active = await listActiveAccountWriteLeases(db, 'user-a') + expect(active).toHaveLength(1) + expect(active[0]).toEqual( + expect.objectContaining({ holder: 'test:outer' }), + ) + return 'nested-result' + }, + }) + }, + }) + expect(result).toBe('nested-result') + // The outer release is the only release. + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) +}) + +test('nested lease for a different user still acquires its own lease', async () => { + const { sqlite, db } = createLeaseTestDb() + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:outer', + async write() { + await withAccountWriteLease({ + db, + stableUserId: 'user-b', + holder: 'test:other-user', + async write() { + expect(countLeaseRows(sqlite)).toBe(2) + expect(await listActiveAccountWriteLeases(db, 'user-b')).toHaveLength( + 1, + ) + }, + }) + }, + }) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + expect(await listActiveAccountWriteLeases(db, 'user-b')).toHaveLength(0) +}) + +test('sequential sibling leases each acquire after the previous released', async () => { + const { sqlite, db } = createLeaseTestDb() + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + async write() { + expect(countLeaseRows(sqlite)).toBe(1) + }, + }) + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + async write() { + expect(countLeaseRows(sqlite)).toBe(2) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1) + }, + }) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) +}) diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index e76f547105..0a330d5457 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from 'node:async_hooks' import { utcSqliteTimestamp } from '@kody-internal/shared/date-keys.ts' export class AccountDeletionInProgressError extends Error { @@ -80,11 +81,42 @@ export async function assertAccountWritable(env: Env, stableUserId: string) { await assertAccountWritableDb(env.APP_DB, stableUserId) } +/** + * Stable user ids whose write lease is already held somewhere up the current + * async call chain. Nested {@link withAccountWriteLease} calls for the same + * user reuse the outer lease 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. + */ +const heldAccountWriteLeaseStorage = new AsyncLocalStorage< + ReadonlySet +>() + export async function withAccountWriteLease(input: { db: D1Database stableUserId: string holder?: string write: () => Promise +}) { + 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), + ) +} + +async function acquireAccountWriteLeaseAndWrite(input: { + db: D1Database + stableUserId: string + holder?: string + write: () => Promise }) { if (typeof input.db.batch !== 'function') { const acquired = await input.db diff --git a/packages/worker/src/entitlements/entitlements.node.test.ts b/packages/worker/src/entitlements/entitlements.node.test.ts index 972ec2b159..5897a5c4b1 100644 --- a/packages/worker/src/entitlements/entitlements.node.test.ts +++ b/packages/worker/src/entitlements/entitlements.node.test.ts @@ -31,6 +31,8 @@ import { countRunningPackageServices, estimateEntitlementStorageEntryByteDelta, estimateEntitlementStorageEntryBytes, + findCachedUserAccountByStableUserId, + getCachedUserPlan, getUserPlan, incrementDailyEntitlementCounter, packageServiceStateStaleMs, @@ -136,6 +138,21 @@ function createEntitlementsTestDb( : null ) as T | null } + if (query.includes('SELECT email, plan, email_verified_at')) { + const stableUserId = params[0] + const user = users.find( + (row) => row.stable_user_id === stableUserId, + ) + return ( + user + ? { + email: user.email, + plan: user.plan, + email_verified_at: null, + } + : null + ) as T | null + } if (query.includes('FROM entitlement_daily_counters')) { const row = counters.find( (counter) => @@ -374,6 +391,95 @@ test('getUserPlan resolves plans via email+stable id and short-circuits invalid expect(consoleWarn).toHaveBeenCalledWith(unknownStoredPlanWarningTag) }) +test('getCachedUserPlan caches per db binding and never caches failures', async () => { + const userId = await createStableUserIdFromEmail(plannedEmail) + const users = [{ email: plannedEmail, plan: 'pro', stable_user_id: userId }] + const { db, queries } = createEntitlementsTestDb({ users }) + + expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe( + 'pro', + ) + expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe( + 'pro', + ) + const planQueries = () => + queries.filter((query) => + query.sql.includes('email = ? AND stable_user_id = ?'), + ) + expect(planQueries()).toHaveLength(1) + + // A plan change is visible to the uncached lookup immediately and to the + // cached lookup only after the TTL: quota checks tolerate that staleness. + users[0]!.plan = 'free' + expect(await getUserPlan(db, { userId, email: plannedEmail })).toBe('free') + expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe( + 'pro', + ) + + // Another db binding (fresh test database) never shares cache entries. + const second = createEntitlementsTestDb({ + users: [{ email: plannedEmail, plan: 'free', stable_user_id: userId }], + }) + expect( + await getCachedUserPlan(second.db, { userId, email: plannedEmail }), + ).toBe('free') + + // Anonymous / invalid ids short-circuit without touching the cache or D1. + expect(await getCachedUserPlan(db, { userId, email: null })).toBe('max') + expect( + await getCachedUserPlan(db, { userId: 'user-1', email: plannedEmail }), + ).toBe('max') + + // Failures are not pinned for the TTL: the next call retries D1. + let firstCall = true + const flaky = { + prepare() { + return { + bind() { + return { + async first() { + if (firstCall) { + firstCall = false + throw new Error('D1 blip') + } + return { plan: 'pro', stripe_plan: null } + }, + } + }, + } + }, + } as unknown as D1Database + await expect( + getCachedUserPlan(flaky, { userId, email: plannedEmail }), + ).rejects.toThrow('D1 blip') + expect(await getCachedUserPlan(flaky, { userId, email: plannedEmail })).toBe( + 'pro', + ) +}) + +test('findCachedUserAccountByStableUserId caches the account reverse-resolution per db', async () => { + const userId = await createStableUserIdFromEmail(plannedEmail) + const { db, queries } = createEntitlementsTestDb({ + users: [{ email: plannedEmail, plan: 'pro', stable_user_id: userId }], + }) + const accountQueries = () => + queries.filter((query) => + query.sql.includes('SELECT email, plan, email_verified_at'), + ) + expect(await findCachedUserAccountByStableUserId(db, userId)).toEqual({ + email: plannedEmail, + plan: 'pro', + emailVerified: false, + }) + expect(await findCachedUserAccountByStableUserId(db, userId)).toEqual({ + email: plannedEmail, + plan: 'pro', + emailVerified: false, + }) + expect(accountQueries()).toHaveLength(1) + expect(await findCachedUserAccountByStableUserId(db, ' ')).toBeNull() +}) + test('assertWithinEntitlement passes under the limit, throws at it, and enforces finite max ordinary limits', async () => { const userId = await createStableUserIdFromEmail(plannedEmail) const freeLimit = planLimits.free.maxScheduledJobs diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index 00d7787738..e1bb2bc1ba 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -54,6 +54,101 @@ export type StableUserAccount = { emailVerified: boolean } +/** + * Short-TTL caches for the plan / account lookups that run on metered hot + * paths (every execute call and every sandbox outbound fetch pays one). The + * quota counter write itself stays atomic and uncached; only the plan-limit + * resolution tolerates staleness, so a plan change takes effect for quota + * checks within {@link entitlementLookupCacheTtlMs} instead of immediately. + * Keyed by the `D1Database` binding so test databases never share entries. + */ +const entitlementLookupCacheTtlMs = 60_000 +const entitlementLookupCacheMaxEntries = 1_000 + +type EntitlementLookupCacheEntry = { + value: Promise + expiresAtMs: number +} + +function createEntitlementLookupCache() { + const cachesByDb = new WeakMap< + D1Database, + Map> + >() + return { + async getOrCreate( + db: D1Database, + key: string, + create: () => Promise, + ): Promise { + let cache = cachesByDb.get(db) + if (!cache) { + cache = new Map() + cachesByDb.set(db, cache) + } + const nowMs = Date.now() + const existing = cache.get(key) + if (existing && existing.expiresAtMs > nowMs) { + return await existing.value + } + const value = create() + // Never cache failures: a D1 blip must not pin an error for the TTL. + value.catch(() => { + if (cache.get(key)?.value === value) cache.delete(key) + }) + if (cache.size >= entitlementLookupCacheMaxEntries) { + const oldestKey = cache.keys().next().value + if (oldestKey !== undefined) cache.delete(oldestKey) + } + cache.set(key, { + value, + expiresAtMs: nowMs + entitlementLookupCacheTtlMs, + }) + return await value + }, + } +} + +const cachedUserPlans = createEntitlementLookupCache() +const cachedStableUserAccounts = + createEntitlementLookupCache() + +/** + * {@link getUserPlan} behind the short-TTL hot-path cache. Use only where a + * plan change may take up to a minute to apply (quota limit resolution); + * interactive plan displays should keep calling {@link getUserPlan}. + */ +export async function getCachedUserPlan( + db: D1Database, + input: { userId: string; email: string | null | undefined }, +): Promise { + const email = input.email?.trim().toLowerCase() + if (!email || !input.userId) return 'max' + if (!stableUserIdPattern.test(input.userId)) return 'max' + return await cachedUserPlans.getOrCreate( + db, + `${input.userId}\n${email}`, + async () => await getUserPlan(db, input), + ) +} + +/** + * {@link findUserAccountByStableUserId} behind the short-TTL hot-path cache, + * for per-fetch account reverse-resolution in the fetch gateway. + */ +export async function findCachedUserAccountByStableUserId( + db: D1Database, + stableUserId: string, +): Promise { + const trimmed = normalizeStableUserId(stableUserId) + if (!trimmed) return null + return await cachedStableUserAccounts.getOrCreate( + db, + trimmed, + async () => await findUserAccountByStableUserId(db, trimmed), + ) +} + /** * Reverse-resolve a stable MCP userId back to the account email, plan, and * verified-email state via one indexed `users.stable_user_id` point read. @@ -658,7 +753,10 @@ export async function consumeDailyEntitlement(input: { now?: Date }): Promise { const now = input.now ?? new Date() - const plan = await getUserPlan(input.db, { + // 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. + const plan = await getCachedUserPlan(input.db, { userId: input.userId, email: input.email, }) diff --git a/packages/worker/src/mcp/fetch-gateway.ts b/packages/worker/src/mcp/fetch-gateway.ts index d646e17315..d3347df921 100644 --- a/packages/worker/src/mcp/fetch-gateway.ts +++ b/packages/worker/src/mcp/fetch-gateway.ts @@ -24,7 +24,7 @@ import { assertPackageCanAccessResolvedSecret } from '#mcp/secrets/package-acces import { type StorageContext } from '#mcp/storage.ts' import { consumeDailyEntitlement, - findUserAccountByStableUserId, + findCachedUserAccountByStableUserId, } from '#worker/entitlements/service.ts' import { recordUsage, type UsageEnv } from '#worker/usage/record-usage.ts' @@ -94,7 +94,7 @@ export async function executeGatewayFetch(input: { const email = input.props.email ?? ( - await findUserAccountByStableUserId( + await findCachedUserAccountByStableUserId( input.env.APP_DB, input.props.userId, ) diff --git a/packages/worker/src/mcp/index.ts b/packages/worker/src/mcp/index.ts index 5d7a231d3b..9ff329886c 100644 --- a/packages/worker/src/mcp/index.ts +++ b/packages/worker/src/mcp/index.ts @@ -82,6 +82,9 @@ class MCPBase extends McpAgent { getLoopbackExports() { return this.ctx.exports as typeof workerExports } + waitUntil(promise: Promise) { + this.ctx.waitUntil(promise) + } requireDomain() { const { baseUrl } = this.getCallerContext() invariant( diff --git a/packages/worker/src/mcp/mcp-registration-agent.ts b/packages/worker/src/mcp/mcp-registration-agent.ts index 7ecb80869e..d47fa94bbd 100644 --- a/packages/worker/src/mcp/mcp-registration-agent.ts +++ b/packages/worker/src/mcp/mcp-registration-agent.ts @@ -7,4 +7,10 @@ export type McpRegistrationAgent = { getCallerContext(): McpCallerContext requireDomain(): string getLoopbackExports(): Cloudflare.Exports + /** + * Durable Object `ctx.waitUntil`, when the agent runs inside one. Tool + * handlers hand observability writes (run records, usage) to it so those + * writes do not serialize the response they observe. + */ + waitUntil?(promise: Promise): void } diff --git a/packages/worker/src/mcp/tools/execute.ts b/packages/worker/src/mcp/tools/execute.ts index 35dda336c2..8818869f87 100644 --- a/packages/worker/src/mcp/tools/execute.ts +++ b/packages/worker/src/mcp/tools/execute.ts @@ -156,6 +156,10 @@ export async function registerExecuteTool(agent: McpRegistrationAgent) { }) => { const timingStart = startToolTiming() const env = agent.getEnv() + // Hand terminal run-record writes and nested-invocation + // observability to the Durable Object's waitUntil so they stop + // serializing the execute response they observe. + const waitUntil = agent.waitUntil?.bind(agent) const baseCallerContext = agent.getCallerContext() const resolvedStorageId = storageId?.trim() || null const callerContext = { @@ -328,6 +332,7 @@ export async function registerExecuteTool(agent: McpRegistrationAgent) { baseUrl: callerContext.baseUrl, callerContext, conversationId: resolvedConversationId, + waitUntil, }) : undefined try { @@ -352,6 +357,7 @@ export async function registerExecuteTool(agent: McpRegistrationAgent) { preExecTypecheck: featureFlags['execute-pre-exec-typecheck'] === true, runRecordHandle: claimedRunHandle, + waitUntil, runRecord: { surface: 'execute', name: null, diff --git a/packages/worker/src/package-invocations/http-invoke.ts b/packages/worker/src/package-invocations/http-invoke.ts index 4d463a9127..4280243dd2 100644 --- a/packages/worker/src/package-invocations/http-invoke.ts +++ b/packages/worker/src/package-invocations/http-invoke.ts @@ -14,6 +14,7 @@ import { type PackageRuntimeToolFactories, } from './common.ts' import { invokeSavedPackageModule } from './idempotent-module-invocation.ts' +import { type PackageInvokeCheckPreloads } from './invoke-check.ts' import { resolveSavedPackage } from './module-artifacts.ts' import { buildJsonErrorResponse } from './responses.ts' @@ -76,6 +77,8 @@ export async function invokePackageExportForExecuteRuntime(input: { conversationId?: string | null toolFactories: PackageRuntimeToolFactories waitUntil?: (promise: Promise) => void + /** Check-phase loads from `packages.invokeChecked`; see invoke-check.ts. */ + preloads?: PackageInvokeCheckPreloads | null }): Promise { const packageIdOrKodyId = input.request.packageIdOrKodyId.trim() if (!packageIdOrKodyId) { @@ -94,11 +97,13 @@ export async function invokePackageExportForExecuteRuntime(input: { message: 'Package invocations require a non-empty idempotencyKey.', }) } - const savedPackage = await resolveSavedPackage({ - db: input.env.APP_DB, - userId: input.caller.userId, - packageIdOrKodyId, - }) + const savedPackage = + input.preloads?.savedPackage ?? + (await resolveSavedPackage({ + db: input.env.APP_DB, + userId: input.caller.userId, + packageIdOrKodyId, + })) if (!savedPackage) { return buildJsonErrorResponse({ status: 404, @@ -140,6 +145,7 @@ export async function invokePackageExportForExecuteRuntime(input: { runtimeInvokeDepth: input.runtimeInvokeDepth ?? 0, toolFactories: input.toolFactories, waitUntil: input.waitUntil, + preloadedModuleArtifact: input.preloads?.moduleArtifact ?? null, }) } @@ -157,6 +163,8 @@ export async function invokePackageExportForPackageRuntime(input: { runtimeInvokeDepth?: number toolFactories: PackageRuntimeToolFactories waitUntil?: (promise: Promise) => void + /** Check-phase loads from `packages.invokeChecked`; see invoke-check.ts. */ + preloads?: PackageInvokeCheckPreloads | null }): Promise { const packageIdOrKodyId = input.request.packageIdOrKodyId.trim() if (!packageIdOrKodyId) { @@ -175,11 +183,13 @@ export async function invokePackageExportForPackageRuntime(input: { message: 'Package invocations require a non-empty idempotencyKey.', }) } - const savedPackage = await resolveSavedPackage({ - db: input.env.APP_DB, - userId: input.caller.userId, - packageIdOrKodyId, - }) + const savedPackage = + input.preloads?.savedPackage ?? + (await resolveSavedPackage({ + db: input.env.APP_DB, + userId: input.caller.userId, + packageIdOrKodyId, + })) if (!savedPackage) { return buildJsonErrorResponse({ status: 404, @@ -216,6 +226,7 @@ export async function invokePackageExportForPackageRuntime(input: { runtimeInvokeDepth: input.runtimeInvokeDepth ?? 0, toolFactories: input.toolFactories, waitUntil: input.waitUntil, + preloadedModuleArtifact: input.preloads?.moduleArtifact ?? null, }) } diff --git a/packages/worker/src/package-invocations/idempotent-module-invocation.ts b/packages/worker/src/package-invocations/idempotent-module-invocation.ts index f5f0e7f204..a49ee51378 100644 --- a/packages/worker/src/package-invocations/idempotent-module-invocation.ts +++ b/packages/worker/src/package-invocations/idempotent-module-invocation.ts @@ -63,6 +63,14 @@ export async function invokeSavedPackageModule(input: { runtimeInvokeDepth?: number toolFactories: PackageRuntimeToolFactories waitUntil?: (promise: Promise) => void + /** + * Artifact already prepared by a `packages.invokeChecked` check phase + * moments earlier; skips a second manifest + artifact load. The claim + * still happens first, so replay semantics are unchanged. + */ + preloadedModuleArtifact?: Awaited< + ReturnType + > | null }) { return await withAccountWriteLease({ db: input.env.APP_DB, @@ -276,20 +284,23 @@ export async function invokeSavedPackageModule(input: { } try { - const { artifact, source: sourceRow } = await ensureModuleArtifact({ - env: input.env, - baseUrl: input.baseUrl, - savedPackage: input.savedPackage, - selector: input.moduleSelector, - userId: input.actor.userId, - }) + const { artifact, source: sourceRow } = + input.preloadedModuleArtifact ?? + (await ensureModuleArtifact({ + env: input.env, + baseUrl: input.baseUrl, + savedPackage: input.savedPackage, + selector: input.moduleSelector, + userId: input.actor.userId, + })) const repoSource = - artifact.packageContext?.sourceId != null - ? await getEntitySourceById( + artifact.packageContext?.sourceId == null || + artifact.packageContext.sourceId === sourceRow.id + ? sourceRow + : await getEntitySourceById( input.env.APP_DB, artifact.packageContext.sourceId, ) - : sourceRow const callerContext = createMcpCallerContext({ baseUrl: input.baseUrl, executionOrigin: 'background', diff --git a/packages/worker/src/package-invocations/invoke-check.ts b/packages/worker/src/package-invocations/invoke-check.ts index 95970d0cca..f4adc7b139 100644 --- a/packages/worker/src/package-invocations/invoke-check.ts +++ b/packages/worker/src/package-invocations/invoke-check.ts @@ -70,6 +70,23 @@ function buildPackageInvokeCheckWarnings(input: { return warnings } +/** + * Everything the check phase already loaded that the invoke phase would + * otherwise reload from D1/KV: the saved-package row, the current manifest, + * the resolved module target, and the prepared bundle artifact. + * `packages.invokeChecked` passes these straight into the invocation so one + * logical call resolves its package exactly once. + */ +export type PackageInvokeCheckPreloads = { + savedPackage: NonNullable>> + moduleArtifact: Awaited> +} + +export type PackageInvokeCheckOutcome = { + result: PackageInvokeCheckResult + preloads: PackageInvokeCheckPreloads | null +} + export async function checkPackageInvokeForRuntime(input: { env: Env baseUrl: string @@ -77,15 +94,39 @@ export async function checkPackageInvokeForRuntime(input: { userId: string rawInput: PackageInvokeInput }): Promise { + return ( + await checkPackageInvokeForRuntimeWithPreloads({ + ...input, + includeExportProjection: true, + }) + ).result +} + +export async function checkPackageInvokeForRuntimeWithPreloads(input: { + env: Env + baseUrl: string + operationName: 'packages.check' | 'packages.invokeChecked' + userId: string + rawInput: PackageInvokeInput + /** + * `packages.check` surfaces description / type-definition detail, which + * requires the full package source. `packages.invokeChecked` discards the + * success contract, so it skips that source load entirely. + */ + includeExportProjection: boolean +}): Promise { let request: ReturnType try { request = parsePackageInvokeInput(input.rawInput, input.operationName) } catch (error) { const message = getErrorMessage(error) - return createPackageInvokeCheckFailure({ - message, - problems: [message], - }) + return { + result: createPackageInvokeCheckFailure({ + message, + problems: [message], + }), + preloads: null, + } } const exportName = normalizeExportName(request.exportName) const invoke = buildNormalizedPackageInvokeInput({ request, exportName }) @@ -96,11 +137,14 @@ export async function checkPackageInvokeForRuntime(input: { }) if (!savedPackage) { const message = buildSavedPackageNotFoundMessage(request.packageIdOrKodyId) - return createPackageInvokeCheckFailure({ - message, - problems: [message], - contract: { exportName }, - }) + return { + result: createPackageInvokeCheckFailure({ + message, + problems: [message], + contract: { exportName }, + }), + preloads: null, + } } const packageContract = { packageId: savedPackage.id, @@ -119,11 +163,14 @@ export async function checkPackageInvokeForRuntime(input: { }) } catch (error) { const problem = `Could not load current package manifest: ${getErrorMessage(error)}` - return createPackageInvokeCheckFailure({ - message: problem, - problems: [problem], - contract: packageContract, - }) + return { + result: createPackageInvokeCheckFailure({ + message: problem, + problems: [problem], + contract: packageContract, + }), + preloads: null, + } } let resolution: ReturnType try { @@ -136,17 +183,21 @@ export async function checkPackageInvokeForRuntime(input: { }) } catch (error) { const problem = getErrorMessage(error) - return createPackageInvokeCheckFailure({ - message: problem, - problems: [problem], - contract: { - ...packageContract, - publishedCommit: manifestResult.source.published_commit ?? null, - }, - }) + return { + result: createPackageInvokeCheckFailure({ + message: problem, + problems: [problem], + contract: { + ...packageContract, + publishedCommit: manifestResult.source.published_commit ?? null, + }, + }), + preloads: null, + } } + let moduleArtifact: Awaited> try { - await ensureModuleArtifact({ + moduleArtifact = await ensureModuleArtifact({ env: input.env, baseUrl: input.baseUrl, packageManifest: manifestResult, @@ -160,15 +211,42 @@ export async function checkPackageInvokeForRuntime(input: { }) } catch (error) { const problem = `Export "${exportName}" could not be prepared for invocation: ${getErrorMessage(error)}` - return createPackageInvokeCheckFailure({ - message: problem, - problems: [problem], - contract: { - ...packageContract, - publishedCommit: manifestResult.source.published_commit ?? null, - runtimeTarget: resolution.entryPoint, + return { + result: createPackageInvokeCheckFailure({ + message: problem, + problems: [problem], + contract: { + ...packageContract, + publishedCommit: manifestResult.source.published_commit ?? null, + runtimeTarget: resolution.entryPoint, + }, + }), + preloads: null, + } + } + const preloads: PackageInvokeCheckPreloads = { + savedPackage, + moduleArtifact, + } + if (!input.includeExportProjection) { + return { + result: { + ok: true, + invoke, + contract: { + ...packageContract, + publishedCommit: manifestResult.source.published_commit ?? null, + runtimeTarget: resolution.entryPoint, + description: null, + typeDefinition: null, + warnings: buildPackageInvokeCheckWarnings({ + exportDetail: null, + sourceLoadFailed: false, + }), + }, }, - }) + preloads, + } } let files: Record | undefined let sourceLoadFailed = false @@ -197,15 +275,18 @@ export async function checkPackageInvokeForRuntime(input: { sourceLoadFailed, }) return { - ok: true, - invoke, - contract: { - ...packageContract, - publishedCommit: manifestResult.source.published_commit ?? null, - runtimeTarget: exportDetail?.runtimeTarget ?? resolution.entryPoint, - description: exportDetail?.description ?? null, - typeDefinition: exportDetail?.typeDefinition ?? null, - warnings, + result: { + ok: true, + invoke, + contract: { + ...packageContract, + publishedCommit: manifestResult.source.published_commit ?? null, + runtimeTarget: exportDetail?.runtimeTarget ?? resolution.entryPoint, + description: exportDetail?.description ?? null, + typeDefinition: exportDetail?.typeDefinition ?? null, + warnings, + }, }, + preloads, } } diff --git a/packages/worker/src/package-invocations/runtime-tool-factories.ts b/packages/worker/src/package-invocations/runtime-tool-factories.ts index c1ec1ee9b9..17080bb198 100644 --- a/packages/worker/src/package-invocations/runtime-tool-factories.ts +++ b/packages/worker/src/package-invocations/runtime-tool-factories.ts @@ -17,7 +17,11 @@ import { invokePackageExportForPackageRuntime, } from './http-invoke.ts' import { parsePackageInvokeInput } from './input-parsing.ts' -import { checkPackageInvokeForRuntime } from './invoke-check.ts' +import { + checkPackageInvokeForRuntime, + checkPackageInvokeForRuntimeWithPreloads, + type PackageInvokeCheckPreloads, +} from './invoke-check.ts' export function createPackageRuntimeInvokeToolsWithToolFactories(input: { env: Env @@ -77,7 +81,10 @@ function createPackageInvokeTools(input: { } return { user, packageContext: input.packageContext } } - const invoke = async (rawInput: PackageInvokeInput) => { + const invoke = async ( + rawInput: PackageInvokeInput, + preloads?: PackageInvokeCheckPreloads | null, + ) => { const { user, packageContext } = requireRuntimeCaller('packages.invoke') const packageInvokeDepth = input.packageInvokeDepth ?? 0 if (packageInvokeDepth >= maxPackageRuntimeInvokeDepth) { @@ -117,6 +124,7 @@ function createPackageInvokeTools(input: { runtimeInvokeDepth: packageInvokeDepth + 1, toolFactories: input.toolFactories, waitUntil: input.waitUntil, + preloads: preloads ?? null, }) : await invokePackageExportForExecuteRuntime({ env: input.env, @@ -138,6 +146,7 @@ function createPackageInvokeTools(input: { conversationId: input.conversationId ?? null, toolFactories: input.toolFactories, waitUntil: input.waitUntil, + preloads: preloads ?? null, }) if (response.status >= 200 && response.status < 400) { return response.body['result'] @@ -173,23 +182,28 @@ function createPackageInvokeTools(input: { invoke, invokeChecked: async (rawInput) => { const { user } = requireRuntimeCaller('packages.invokeChecked') - const check = await checkPackageInvokeForRuntime({ + // Skip the export projection (full package source) that a bare + // `packages.check` surfaces: invokeChecked discards the success + // contract, and the preloads let the invoke phase reuse the + // package row, manifest, and bundle artifact loaded here. + const check = await checkPackageInvokeForRuntimeWithPreloads({ env: input.env, baseUrl: input.baseUrl, operationName: 'packages.invokeChecked', userId: user.userId, rawInput, + includeExportProjection: false, }) - if (!check.ok) { + if (!check.result.ok) { const error = new Error( - `packages.invokeChecked check failed: ${check.message}`, + `packages.invokeChecked check failed: ${check.result.message}`, ) as Error & { check?: PackageInvokeCheckResult } - error.check = check + error.check = check.result throw error } - return await invoke(check.invoke) + return await invoke(check.result.invoke, check.preloads) }, } } From daf4cab4b868b1f3ee48e86971a458182aa36197 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 30 Jul 2026 02:07:09 +0000 Subject: [PATCH 2/2] Address review: fence lease reuse to the live lease and scope source 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 --- .../src/account/deletion-state.node.test.ts | 41 ++++++++++++++++++ packages/worker/src/account/deletion-state.ts | 43 +++++++++++++------ .../idempotent-module-invocation.ts | 10 ++--- 3 files changed, 75 insertions(+), 19 deletions(-) diff --git a/packages/worker/src/account/deletion-state.node.test.ts b/packages/worker/src/account/deletion-state.node.test.ts index c9f2358c1d..cff18d4638 100644 --- a/packages/worker/src/account/deletion-state.node.test.ts +++ b/packages/worker/src/account/deletion-state.node.test.ts @@ -175,6 +175,47 @@ test('nested lease for a different user still acquires its own lease', async () expect(await listActiveAccountWriteLeases(db, 'user-b')).toHaveLength(0) }) +test('detached work spawned inside write re-acquires after the outer lease releases', async () => { + const { sqlite, db } = createLeaseTestDb() + 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', + async write() { + // Simulates a waitUntil callback: created inside the lease scope + // (inheriting the AsyncLocalStorage context) but running only + // after the outer lease has been released. + detached = (async () => { + await outerReleased + await withAccountWriteLease({ + db, + stableUserId: 'user-a', + holder: 'test:detached', + async write() { + const active = await listActiveAccountWriteLeases(db, 'user-a') + expect(active).toHaveLength(1) + expect(active[0]).toEqual( + expect.objectContaining({ holder: 'test:detached' }), + ) + }, + }) + })() + }, + }) + releaseOuter() + await detached + expect(countLeaseRows(sqlite)).toBe(2) + expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0) + expect( + sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(), + ).toEqual({ active_write_count: 0 }) +}) + test('sequential sibling leases each acquire after the previous released', async () => { const { sqlite, db } = createLeaseTestDb() await withAccountWriteLease({ diff --git a/packages/worker/src/account/deletion-state.ts b/packages/worker/src/account/deletion-state.ts index 0a330d5457..31110ce6fe 100644 --- a/packages/worker/src/account/deletion-state.ts +++ b/packages/worker/src/account/deletion-state.ts @@ -82,17 +82,24 @@ export async function assertAccountWritable(env: Env, stableUserId: string) { } /** - * Stable user ids whose write lease is already held somewhere up the current + * 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 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. + * 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. + * + * Frames deactivate when the outer lease releases. Detached work (for + * example `waitUntil` callbacks spawned inside `write`) inherits this + * AsyncLocalStorage context, and without the active flag it would keep + * skipping acquisition after the lease row was already released. */ +type AccountWriteLeaseFrame = { active: boolean } const heldAccountWriteLeaseStorage = new AsyncLocalStorage< - ReadonlySet + ReadonlyMap >() export async function withAccountWriteLease(input: { @@ -102,20 +109,26 @@ export async function withAccountWriteLease(input: { write: () => Promise }) { const heldLeases = heldAccountWriteLeaseStorage.getStore() - if (heldLeases?.has(input.stableUserId)) { + if (heldLeases?.get(input.stableUserId)?.active) { return await input.write() } - const nextHeldLeases = new Set(heldLeases) - nextHeldLeases.add(input.stableUserId) - return await heldAccountWriteLeaseStorage.run(nextHeldLeases, async () => - acquireAccountWriteLeaseAndWrite(input), - ) + const frame: AccountWriteLeaseFrame = { active: true } + const nextHeldLeases = new Map(heldLeases) + nextHeldLeases.set(input.stableUserId, frame) + try { + return await heldAccountWriteLeaseStorage.run(nextHeldLeases, async () => + acquireAccountWriteLeaseAndWrite({ ...input, frame }), + ) + } finally { + frame.active = false + } } async function acquireAccountWriteLeaseAndWrite(input: { db: D1Database stableUserId: string holder?: string + frame: AccountWriteLeaseFrame write: () => Promise }) { if (typeof input.db.batch !== 'function') { @@ -133,6 +146,7 @@ async function acquireAccountWriteLeaseAndWrite(input: { try { return await input.write() } finally { + input.frame.active = false await input.db .prepare( `UPDATE users @@ -191,6 +205,7 @@ async function acquireAccountWriteLeaseAndWrite(input: { if (held?.held !== 1) throw new AccountWriteLeaseLostError() return result } finally { + input.frame.active = false const releasedAt = utcSqliteTimestamp() const released = await input.db .prepare( diff --git a/packages/worker/src/package-invocations/idempotent-module-invocation.ts b/packages/worker/src/package-invocations/idempotent-module-invocation.ts index a49ee51378..e780c4c19a 100644 --- a/packages/worker/src/package-invocations/idempotent-module-invocation.ts +++ b/packages/worker/src/package-invocations/idempotent-module-invocation.ts @@ -5,7 +5,7 @@ import { runBundledModuleWithRegistry } from '#mcp/run-kody-registry.ts' import { withAccountWriteLease } from '#worker/account/deletion-state.ts' import { type RunRecordContext } from '#worker/run-records/types.ts' import { type SavedPackageRecord } from '#worker/package-registry/types.ts' -import { getEntitySourceById } from '#worker/repo/entity-sources.ts' +import { getEntitySourceByIdForUser } from '#worker/repo/entity-sources.ts' import { getEmailAttachmentById, getEmailMessageById, @@ -297,10 +297,10 @@ export async function invokeSavedPackageModule(input: { artifact.packageContext?.sourceId == null || artifact.packageContext.sourceId === sourceRow.id ? sourceRow - : await getEntitySourceById( - input.env.APP_DB, - artifact.packageContext.sourceId, - ) + : await getEntitySourceByIdForUser(input.env.APP_DB, { + id: artifact.packageContext.sourceId, + userId: input.actor.userId, + }) const callerContext = createMcpCallerContext({ baseUrl: input.baseUrl, executionOrigin: 'background',