From 8576ae9dca083f33deea0d61127e779ed87c65e4 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:11:49 +0000 Subject: [PATCH 01/14] Offload raw email MIME to EMAIL_BLOBS R2 bucket Store inbound raw MIME in R2 (email-raw:v1:{userId}/{messageId}) with raw_mime_key on email_messages, falling back inline when the binding is missing. Reads go through loadRawMime, deletes remove the blob, account deletion enumerates and purges user blobs, and CI resource tooling provisions the bucket. Co-authored-by: Kent C. Dodds --- .../migrations/0054-email-blob-storage.sql | 4 + .../src/app/account-deletion.node.test.ts | 59 ++++- packages/worker/src/app/account-deletion.ts | 53 +++++ .../worker/src/app/admin-system-email-data.ts | 17 +- packages/worker/src/email/inbound.ts | 11 +- .../worker/src/email/inbound.workers.test.ts | 143 ++++++++++++ packages/worker/src/email/outbound.ts | 2 + .../worker/src/email/outbound.workers.test.ts | 4 + packages/worker/src/email/repo.ts | 205 ++++++++++++++---- packages/worker/src/email/test-schema.ts | 1 + packages/worker/src/email/types.ts | 1 + .../admin/admin-system-email-get.ts | 1 + .../email/email-attachment-get.ts | 1 + .../worker/src/package-invocations/service.ts | 1 + packages/worker/worker-configuration.d.ts | 3 +- packages/worker/wrangler.jsonc | 20 ++ tools/ci/preview-resources.ts | 32 ++- tools/ci/production-resources.ts | 38 +++- tools/ci/resource-utils.node.test.ts | 42 +++- tools/ci/resource-utils.ts | 120 ++++++++++ 20 files changed, 702 insertions(+), 56 deletions(-) create mode 100644 packages/worker/migrations/0054-email-blob-storage.sql diff --git a/packages/worker/migrations/0054-email-blob-storage.sql b/packages/worker/migrations/0054-email-blob-storage.sql new file mode 100644 index 0000000000..97d562b627 --- /dev/null +++ b/packages/worker/migrations/0054-email-blob-storage.sql @@ -0,0 +1,4 @@ +-- Move raw email MIME out of D1: a message may now store an R2 object key +-- (EMAIL_BLOBS bucket, `email-raw:v1:{userId}/{messageId}`) instead of the +-- inline raw_mime payload. Additive only; legacy rows keep raw_mime inline. +ALTER TABLE email_messages ADD COLUMN raw_mime_key TEXT; diff --git a/packages/worker/src/app/account-deletion.node.test.ts b/packages/worker/src/app/account-deletion.node.test.ts index 55d8d4d628..e4c07ced9f 100644 --- a/packages/worker/src/app/account-deletion.node.test.ts +++ b/packages/worker/src/app/account-deletion.node.test.ts @@ -190,6 +190,18 @@ function createTestDb(initial: RowMap): { .map((row) => ({ storage_id: row['storage_id'] })) return { results: results as Array, meta: { changes: 0 } } } + if ( + lower === + 'select raw_mime_key from email_messages where user_id = ? and raw_mime_key is not null' + ) { + results = (rows.email_messages ?? []) + .filter( + (row) => + row['user_id'] === userId && row['raw_mime_key'] != null, + ) + .map((row) => ({ raw_mime_key: row['raw_mime_key'] })) + return { results: results as Array, meta: { changes: 0 } } + } const kvMatch = lower.match( /^select kv_key from (\w+) where user_id = \?/, ) @@ -558,7 +570,19 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async email_inboxes: [{ id: 'in-1', user_id: userAaa }], email_inbox_addresses: [{ id: 'ia-1', user_id: userAaa }], email_threads: [{ id: 'et-1', user_id: userAaa }], - email_messages: [{ id: 'em-1', user_id: userAaa }], + email_messages: [ + { + id: 'em-1', + user_id: userAaa, + raw_mime_key: 'email-raw:v1:user-aaa/em-1', + }, + { id: 'em-2', user_id: userAaa, raw_mime_key: null }, + { + id: 'em-3', + user_id: userBbb, + raw_mime_key: 'email-raw:v1:user-bbb/em-3', + }, + ], 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 }], @@ -637,6 +661,13 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async }, } as unknown as KVNamespace + const deletedEmailBlobKeys: Array = [] + const emailBlobs = { + async delete(key: string) { + deletedEmailBlobKeys.push(key) + }, + } as unknown as R2Bucket + const clearStorageMock = 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 })) @@ -647,6 +678,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async const env = { APP_DB: db, BUNDLE_ARTIFACTS_KV: kv, + EMAIL_BLOBS: emailBlobs, CAPABILITY_VECTOR_INDEX: { deleteByIds: deleteVectorsMock, }, @@ -725,7 +757,15 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async ]) expect(rows.repo_sessions).toEqual([]) expect(rows.email_attachments).toEqual([]) - expect(rows.email_messages).toEqual([]) + // The other user's message and its R2 raw-MIME blob are untouched. + expect(rows.email_messages).toEqual([ + { + id: 'em-3', + user_id: userBbb, + raw_mime_key: 'email-raw:v1:user-bbb/em-3', + }, + ]) + expect(deletedEmailBlobKeys).toEqual(['email-raw:v1:user-aaa/em-1']) expect(rows.entitlement_daily_counters).toEqual([ { user_id: userBbb, resource: 'email_sends_per_day', day: '2026-07-05' }, ]) @@ -822,6 +862,7 @@ test('deleteUserAccount cascades user-scoped rows for the requested user', async expect(result.deletedRowCounts.community_bans).toBe(1) expect(result.updatedRowCounts.community_bans).toBe(1) expect(result.deletedKvKeys).toBe(11) + expect(result.deletedEmailBlobs).toBe(1) expect(result.deletedVectors).toBe(5) expect(result.clearedDurableObjects).toMatchObject({ storageRunners: 6, @@ -907,6 +948,13 @@ test('deleteUserAccount handles OAuth grant revocation and warning-only edge cas archived_job_artifacts: [ { id: 'aja-1', user_id: 'user-aaa', kv_key: 'archived:src-1' }, ], + email_messages: [ + { + id: 'em-1', + user_id: 'user-aaa', + raw_mime_key: 'email-raw:v1:user-aaa/em-1', + }, + ], }) const kvFailureResult = await deleteUserAccount({ env: { @@ -921,6 +969,13 @@ test('deleteUserAccount handles OAuth grant revocation and warning-only edge cas }) expect(kvFailureRows.published_bundle_artifacts).toEqual([]) expect(kvFailureRows.archived_job_artifacts).toEqual([]) + expect(kvFailureRows.email_messages).toEqual([]) expect(kvFailureResult.deletedKvKeys).toBe(0) + // The D1 rows are still removed when EMAIL_BLOBS is unavailable; the + // stranded blobs are reported as a warning instead of aborting. + expect(kvFailureResult.deletedEmailBlobs).toBe(0) + expect(kvFailureResult.warnings).toContain( + 'EMAIL_BLOBS binding was unavailable; 1 raw email MIME blob(s) referenced by the deleted user were not removed and must be cleaned up manually.', + ) expect(kvFailureResult.warnings.length).toBeGreaterThan(0) }) diff --git a/packages/worker/src/app/account-deletion.ts b/packages/worker/src/app/account-deletion.ts index 7a114f2dc2..08fe512a18 100644 --- a/packages/worker/src/app/account-deletion.ts +++ b/packages/worker/src/app/account-deletion.ts @@ -47,6 +47,7 @@ export type AccountDeletionResult = { deletedRowCounts: Record updatedRowCounts: Record deletedKvKeys: number + deletedEmailBlobs: number deletedArtifactRepos: number revokedOAuthGrants: number clearedDurableObjects: Record @@ -93,6 +94,7 @@ type UserDeletionInventory = { vectorIds: Array storageIds: Array bundleKvKeys: Array + emailRawMimeKeys: Array sourceSnapshots: Array savedPackages: Array repoSessions: Array @@ -280,6 +282,17 @@ async function listUserPackageServices(env: Env, userId: string) { })) } +async function listUserEmailRawMimeKeys(env: Env, userId: string) { + const rows = await env.APP_DB.prepare( + `SELECT raw_mime_key + FROM email_messages + WHERE user_id = ? AND raw_mime_key IS NOT NULL`, + ) + .bind(userId) + .all<{ raw_mime_key: string }>() + return uniqueStrings((rows.results ?? []).map((row) => row.raw_mime_key)) +} + async function listUserCommunityListingIds(env: Env, userId: string) { const rows = await env.APP_DB.prepare( `SELECT id FROM community_listings WHERE owner_user_id = ?`, @@ -328,6 +341,7 @@ async function collectUserDeletionInventory(input: { const [ vectorIds, storageIds, + emailRawMimeKeys, sourceSnapshots, savedPackages, repoSessions, @@ -346,6 +360,13 @@ async function collectUserDeletionInventory(input: { input.warnings.push(`Failed to enumerate storage ids: ${message}`) return [] as Array }), + listUserEmailRawMimeKeys(input.env, input.userId).catch((error) => { + const message = error instanceof Error ? error.message : String(error) + input.warnings.push( + `Failed to enumerate email raw-MIME blob keys: ${message}`, + ) + return [] as Array + }), listUserSourceSnapshots(input.env, input.userId).catch((error) => { const message = error instanceof Error ? error.message : String(error) input.warnings.push(`Failed to enumerate source snapshots: ${message}`) @@ -396,6 +417,7 @@ async function collectUserDeletionInventory(input: { vectorIds, storageIds, bundleKvKeys, + emailRawMimeKeys, sourceSnapshots, savedPackages, repoSessions, @@ -690,6 +712,24 @@ async function deleteKvKeys(input: { return deleted } +async function deleteEmailBlobs(input: { + blobs: R2Bucket + keys: ReadonlyArray + warnings: Array +}): Promise { + let deleted = 0 + for (const key of input.keys) { + try { + await input.blobs.delete(key) + deleted += 1 + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + input.warnings.push(`Email blob delete failed for ${key}: ${message}`) + } + } + return deleted +} + async function listKvKeysByPrefix(input: { kv: KVNamespace prefixes: ReadonlyArray @@ -919,6 +959,7 @@ export async function deleteUserAccount(input: { deletedRowCounts: {}, updatedRowCounts: {}, deletedKvKeys: 0, + deletedEmailBlobs: 0, deletedArtifactRepos: 0, revokedOAuthGrants: 0, clearedDurableObjects: {}, @@ -1024,6 +1065,18 @@ export async function deleteUserAccount(input: { ) } + if (input.env.EMAIL_BLOBS) { + result.deletedEmailBlobs = await deleteEmailBlobs({ + blobs: input.env.EMAIL_BLOBS, + keys: inventory.emailRawMimeKeys, + warnings, + }) + } else if (inventory.emailRawMimeKeys.length > 0) { + warnings.push( + `EMAIL_BLOBS binding was unavailable; ${inventory.emailRawMimeKeys.length} raw email MIME blob(s) referenced by the deleted user were not removed and must be cleaned up manually.`, + ) + } + const d1Cleanup = await deleteUserScopedRows({ env: input.env, mcpUserId: input.mcpUserId, diff --git a/packages/worker/src/app/admin-system-email-data.ts b/packages/worker/src/app/admin-system-email-data.ts index 1cb4662799..0781cc9f5f 100644 --- a/packages/worker/src/app/admin-system-email-data.ts +++ b/packages/worker/src/app/admin-system-email-data.ts @@ -1,4 +1,5 @@ import { type AdminSystemEmailLoaderData } from '#app/loader-data.ts' +import { loadRawMime } from '#worker/email/repo.ts' import { systemEmailLimits, systemEmailLocals, @@ -129,6 +130,7 @@ async function listAttachments(db: D1Database, messageId: string) { export async function loadAdminSystemEmailMessageById( db: D1Database, messageId: string, + blobs?: R2Bucket | null, ): Promise { const row = await db .prepare( @@ -166,7 +168,14 @@ export async function loadAdminSystemEmailMessageById( ), text_body: row['text_body'] == null ? null : String(row['text_body']), html_body: row['html_body'] == null ? null : String(row['html_body']), - raw_mime: row['raw_mime'] == null ? null : String(row['raw_mime']), + raw_mime: await loadRawMime({ + blobs, + message: { + rawMime: row['raw_mime'] == null ? null : String(row['raw_mime']), + rawMimeKey: + row['raw_mime_key'] == null ? null : String(row['raw_mime_key']), + }, + }), attachments: await listAttachments(db, messageId), } } @@ -209,7 +218,11 @@ export async function loadAdminSystemEmailData( .bind(systemEmailOwnerId, pageSize, offset) .all>(), selectedMessageId - ? loadAdminSystemEmailMessageById(env.APP_DB, selectedMessageId) + ? loadAdminSystemEmailMessageById( + env.APP_DB, + selectedMessageId, + env.EMAIL_BLOBS, + ) : Promise.resolve(null), ]) return { diff --git a/packages/worker/src/email/inbound.ts b/packages/worker/src/email/inbound.ts index 55fae52b7e..0e05944aa1 100644 --- a/packages/worker/src/email/inbound.ts +++ b/packages/worker/src/email/inbound.ts @@ -50,6 +50,7 @@ type ParsedInboundEmail = Awaited< async function parseAndStoreInboundEmail(input: { db: D1Database + blobs?: R2Bucket | null message: ForwardableEmailMessage recipient: string userId: string @@ -90,6 +91,7 @@ async function parseAndStoreInboundEmail(input: { })) const stored = await insertEmailMessageWithAttachments({ db: input.db, + blobs: input.blobs, message: { direction: 'inbound', userId: input.userId, @@ -154,6 +156,7 @@ export async function handleInboundEmail( env: Pick< Env, | 'APP_DB' + | 'EMAIL_BLOBS' | 'BUNDLE_ARTIFACTS_KV' | 'APP_BASE_URL' | 'USER_EMAIL_DOMAIN' @@ -350,6 +353,7 @@ export async function handleInboundEmail( const stored = await parseAndStoreInboundEmail({ db: env.APP_DB, + blobs: env.EMAIL_BLOBS, message, recipient, userId, @@ -395,7 +399,11 @@ async function handleSystemInboundEmail(input: { message: ForwardableEmailMessage env: Pick< Env, - 'APP_DB' | 'BUNDLE_ARTIFACTS_KV' | 'APP_BASE_URL' | 'USAGE_EVENTS' + | 'APP_DB' + | 'EMAIL_BLOBS' + | 'BUNDLE_ARTIFACTS_KV' + | 'APP_BASE_URL' + | 'USAGE_EVENTS' > recipient: string localPart: SystemEmailLocal @@ -477,6 +485,7 @@ async function handleSystemInboundEmail(input: { const stored = await parseAndStoreInboundEmail({ db: input.env.APP_DB, + blobs: input.env.EMAIL_BLOBS, message: input.message, recipient: input.recipient, userId: systemEmailOwnerId, diff --git a/packages/worker/src/email/inbound.workers.test.ts b/packages/worker/src/email/inbound.workers.test.ts index c69ac06fd1..1752e393a9 100644 --- a/packages/worker/src/email/inbound.workers.test.ts +++ b/packages/worker/src/email/inbound.workers.test.ts @@ -4,12 +4,16 @@ import { handleInboundEmail } from './inbound.ts' import { defaultEmailInboxName } from './default-inbox.ts' import { createEmailThread, + deleteEmailMessageById, + emailRawMimeKey, + getEmailMessageById, insertEmailMessageWithAttachments, getEmailAttachmentById, listEmailInboxesForUser, listEmailInboxAddressesForUser, listEmailMessages, listEmailAttachmentsForMessage, + loadRawMime, } from './repo.ts' import { createForwardableEmailMessage } from './test-fixtures.ts' import { ensureEmailTestSchema } from './test-schema.ts' @@ -627,6 +631,145 @@ test('getEmailAttachmentById reconstructs unnamed attachments from raw MIME', as ) }) +test('inbound email offloads raw MIME to R2 and readers and deletes follow the blob', async () => { + await ensureEmailTestSchema(env.APP_DB) + const username = `blob-${crypto.randomUUID().slice(0, 8)}` + const accountEmail = `blob-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(accountEmail) + const address = `${username}@${platformDomain}` + await seedVerifiedAccount({ + db: env.APP_DB, + email: accountEmail, + username, + }) + + const raw = [ + 'From: Sender ', + `To: ${address}`, + 'Subject: Blob stored mail', + 'Message-ID: ', + 'Content-Type: multipart/mixed; boundary="mail-boundary"', + '', + '--mail-boundary', + 'Content-Type: text/plain; charset="utf-8"', + '', + 'Blob body.', + '--mail-boundary', + 'Content-Type: text/plain; name="note.txt"', + 'Content-Disposition: attachment; filename="note.txt"', + '', + 'Attachment text', + '--mail-boundary--', + ].join('\r\n') + const message = createForwardableEmailMessage({ + from: 'sender@example.net', + to: address, + raw, + }) + await handleInboundEmail(message, createInboundEnv()) + expect(message.rejectedReason).toBeNull() + + const [stored] = await listEmailMessages({ + db: env.APP_DB, + userId, + limit: 1, + }) + expect(stored).toBeDefined() + if (!stored) throw new Error('Expected stored inbound message') + + // The raw MIME left D1 and lives at the per-user R2 key. + expect(stored.rawMime).toBeNull() + expect(stored.rawMimeKey).toBe(emailRawMimeKey(userId, stored.id)) + const blob = await env.EMAIL_BLOBS.get(stored.rawMimeKey!) + expect(await blob?.text()).toBe(raw) + + // The shared read helper resolves the blob for offloaded rows. + expect(await loadRawMime({ blobs: env.EMAIL_BLOBS, message: stored })).toBe( + raw, + ) + + // Attachment content is reconstructed from the R2-stored MIME. + const attachments = await listEmailAttachmentsForMessage({ + db: env.APP_DB, + messageId: stored.id, + }) + expect(attachments).toHaveLength(1) + const loaded = await getEmailAttachmentById({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + userId, + attachmentId: attachments[0]!.id, + }) + expect(loaded?.contentBase64 ? atob(loaded.contentBase64) : null).toBe( + 'Attachment text\n', + ) + + // Deleting the message also removes its blob. + await deleteEmailMessageById({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + messageId: stored.id, + }) + expect( + await getEmailMessageById({ + db: env.APP_DB, + userId, + messageId: stored.id, + }), + ).toBeNull() + expect(await env.EMAIL_BLOBS.get(stored.rawMimeKey!)).toBeNull() +}) + +test('raw MIME stays inline when the EMAIL_BLOBS binding is unavailable', async () => { + await ensureEmailTestSchema(env.APP_DB) + const username = `inline-${crypto.randomUUID().slice(0, 8)}` + const accountEmail = `inline-${crypto.randomUUID()}@example.com` + const userId = await createStableUserIdFromEmail(accountEmail) + const address = `${username}@${platformDomain}` + await seedVerifiedAccount({ + db: env.APP_DB, + email: accountEmail, + username, + }) + + const raw = [ + 'From: Sender ', + `To: ${address}`, + 'Subject: Inline fallback mail', + 'Message-ID: ', + '', + 'Inline body.', + ].join('\r\n') + const message = createForwardableEmailMessage({ + from: 'sender@example.net', + to: address, + raw, + }) + const envWithoutBlobs = { + ...createInboundEnv(), + EMAIL_BLOBS: undefined, + } as unknown as Parameters[1] + await handleInboundEmail(message, envWithoutBlobs) + expect(message.rejectedReason).toBeNull() + + const [stored] = await listEmailMessages({ + db: env.APP_DB, + userId, + limit: 1, + }) + expect(stored).toBeDefined() + if (!stored) throw new Error('Expected stored inbound message') + + // Legacy inline behavior: the fallback never throws and keeps mail. + expect(stored.rawMime).toBe(raw) + expect(stored.rawMimeKey).toBeNull() + // The read helper prefers the inline payload and needs no bucket. + expect(await loadRawMime({ message: stored })).toBe(raw) + expect(await loadRawMime({ blobs: env.EMAIL_BLOBS, message: stored })).toBe( + raw, + ) +}) + test('inbound email handler dispatches package subscriptions for stored inbound email', async () => { await ensureEmailTestSchema(env.APP_DB) const username = `subscriber-${crypto.randomUUID().slice(0, 8)}` diff --git a/packages/worker/src/email/outbound.ts b/packages/worker/src/email/outbound.ts index 1bce41dd01..532771f0b4 100644 --- a/packages/worker/src/email/outbound.ts +++ b/packages/worker/src/email/outbound.ts @@ -21,6 +21,7 @@ type SendEmailEnv = Pick< Env, | 'APP_DB' | 'EMAIL' + | 'EMAIL_BLOBS' | 'USAGE_EVENTS' | 'APP_BASE_URL' | 'USER_EMAIL_DOMAIN' @@ -333,6 +334,7 @@ export async function sendOutboundEmail( const providerHeaders = buildProviderHeaders(storedHeaders) const message = await insertEmailMessage({ db: input.env.APP_DB, + blobs: input.env.EMAIL_BLOBS, message: { direction: 'outbound', userId: input.userId, diff --git a/packages/worker/src/email/outbound.workers.test.ts b/packages/worker/src/email/outbound.workers.test.ts index 508c36c7bf..15e262e035 100644 --- a/packages/worker/src/email/outbound.workers.test.ts +++ b/packages/worker/src/email/outbound.workers.test.ts @@ -154,6 +154,10 @@ test('sendOutboundEmail sends from the platform-assigned username address to the 'X-Kody-Email-Message-Id': result.message.messageIdHeader, }, }) + // Outbound sends store parsed bodies only: nothing to offload to the + // EMAIL_BLOBS bucket even though the binding is available. + expect(stored?.rawMime).toBeNull() + expect(stored?.rawMimeKey).toBeNull() // The send auto-provisioned (and referenced) the platform sender // identity — no self-service verify step exists. expect(stored?.senderIdentityId).toBeTruthy() diff --git a/packages/worker/src/email/repo.ts b/packages/worker/src/email/repo.ts index 0dbccb6ca9..b64a1a4de2 100644 --- a/packages/worker/src/email/repo.ts +++ b/packages/worker/src/email/repo.ts @@ -15,6 +15,73 @@ function nowIso() { return new Date().toISOString() } +/** + * R2 object key for a message's raw MIME payload in the EMAIL_BLOBS + * bucket. The userId prefix is part of the per-user isolation contract: + * account deletion enumerates and deletes a user's blobs by these stored + * keys, and the key can never be forged to point at another user's mail. + */ +export function emailRawMimeKey(userId: string, messageId: string) { + return `email-raw:v1:${userId}/${messageId}` +} + +/** + * Best-effort raw-MIME offload to R2. Returns the stored object key, or + * null when the payload must stay inline in D1 (binding unavailable or + * the put failed). Never throws: falling back to the legacy inline + * raw_mime column must not lose mail. + */ +async function offloadRawMimeToBlobs(input: { + blobs: R2Bucket | null | undefined + userId: string + messageId: string + rawMime: string +}): Promise { + if (!input.blobs) { + console.debug( + 'email-raw-mime-inline-fallback', + 'EMAIL_BLOBS binding unavailable', + input.messageId, + ) + return null + } + const key = emailRawMimeKey(input.userId, input.messageId) + try { + await input.blobs.put(key, input.rawMime) + return key + } catch (error) { + console.debug( + 'email-raw-mime-inline-fallback', + 'R2 put failed', + input.messageId, + error, + ) + return null + } +} + +/** + * Resolve a message's raw MIME regardless of where it is stored: legacy + * rows keep the payload inline in raw_mime, offloaded rows store an R2 + * key in raw_mime_key. Returns null when the message has no raw MIME or + * the blob is unreachable. + */ +export async function loadRawMime(input: { + blobs?: R2Bucket | null + message: Pick +}): Promise { + if (input.message.rawMime != null) return input.message.rawMime + const key = input.message.rawMimeKey + if (!key) return null + if (!input.blobs) { + console.debug('email-raw-mime-blob-unavailable', key) + return null + } + const object = await input.blobs.get(key) + if (!object) return null + return await object.text() +} + function bytesToBase64(bytes: Uint8Array) { let binary = '' for (const byte of bytes) { @@ -132,6 +199,8 @@ function mapMessageRow(row: Record): EmailMessageRecord { textBody: row['text_body'] == null ? null : String(row['text_body']), htmlBody: row['html_body'] == null ? null : String(row['html_body']), rawMime: row['raw_mime'] == null ? null : String(row['raw_mime']), + rawMimeKey: + row['raw_mime_key'] == null ? null : String(row['raw_mime_key']), rawSize: row['raw_size'] == null ? null : Number(row['raw_size']), processingStatus: String(row['processing_status']) as EmailProcessingStatus, providerMessageId: @@ -509,6 +578,12 @@ export async function touchEmailThread(input: { export async function insertEmailMessage(input: { db: D1Database + /** + * EMAIL_BLOBS bucket. When present, raw MIME is offloaded to R2 and + * only raw_mime_key is persisted; when absent, raw MIME stays inline + * in the legacy raw_mime column. + */ + blobs?: R2Bucket | null message: { id?: string direction: EmailDirection @@ -540,8 +615,20 @@ export async function insertEmailMessage(input: { } }) { const timestamp = nowIso() + const messageId = input.message.id ?? crypto.randomUUID() + let rawMime = input.message.rawMime ?? null + let rawMimeKey: string | null = null + if (rawMime != null) { + rawMimeKey = await offloadRawMimeToBlobs({ + blobs: input.blobs, + userId: input.message.userId, + messageId, + rawMime, + }) + if (rawMimeKey != null) rawMime = null + } const row = { - id: input.message.id ?? crypto.randomUUID(), + id: messageId, direction: input.message.direction, user_id: input.message.userId, inbox_id: input.message.inboxId ?? null, @@ -565,7 +652,8 @@ export async function insertEmailMessage(input: { auth_results: input.message.authResults ?? null, text_body: input.message.textBody ?? null, html_body: input.message.htmlBody ?? null, - raw_mime: input.message.rawMime ?? null, + raw_mime: rawMime, + raw_mime_key: rawMimeKey, raw_size: input.message.rawSize ?? 0, processing_status: input.message.processingStatus, provider_message_id: input.message.providerMessageId ?? null, @@ -575,50 +663,59 @@ export async function insertEmailMessage(input: { created_at: timestamp, updated_at: timestamp, } - await input.db - .prepare( - `INSERT INTO email_messages ( + try { + await input.db + .prepare( + `INSERT INTO email_messages ( id, direction, user_id, inbox_id, thread_id, sender_identity_id, from_address, envelope_from, to_addresses_json, cc_addresses_json, bcc_addresses_json, reply_to_addresses_json, subject, message_id_header, in_reply_to_header, references_json, headers_json, auth_results, - text_body, html_body, raw_mime, raw_size, processing_status, + text_body, html_body, raw_mime, raw_mime_key, raw_size, processing_status, provider_message_id, error, received_at, sent_at, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - ) - .bind( - row.id, - row.direction, - row.user_id, - row.inbox_id, - row.thread_id, - row.sender_identity_id, - row.from_address, - row.envelope_from, - row.to_addresses_json, - row.cc_addresses_json, - row.bcc_addresses_json, - row.reply_to_addresses_json, - row.subject, - row.message_id_header, - row.in_reply_to_header, - row.references_json, - row.headers_json, - row.auth_results, - row.text_body, - row.html_body, - row.raw_mime, - row.raw_size, - row.processing_status, - row.provider_message_id, - row.error, - row.received_at, - row.sent_at, - row.created_at, - row.updated_at, - ) - .run() + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .bind( + row.id, + row.direction, + row.user_id, + row.inbox_id, + row.thread_id, + row.sender_identity_id, + row.from_address, + row.envelope_from, + row.to_addresses_json, + row.cc_addresses_json, + row.bcc_addresses_json, + row.reply_to_addresses_json, + row.subject, + row.message_id_header, + row.in_reply_to_header, + row.references_json, + row.headers_json, + row.auth_results, + row.text_body, + row.html_body, + row.raw_mime, + row.raw_mime_key, + row.raw_size, + row.processing_status, + row.provider_message_id, + row.error, + row.received_at, + row.sent_at, + row.created_at, + row.updated_at, + ) + .run() + } catch (error) { + // Best-effort orphan cleanup: the blob was written before the row. + if (rawMimeKey != null && input.blobs) { + await input.blobs.delete(rawMimeKey).catch(() => undefined) + } + throw error + } return mapMessageRow(row) } @@ -796,12 +893,30 @@ export async function searchEmailMessages(input: { export async function deleteEmailMessageById(input: { db: D1Database + /** EMAIL_BLOBS bucket; when present the raw-MIME blob is also deleted. */ + blobs?: R2Bucket | null messageId: string }) { + // Capture the blob key before the row disappears. Blob cleanup is + // best-effort and must never block the D1 delete. + let rawMimeKey: string | null = null + if (input.blobs) { + const row = await input.db + .prepare(`SELECT raw_mime_key FROM email_messages WHERE id = ?`) + .bind(input.messageId) + .first<{ raw_mime_key: string | null }>() + .catch(() => null) + rawMimeKey = row?.raw_mime_key ?? null + } await input.db .prepare(`DELETE FROM email_messages WHERE id = ?`) .bind(input.messageId) .run() + if (rawMimeKey != null && input.blobs) { + await input.blobs.delete(rawMimeKey).catch((error: unknown) => { + console.warn('email-raw-mime-blob-delete-failed', rawMimeKey, error) + }) + } } export async function insertEmailAttachments(input: { @@ -857,6 +972,7 @@ export async function insertEmailMessageWithAttachments( } catch (error) { await deleteEmailMessageById({ db: input.db, + blobs: input.blobs, messageId: message.id, }).catch(() => undefined) throw error @@ -881,6 +997,8 @@ export async function listEmailAttachmentsForMessage(input: { export async function getEmailAttachmentById(input: { db: D1Database + /** EMAIL_BLOBS bucket for messages whose raw MIME lives in R2. */ + blobs?: R2Bucket | null userId: string attachmentId: string }) { @@ -902,14 +1020,17 @@ export async function getEmailAttachmentById(input: { userId: input.userId, messageId: attachment.messageId, }) - if (!message?.rawMime) { + const rawMime = message + ? await loadRawMime({ blobs: input.blobs, message }) + : null + if (!rawMime) { return { ...attachment, content: null, contentBase64: null, } } - const parsed = await PostalMime.parse(message.rawMime, { + const parsed = await PostalMime.parse(rawMime, { attachmentEncoding: 'arraybuffer', }) const matched = parsed.attachments.find((candidate) => { diff --git a/packages/worker/src/email/test-schema.ts b/packages/worker/src/email/test-schema.ts index 172bfa3c94..4302a8d498 100644 --- a/packages/worker/src/email/test-schema.ts +++ b/packages/worker/src/email/test-schema.ts @@ -87,6 +87,7 @@ ON email_sender_identities(user_id, email);`, text_body TEXT, html_body TEXT, raw_mime TEXT, + raw_mime_key TEXT, raw_size INTEGER NOT NULL DEFAULT 0, processing_status TEXT NOT NULL CHECK (processing_status IN ('stored', 'sent', 'failed')), provider_message_id TEXT, diff --git a/packages/worker/src/email/types.ts b/packages/worker/src/email/types.ts index f2cf5e74e3..c22b122693 100644 --- a/packages/worker/src/email/types.ts +++ b/packages/worker/src/email/types.ts @@ -117,6 +117,7 @@ export type EmailMessageRecord = { textBody: string | null htmlBody: string | null rawMime: string | null + rawMimeKey: string | null rawSize: number | null processingStatus: EmailProcessingStatus providerMessageId: string | null diff --git a/packages/worker/src/mcp/capabilities/admin/admin-system-email-get.ts b/packages/worker/src/mcp/capabilities/admin/admin-system-email-get.ts index 8c02f2f8e1..bf7d8ee6f4 100644 --- a/packages/worker/src/mcp/capabilities/admin/admin-system-email-get.ts +++ b/packages/worker/src/mcp/capabilities/admin/admin-system-email-get.ts @@ -62,6 +62,7 @@ export const adminSystemEmailGetCapability = defineDomainCapability( const message = await loadAdminSystemEmailMessageById( ctx.env.APP_DB, args.id, + ctx.env.EMAIL_BLOBS, ) return { message } }, diff --git a/packages/worker/src/mcp/capabilities/email/email-attachment-get.ts b/packages/worker/src/mcp/capabilities/email/email-attachment-get.ts index 622333ef7b..a1edc6423f 100644 --- a/packages/worker/src/mcp/capabilities/email/email-attachment-get.ts +++ b/packages/worker/src/mcp/capabilities/email/email-attachment-get.ts @@ -33,6 +33,7 @@ export const emailAttachmentGetCapability = defineDomainCapability( const user = await requireVerifiedEmailAccountUser(ctx) const attachment = await getEmailAttachmentById({ db: ctx.env.APP_DB, + blobs: ctx.env.EMAIL_BLOBS, userId: user.userId, attachmentId: args.attachment_id, }) diff --git a/packages/worker/src/package-invocations/service.ts b/packages/worker/src/package-invocations/service.ts index f930b64592..92cf7f31ad 100644 --- a/packages/worker/src/package-invocations/service.ts +++ b/packages/worker/src/package-invocations/service.ts @@ -1173,6 +1173,7 @@ async function invokeSavedPackageModule(input: { getAttachment: async (attachmentId) => { const attachment = await getEmailAttachmentById({ db: input.env.APP_DB, + blobs: input.env.EMAIL_BLOBS, userId: input.actor.userId, attachmentId, }) diff --git a/packages/worker/worker-configuration.d.ts b/packages/worker/worker-configuration.d.ts index accfc17ad0..0a2abe6415 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 ./packages/worker/worker-configuration.d.ts --config packages/worker/wrangler.jsonc --env production` (hash: bf20216b0d00b726c40a80e44b0af42e) +// Generated by Wrangler by running `wrangler types ./packages/worker/worker-configuration.d.ts --config packages/worker/wrangler.jsonc --env production` (hash: f460224c45cb9da71a1b3fda65a62ffd) // Runtime types generated with workerd@1.20260415.1 2026-04-16 global_fetch_strictly_public,nodejs_compat declare namespace Cloudflare { interface GlobalProps { @@ -9,6 +9,7 @@ declare namespace Cloudflare { interface Env { OAUTH_KV: KVNamespace; BUNDLE_ARTIFACTS_KV: KVNamespace; + EMAIL_BLOBS: R2Bucket; APP_DB: D1Database; CAPABILITY_VECTOR_INDEX: VectorizeIndex; EMAIL: SendEmail; diff --git a/packages/worker/wrangler.jsonc b/packages/worker/wrangler.jsonc index 7221feedcc..7a2a393be0 100644 --- a/packages/worker/wrangler.jsonc +++ b/packages/worker/wrangler.jsonc @@ -169,6 +169,12 @@ "binding": "BUNDLE_ARTIFACTS_KV", }, ], + "r2_buckets": [ + { + "binding": "EMAIL_BLOBS", + "bucket_name": "kody-email-blobs", + }, + ], "vectorize": [ { "binding": "CAPABILITY_VECTOR_INDEX", @@ -275,6 +281,12 @@ "binding": "BUNDLE_ARTIFACTS_KV", }, ], + "r2_buckets": [ + { + "binding": "EMAIL_BLOBS", + "bucket_name": "kody-preview-email-blobs", + }, + ], "vectorize": [ { "binding": "CAPABILITY_VECTOR_INDEX", @@ -381,6 +393,14 @@ "binding": "BUNDLE_ARTIFACTS_KV", }, ], + // Like the D1 database above, the test env reuses the preview + // bucket name; local runs and vitest-pool-workers simulate it. + "r2_buckets": [ + { + "binding": "EMAIL_BLOBS", + "bucket_name": "kody-preview-email-blobs", + }, + ], "vars": { "SENTRY_ENVIRONMENT": "test", "ARTIFACTS_NAMESPACE": "default", diff --git a/tools/ci/preview-resources.ts b/tools/ci/preview-resources.ts index 72ad227bcf..910aae5206 100644 --- a/tools/ci/preview-resources.ts +++ b/tools/ci/preview-resources.ts @@ -1,4 +1,6 @@ import { + deleteR2Bucket, + ensureR2Bucket, fail, listD1Databases, listKvNamespaces, @@ -88,6 +90,7 @@ function buildPreviewResourceNames(workerName: string) { const d1Suffix = '-db' const oauthKvSuffix = '-oauth-kv' const bundleKvSuffix = '-bundle-artifacts-kv' + const emailBlobsSuffix = '-email-blobs' const d1DatabaseName = truncateWithSuffix(workerName, d1Suffix, maxLen) const oauthKvTitle = truncateWithSuffix(workerName, oauthKvSuffix, maxLen) @@ -96,11 +99,17 @@ function buildPreviewResourceNames(workerName: string) { bundleKvSuffix, maxLen, ) + const emailBlobsBucketName = truncateWithSuffix( + workerName, + emailBlobsSuffix, + maxLen, + ) return { d1DatabaseName, oauthKvTitle, bundleArtifactsKvTitle, + emailBlobsBucketName, } } @@ -234,8 +243,12 @@ function deleteKvNamespace({ } async function ensurePreviewResources(options: CliOptions) { - const { d1DatabaseName, oauthKvTitle, bundleArtifactsKvTitle } = - buildPreviewResourceNames(options.workerName) + const { + d1DatabaseName, + oauthKvTitle, + bundleArtifactsKvTitle, + emailBlobsBucketName, + } = buildPreviewResourceNames(options.workerName) const d1 = ensureD1Database({ name: d1DatabaseName, location: options.d1Location, @@ -249,6 +262,10 @@ async function ensurePreviewResources(options: CliOptions) { title: bundleArtifactsKvTitle, dryRun: options.dryRun, }) + const emailBlobs = ensureR2Bucket({ + name: emailBlobsBucketName, + dryRun: options.dryRun, + }) const generatedConfigPath = await writeGeneratedWranglerConfig({ baseConfigPath: options.wranglerConfigPath, @@ -259,6 +276,7 @@ async function ensurePreviewResources(options: CliOptions) { d1DatabaseId: d1.id, oauthKvId: oauthKv.id, bundleArtifactsKvId: bundleArtifactsKv.id, + emailBlobsBucketName: emailBlobs.name, workerVars: { CLOUDFLARE_ACCOUNT_ID: process.env.CLOUDFLARE_ACCOUNT_ID, }, @@ -272,11 +290,17 @@ async function ensurePreviewResources(options: CliOptions) { console.log(`oauth_kv_id=${oauthKv.id}`) console.log(`bundle_artifacts_kv_title=${bundleArtifactsKv.title}`) console.log(`bundle_artifacts_kv_id=${bundleArtifactsKv.id}`) + console.log(`email_blobs_bucket_name=${emailBlobs.name}`) } async function cleanupPreviewResources(options: CliOptions) { - const { d1DatabaseName, oauthKvTitle, bundleArtifactsKvTitle } = - buildPreviewResourceNames(options.workerName) + const { + d1DatabaseName, + oauthKvTitle, + bundleArtifactsKvTitle, + emailBlobsBucketName, + } = buildPreviewResourceNames(options.workerName) + deleteR2Bucket({ name: emailBlobsBucketName, dryRun: options.dryRun }) deleteKvNamespace({ title: bundleArtifactsKvTitle, dryRun: options.dryRun }) deleteKvNamespace({ title: oauthKvTitle, dryRun: options.dryRun }) deleteD1Database({ name: d1DatabaseName, dryRun: options.dryRun }) diff --git a/tools/ci/production-resources.ts b/tools/ci/production-resources.ts index 0ee686150c..73e2b39861 100644 --- a/tools/ci/production-resources.ts +++ b/tools/ci/production-resources.ts @@ -1,5 +1,6 @@ import { readFile } from 'node:fs/promises' import { + ensureR2Bucket, fail, listD1Databases, listKvNamespaces, @@ -27,6 +28,7 @@ type ResolvedProductionBindings = { oauthKvConfiguredId: string bundleArtifactsKvTitle: string bundleArtifactsKvConfiguredId: string + emailBlobsBucketName: string } function parseArgs(argv: Array): { @@ -305,6 +307,33 @@ async function resolveProductionBindings({ bundleArtifactsKvTitleFromConfig || defaultBundleArtifactsKvTitle(workerName) + const r2Buckets = (productionEnv as Record).r2_buckets + if (!Array.isArray(r2Buckets)) { + fail( + `wrangler config "${wranglerConfigPath}" is missing "env.production.r2_buckets".`, + ) + } + + const emailBlobsEntry = r2Buckets.find((entry) => { + if (!entry || typeof entry !== 'object') return false + return (entry as Record).binding === 'EMAIL_BLOBS' + }) as Record | undefined + if (!emailBlobsEntry) { + fail( + `wrangler config "${wranglerConfigPath}" has no production R2 binding for "EMAIL_BLOBS".`, + ) + } + + const emailBlobsBucketName = emailBlobsEntry.bucket_name + if ( + typeof emailBlobsBucketName !== 'string' || + emailBlobsBucketName.length === 0 + ) { + fail( + `wrangler config "${wranglerConfigPath}" is missing "bucket_name" for production "EMAIL_BLOBS".`, + ) + } + const resolved: ResolvedProductionBindings = { workerName, d1DatabaseName, @@ -313,6 +342,7 @@ async function resolveProductionBindings({ oauthKvConfiguredId, bundleArtifactsKvTitle, bundleArtifactsKvConfiguredId, + emailBlobsBucketName, } return resolved @@ -324,7 +354,7 @@ async function ensureProductionResources(options: CliOptions) { kvTitleOverride: options.kvTitleOverride, }) console.error( - `Ensuring production resources for worker: ${bindings.workerName} (D1: ${bindings.d1DatabaseName}, OAuth KV: ${bindings.oauthKvTitle}, Bundle KV: ${bindings.bundleArtifactsKvTitle})`, + `Ensuring production resources for worker: ${bindings.workerName} (D1: ${bindings.d1DatabaseName}, OAuth KV: ${bindings.oauthKvTitle}, Bundle KV: ${bindings.bundleArtifactsKvTitle}, Email R2: ${bindings.emailBlobsBucketName})`, ) const d1 = ensureD1Database({ @@ -343,6 +373,10 @@ async function ensureProductionResources(options: CliOptions) { configuredId: bindings.bundleArtifactsKvConfiguredId, dryRun: options.dryRun, }) + const emailBlobs = ensureR2Bucket({ + name: bindings.emailBlobsBucketName, + dryRun: options.dryRun, + }) const generatedConfigPath = await writeGeneratedWranglerConfig({ baseConfigPath: options.wranglerConfigPath, @@ -353,6 +387,7 @@ async function ensureProductionResources(options: CliOptions) { d1DatabaseId: d1.id, oauthKvId: oauthKv.id, bundleArtifactsKvId: bundleArtifactsKv.id, + emailBlobsBucketName: emailBlobs.name, workerVars: { APP_BASE_URL: process.env.APP_BASE_URL, CLOUDFLARE_ACCOUNT_ID: process.env.CLOUDFLARE_ACCOUNT_ID, @@ -367,6 +402,7 @@ async function ensureProductionResources(options: CliOptions) { console.log(`oauth_kv_id=${oauthKv.id}`) console.log(`bundle_artifacts_kv_title=${bundleArtifactsKv.title}`) console.log(`bundle_artifacts_kv_id=${bundleArtifactsKv.id}`) + console.log(`email_blobs_bucket_name=${emailBlobs.name}`) } async function main() { diff --git a/tools/ci/resource-utils.node.test.ts b/tools/ci/resource-utils.node.test.ts index 9765bf2ec7..735fcf0131 100644 --- a/tools/ci/resource-utils.node.test.ts +++ b/tools/ci/resource-utils.node.test.ts @@ -4,7 +4,11 @@ import path from 'node:path' import { fileURLToPath } from 'node:url' import { expect, test, vi } from 'vitest' -import { parseJsonc, writeGeneratedWranglerConfig } from './resource-utils.ts' +import { + parseJsonc, + parseR2BucketListOutput, + writeGeneratedWranglerConfig, +} from './resource-utils.ts' const thisDir = path.dirname(fileURLToPath(import.meta.url)) const workerWranglerConfigPath = path.resolve( @@ -28,6 +32,7 @@ test('writeGeneratedWranglerConfig preserves migrations and copies environment a d1DatabaseId: 'dry-run-kody', oauthKvId: 'dry-run-kody-oauth', bundleArtifactsKvId: 'dry-run-kody-bundle-artifacts', + emailBlobsBucketName: 'kody-email-blobs', }) const productionConfig = parseJsonc<{ @@ -38,7 +43,10 @@ test('writeGeneratedWranglerConfig preserves migrations and copies environment a }> assets?: { run_worker_first?: Array } env?: { - production?: { assets?: { run_worker_first?: Array } } + production?: { + assets?: { run_worker_first?: Array } + r2_buckets?: Array<{ binding: string; bucket_name: string }> + } } }>(await readFile(productionOutPath, 'utf8')) const migrationTags = productionConfig.migrations.map( @@ -61,6 +69,9 @@ test('writeGeneratedWranglerConfig preserves migrations and copies environment a productionConfig.env?.production?.assets, ) expect(productionConfig.assets?.run_worker_first?.length).toBeGreaterThan(0) + expect(productionConfig.env?.production?.r2_buckets).toEqual([ + { binding: 'EMAIL_BLOBS', bucket_name: 'kody-email-blobs' }, + ]) const previewOutPath = path.join(tempDir, 'wrangler-preview.generated.json') await writeGeneratedWranglerConfig({ @@ -72,21 +83,45 @@ test('writeGeneratedWranglerConfig preserves migrations and copies environment a d1DatabaseId: 'dry-run-kody-pr-123-db', oauthKvId: 'dry-run-kody-pr-123-oauth', bundleArtifactsKvId: 'dry-run-kody-pr-123-bundle-artifacts', + emailBlobsBucketName: 'kody-pr-123-email-blobs', }) const previewConfig = parseJsonc<{ assets?: { run_worker_first?: Array } env?: { - preview?: { assets?: { run_worker_first?: Array } } + preview?: { + assets?: { run_worker_first?: Array } + r2_buckets?: Array<{ binding: string; bucket_name: string }> + } } }>(await readFile(previewOutPath, 'utf8')) expect(previewConfig.assets).toEqual(previewConfig.env?.preview?.assets) expect(previewConfig.assets?.run_worker_first?.length).toBeGreaterThan(0) + // The preview R2 bucket name is overridden per preview deploy. + expect(previewConfig.env?.preview?.r2_buckets).toEqual([ + { binding: 'EMAIL_BLOBS', bucket_name: 'kody-pr-123-email-blobs' }, + ]) } finally { await rm(tempDir, { force: true, recursive: true }) } }) +test('parseR2BucketListOutput reads bucket names from labelled wrangler output', () => { + const output = [ + 'Listing buckets...', + 'name: kody-email-blobs', + 'creation_date: 2026-07-01T00:00:00.000Z', + '', + 'name: kody-pr-42-email-blobs', + 'creation_date: 2026-07-02T00:00:00.000Z', + ].join('\n') + expect(parseR2BucketListOutput(output)).toEqual([ + 'kody-email-blobs', + 'kody-pr-42-email-blobs', + ]) + expect(parseR2BucketListOutput('')).toEqual([]) +}) + test('writeGeneratedWranglerConfig rejects invalid environment asset config', async () => { const tempDir = await mkdtemp(path.join(os.tmpdir(), 'kody-resource-utils-')) const exit = vi.spyOn(process, 'exit').mockImplementation((() => { @@ -126,6 +161,7 @@ test('writeGeneratedWranglerConfig rejects invalid environment asset config', as d1DatabaseId: 'dry-run-kody', oauthKvId: 'dry-run-kody-oauth', bundleArtifactsKvId: 'dry-run-kody-bundle-artifacts', + emailBlobsBucketName: 'kody-email-blobs', }), ).rejects.toThrow('process.exit') expect(error).toHaveBeenCalledWith( diff --git a/tools/ci/resource-utils.ts b/tools/ci/resource-utils.ts index 31c79fa87d..08c88f09b7 100644 --- a/tools/ci/resource-utils.ts +++ b/tools/ci/resource-utils.ts @@ -108,6 +108,98 @@ export function listKvNamespaces(): Array { } } +const ansiEscape = String.fromCharCode(27) + +/** + * `wrangler r2 bucket list` has no JSON output; it prints labelled-value + * blocks (`name: ` / `creation_date: ...`). Parse the bucket + * names out of that listing. Piped wrangler output is normally + * color-free, but ANSI sequences are stripped defensively. + */ +export function parseR2BucketListOutput(output: string): Array { + const names: Array = [] + for (const line of output.split('\n')) { + const plain = line + .split(ansiEscape) + .map((part, index) => + index === 0 ? part : part.replace(/^\[[0-9;]*m/, ''), + ) + .join('') + const match = /^name:\s*(\S+)\s*$/.exec(plain.trim()) + if (match?.[1]) names.push(match[1]) + } + return names +} + +export function listR2BucketNames(): Array { + const result = runWrangler(['r2', 'bucket', 'list'], { quiet: true }) + if (result.status !== 0) { + fail('Failed to list R2 buckets (wrangler r2 bucket list).') + } + return parseR2BucketListOutput(result.stdout) +} + +export function ensureR2Bucket({ + name, + dryRun, +}: { + name: string + dryRun: boolean +}) { + if (dryRun) { + console.error(`[dry-run] ensure R2 bucket: ${name}`) + return { name } + } + + if (listR2BucketNames().includes(name)) { + console.error(`R2 bucket exists: ${name}`) + return { name } + } + + const createResult = runWrangler(['r2', 'bucket', 'create', name], { + quiet: true, + }) + if (createResult.status !== 0) { + fail(`Failed to create R2 bucket: ${name}`) + } + + if (!listR2BucketNames().includes(name)) { + fail(`Created R2 bucket "${name}" but could not find it via list.`) + } + console.error(`Created R2 bucket: ${name}`) + return { name } +} + +export function deleteR2Bucket({ + name, + dryRun, +}: { + name: string + dryRun: boolean +}) { + if (dryRun) { + console.error(`[dry-run] delete R2 bucket: ${name}`) + return + } + + if (!listR2BucketNames().includes(name)) { + console.error(`R2 bucket already deleted: ${name}`) + return + } + + const result = runWrangler(['r2', 'bucket', 'delete', name], { quiet: true }) + if (result.status !== 0) { + // R2 refuses to delete non-empty buckets and wrangler has no purge + // command; a leftover preview bucket is preferable to a failed + // cleanup run, so warn instead of failing the workflow. + console.error( + `Warning: failed to delete R2 bucket ${name} (it may not be empty); it must be cleaned up manually.`, + ) + return + } + console.error(`Deleted R2 bucket: ${name}`) +} + function stripJsonc(source: string) { let output = '' let inString = false @@ -283,6 +375,7 @@ export async function writeGeneratedWranglerConfig({ d1DatabaseId, oauthKvId, bundleArtifactsKvId, + emailBlobsBucketName, workerVars, extraMigrations, }: { @@ -294,6 +387,7 @@ export async function writeGeneratedWranglerConfig({ d1DatabaseId: string oauthKvId: string bundleArtifactsKvId: string + emailBlobsBucketName: string workerVars?: Record extraMigrations?: Array }) { @@ -396,6 +490,32 @@ export async function writeGeneratedWranglerConfig({ preview_id: bundleArtifactsKvId, } + const r2Buckets = (targetEnv as Record).r2_buckets + if (!Array.isArray(r2Buckets)) { + fail( + `wrangler config "${baseConfigPath}" is missing "env.${envName}.r2_buckets".`, + ) + } + + const emailBlobsEntryIndex = r2Buckets.findIndex((entry) => { + if (!entry || typeof entry !== 'object') return false + return (entry as Record).binding === 'EMAIL_BLOBS' + }) + if (emailBlobsEntryIndex < 0) { + fail( + `wrangler config "${baseConfigPath}" has no ${envName} R2 binding for "EMAIL_BLOBS".`, + ) + } + + const emailBlobsEntry = r2Buckets[emailBlobsEntryIndex] as Record< + string, + unknown + > + r2Buckets[emailBlobsEntryIndex] = { + ...emailBlobsEntry, + bucket_name: emailBlobsBucketName, + } + const existingVars = (targetEnv as Record).vars if ( existingVars !== undefined && From 7f57acbc8099a3494d6ce4ae369f73247b35e458 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:11:59 +0000 Subject: [PATCH 02/14] Derive usage rollups from Analytics Engine; heal stable user ids recordUsage writes only the AE data point when USAGE_EVENTS is bound; usage_rollups become an hourly derived aggregate (aggregateUsageRollups over the AE SQL API) with a cachified KV read cache on the admin usage pages. The stable-user-id scan fallback now writes back healed ids and a maintenance endpoint backfills legacy rows. Co-authored-by: Kent C. Dodds --- .../contributing/architecture/entitlements.md | 32 ++- .../architecture/usage-metering.md | 55 +++- package-lock.json | 7 + packages/worker/package.json | 1 + .../src/app/admin-usage-data.node.test.ts | 74 +++++ packages/worker/src/app/admin-usage-data.ts | 68 ++++- packages/worker/src/entitlements/service.ts | 13 +- packages/worker/src/kv-cachified.ts | 53 ++++ .../src/maintenance-handler.node.test.ts | 145 +++++++++- packages/worker/src/maintenance-handler.ts | 72 +++++ packages/worker/src/mcp/executor.node.test.ts | 17 +- .../src/usage/aggregate-rollups.node.test.ts | 264 ++++++++++++++++++ .../worker/src/usage/aggregate-rollups.ts | 202 ++++++++++++++ packages/worker/src/usage/record-usage.ts | 38 ++- .../src/usage/record-usage.workers.test.ts | 80 ++++-- packages/worker/src/user-id.ts | 51 +++- packages/worker/src/user-id.workers.test.ts | 139 +++++++++ 17 files changed, 1223 insertions(+), 88 deletions(-) create mode 100644 packages/worker/src/kv-cachified.ts create mode 100644 packages/worker/src/usage/aggregate-rollups.node.test.ts create mode 100644 packages/worker/src/usage/aggregate-rollups.ts create mode 100644 packages/worker/src/user-id.workers.test.ts diff --git a/docs/contributing/architecture/entitlements.md b/docs/contributing/architecture/entitlements.md index f061456380..478127c296 100644 --- a/docs/contributing/architecture/entitlements.md +++ b/docs/contributing/architecture/entitlements.md @@ -56,19 +56,25 @@ documented fail-open gaps, acceptable because they are only reachable after a user action that is itself gated. One path cannot fail open: inbound email routing has no caller context but must -enforce receive quotas. `findUserAccountByStableUserId` in `service.ts` -reverse-resolves the stable user id to the account email (and plan) by scanning -the users table and hashing each email — the same pattern -`isAccountEmailVerified` uses. Positive matches are cached per isolate (the -mapping is a content hash, so hits can never go stale) and re-verified with one -point read. Only use it on contextless paths; interactive surfaces already carry -the email. - -The cold-path scan is O(users) per isolate on an attacker-reachable path, which -is acceptable for the current single-digit user count only. **A persisted -`users.stable_user_id` column (indexed, with an app-level backfill — SQLite -cannot compute SHA-256 in a migration) is required before onboarding external -users / design partners**; the scan must not ship into multi-tenant use. +enforce receive quotas. It resolves the owning account via the indexed username +lookup (`findPublicUserIdentityByUsername`) — it does not reverse-resolve stable +user ids. `findUserAccountByStableUserId` in `service.ts` remains the +reverse-resolution helper for other contextless paths (package-runtime contexts +that act with only the hashed userId, mirroring `findUserAccount` in +`email/platform-address.ts`): stored `stable_user_id` values are one indexed +point read (unique partial index from migration 0052, persisted at signup), +positive matches are cached per isolate (the mapping is a content hash, so hits +can never go stale) and re-verified with one point read. Only use it on +contextless paths; interactive surfaces already carry the email. + +Legacy rows with a NULL `stable_user_id` still fall back to a scan that hashes +each email, but the scan is now self-healing: a match writes the computed id +back (`UPDATE ... WHERE stable_user_id IS NULL`), so each legacy row pays the +scan at most once. The authenticated +`POST /__maintenance/backfill-stable-user-ids` endpoint (`backfillStableUserIds` +in `packages/worker/src/maintenance-handler.ts`) backfills all remaining legacy +rows in keyset-paged batches, eliminating the scan entirely for existing +deployments. ## The error shape diff --git a/docs/contributing/architecture/usage-metering.md b/docs/contributing/architecture/usage-metering.md index 3abd690c57..172141d60e 100644 --- a/docs/contributing/architecture/usage-metering.md +++ b/docs/contributing/architecture/usage-metering.md @@ -70,12 +70,14 @@ each chokepoint records exactly one event per metered unit, so sums are safe. ## Sinks -`recordUsage()` writes each event to two sinks, and each sink degrades -independently: - -1. **Workers Analytics Engine** (`USAGE_EVENTS` dataset binding, configured for - the `production` and `preview` environments in - `packages/worker/wrangler.jsonc`). Data point layout: +`recordUsage()` picks its sink by environment: + +1. **Workers Analytics Engine is the write path in production and preview** + (`USAGE_EVENTS` dataset binding, configured for the `production` and + `preview` environments in `packages/worker/wrangler.jsonc`). When the binding + is present, each event is one non-blocking `writeDataPoint` call and nothing + else — a per-event D1 upsert would serialize every metered request (execute, + fetch, email, jobs, ...) on D1's single writer. Data point layout: - `indexes`: `[userId]` - `blobs`: `[userId, eventType, entityId ?? '', outcome, timestamp]` - `doubles`: `[durationMs ?? 0, cpuMs ?? 0, bytes ?? 0]` @@ -83,14 +85,30 @@ independently: Analytics Engine is the analysis store (sampling-tolerant, high cardinality). Do not build enforcement on it. -2. **D1 `usage_rollups`** (migration +2. **D1 `usage_rollups` is a derived aggregate** (migration `packages/worker/migrations/0047-usage-rollups.sql`): per-user, per-metric, - per-month counters, upserted on every event: + per-month counters: - key: `(user_id, metric, month)` where `metric` is the `eventType` and `month` is the UTC `YYYY-MM` prefix of the event timestamp - counters: `event_count`, `error_count`, `total_duration_ms`, `total_cpu_ms`, `total_bytes` + In production/preview the rows are recomputed hourly by + `aggregateUsageRollups` in `packages/worker/src/usage/aggregate-rollups.ts` + (gated by `shouldRunUsageAggregationCron` in the scheduled handler): it + queries the Analytics Engine SQL API for the current UTC month grouped by + user and metric (weighting by `_sample_interval`, since Analytics Engine + samples under load) and batch-upserts absolute values — an idempotent + recompute, not increments. Analytics Engine retention (~90 days) always + covers a full month, so month-to-date recompute is complete; prior months + already in D1 stay untouched. The aggregation needs `CLOUDFLARE_ACCOUNT_ID` + and `CLOUDFLARE_API_TOKEN` and no-ops with a debug log when either (or the + `USAGE_EVENTS` binding) is missing. + + **Local-dev direct fallback:** when `USAGE_EVENTS` is absent (local dev, + tests), `recordUsage` upserts `usage_rollups` directly per event, so local + admin pages and workers-unit tests work without Analytics Engine access. + The rollup is the cheap read path a future quota/entitlements layer will consume: one point lookup per user, metric, and month. @@ -116,15 +134,15 @@ Guarantees and rules: - It accepts any object with optional `USAGE_EVENTS` / `APP_DB` bindings (`UsageEnv`), so the full `Env` can be passed directly. - **Graceful degradation:** in local dev and tests where the Analytics Engine - binding is not present, the event is logged via `console.debug` instead; when - `APP_DB` is missing (or the table does not exist), the rollup write is skipped - with a debug log. + binding is not present, the event is logged via `console.debug` and upserted + into `usage_rollups` directly; when `APP_DB` is missing (or the table does not + exist), the rollup write is skipped with a debug log. - If `userId` is empty, the event is skipped entirely. Callers on paths that can run without a user (for example anonymous gateway fetches) must guard with `if (userId)` and not invent placeholder ids. -- The returned promise resolves quickly (one `writeDataPoint`, one D1 upsert). - `await` it inline, or pass it to `ctx.waitUntil(...)` inside Durable Objects - when the caller must not block. +- The returned promise resolves quickly (one `writeDataPoint`, or one D1 upsert + in local dev). `await` it inline, or pass it to `ctx.waitUntil(...)` inside + Durable Objects when the caller must not block. ## Recipe: instrumenting a new chokepoint @@ -188,6 +206,13 @@ Guarantees and rules: ## Reading the data - Analytics Engine: query the `kody_usage_events` dataset (SQL API) filtered by - the `index1` user id; blob/double positions are listed above. + the `index1` user id; blob/double positions are listed above. Remember that + Analytics Engine samples: count with `sum(_sample_interval)` and sum values + with `sum(doubleN * _sample_interval)`. - D1: `SELECT * FROM usage_rollups WHERE user_id = ?1 AND month = ?2` gives every metric for a user's month in one small scan. +- Admin usage page: `packages/worker/src/app/admin-usage-data.ts` caches its + per-page rollup read model for ~5 minutes in `BUNDLE_ARTIFACTS_KV` via the + `@epic-web/cachified` adapter in `packages/worker/src/kv-cachified.ts` (key + prefix `derived-cache:v1:`), keyed by user ids + month, falling through to + direct D1 queries when KV is unavailable. diff --git a/package-lock.json b/package-lock.json index 00f5383a30..645b0a1c74 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1178,6 +1178,12 @@ "tslib": "^2.4.0" } }, + "node_modules/@epic-web/cachified": { + "version": "5.6.3", + "resolved": "https://registry.npmjs.org/@epic-web/cachified/-/cachified-5.6.3.tgz", + "integrity": "sha512-SHVZOm3h9/iXN3aoONTIouoOVQpBSwgJFo/NPGvTqTI1YHRDer6g9kP4iEhZ7UfWDvNiUCqOvC5f8cQ88yhf5w==", + "license": "MIT" + }, "node_modules/@epic-web/config": { "version": "1.24.1", "dev": true, @@ -15286,6 +15292,7 @@ "@cloudflare/shell": "^0.3.2", "@cloudflare/worker-bundler": "^0.1.3", "@cloudflare/workers-oauth-provider": "^0.4.0", + "@epic-web/cachified": "^5.6.3", "@epic-web/invariant": "^1.0.0", "@epic-web/totp": "^4.0.1", "@kody-bot/connector-kit": "^2.0.0", diff --git a/packages/worker/package.json b/packages/worker/package.json index 3d6b893b94..17b0ec2a2d 100644 --- a/packages/worker/package.json +++ b/packages/worker/package.json @@ -11,6 +11,7 @@ "@cloudflare/shell": "^0.3.2", "@cloudflare/worker-bundler": "^0.1.3", "@cloudflare/workers-oauth-provider": "^0.4.0", + "@epic-web/cachified": "^5.6.3", "@epic-web/invariant": "^1.0.0", "@epic-web/totp": "^4.0.1", "@kody-bot/connector-kit": "^2.0.0", diff --git a/packages/worker/src/app/admin-usage-data.node.test.ts b/packages/worker/src/app/admin-usage-data.node.test.ts index d69d06b590..ef7acfa85a 100644 --- a/packages/worker/src/app/admin-usage-data.node.test.ts +++ b/packages/worker/src/app/admin-usage-data.node.test.ts @@ -369,6 +369,80 @@ test('loadAdminUsageData shows email fallback limits for plan-less users', async ).toEqual(['email_receives_per_day']) }) +function createFakeKv() { + const store = new Map() + const kv = { + async get(key: string, _type: 'json') { + const raw = store.get(key) + return raw === undefined ? null : JSON.parse(raw) + }, + async put(key: string, value: string, _options?: unknown) { + store.set(key, value) + }, + async delete(key: string) { + store.delete(key) + }, + } as unknown as KVNamespace + return { kv, store } +} + +test('loadAdminUsageData caches rollup reads in KV and serves repeat loads from cache', async () => { + const email = 'cached-usage@example.com' + const usageUserId = await createStableUserIdFromEmail(email) + let rollupQueryCount = 0 + const db = createAdminUsageTestDb({ + users: [{ id: 1, username: 'cached', email, plan: null }], + usageRollups: [ + usageRow({ + user_id: usageUserId, + metric: 'execute', + month: '2026-07', + event_count: 7, + }), + ], + resourceCounts: { [usageUserId]: {} }, + }) + const countingDb = new Proxy(db, { + get(target, property, receiver) { + if (property === 'prepare') { + return (query: string) => { + if (normalizeQuery(query).includes('from usage_rollups')) { + rollupQueryCount += 1 + } + return target.prepare(query) + } + } + return Reflect.get(target, property, receiver) + }, + }) + const { kv, store } = createFakeKv() + const env = { APP_DB: countingDb, BUNDLE_ARTIFACTS_KV: kv } as Env + const now = new Date('2026-07-05T12:00:00.000Z') + + const first = await loadAdminUsageData( + env, + 'https://kody.local/admin/usage.json', + now, + ) + expect(getEventCount(first.users[0], 'execute')).toBe(7) + const queriesAfterFirstLoad = rollupQueryCount + expect(queriesAfterFirstLoad).toBeGreaterThan(0) + expect(store.size).toBeGreaterThan(0) + for (const key of store.keys()) { + expect(key.startsWith('derived-cache:v1:usage-rollups:')).toBe(true) + } + + const second = await loadAdminUsageData( + env, + 'https://kody.local/admin/usage.json', + now, + ) + expect(getEventCount(second.users[0], 'execute')).toBe(7) + expect(second.selectedUser?.monthUsage[0]?.month).toBe('2026-07') + // The repeat load is served from KV: no additional rollup queries. + expect(rollupQueryCount).toBe(queriesAfterFirstLoad) +}) + test('loadAdminUsageData keeps current-month and month-over-month rollups on UTC month boundaries', async () => { const email = 'month-boundary@example.com' const usageUserId = await createStableUserIdFromEmail(email) diff --git a/packages/worker/src/app/admin-usage-data.ts b/packages/worker/src/app/admin-usage-data.ts index d727b2719d..93e7189b52 100644 --- a/packages/worker/src/app/admin-usage-data.ts +++ b/packages/worker/src/app/admin-usage-data.ts @@ -1,3 +1,5 @@ +import { cachified, type Cache } from '@epic-web/cachified' +import { toHex } from '@kody-internal/shared/hex.ts' import { entitlementResourceLabels, isEmailFallbackResource, @@ -11,6 +13,7 @@ import { readEntitlementResourceUsage, utcDayKey, } from '#worker/entitlements/service.ts' +import { createKvCachifiedCache } from '#worker/kv-cachified.ts' import { resolveUserStableId } from '#worker/user-id.ts' import { type AdminUsageDailyCounter, @@ -64,6 +67,12 @@ const adminUsageDailyCounterResources = [ const defaultPageSize = 20 const maxPageSize = 100 const warningThreshold = 0.8 +/** + * Rollup rows are derived hourly from Analytics Engine in production, so a + * short KV cache on the per-page read model adds no meaningful staleness + * while keeping repeated admin page loads off D1. + */ +const rollupCacheTtlMs = 5 * 60 * 1000 type AdminUsageUserRow = { id: number @@ -107,6 +116,11 @@ export async function loadAdminUsageData( const selectedUserId = readPositiveInt(url.searchParams.get('userId'), 0) const currentMonth = utcMonthKey(now) const today = utcDayKey(now) + // Fall through to direct D1 queries when KV is unavailable (some tests + // construct a partial Env without the binding). + const rollupCache = env.BUNDLE_ARTIFACTS_KV + ? createKvCachifiedCache(env.BUNDLE_ARTIFACTS_KV) + : null const [totalResult, userRows] = await Promise.all([ env.APP_DB.prepare(`SELECT COUNT(*) AS total FROM users`).first<{ @@ -131,6 +145,7 @@ export async function loadAdminUsageData( const usageRows = await loadCurrentMonthRollups({ db: env.APP_DB, + cache: rollupCache, userIds: users.map((user) => user.usageUserId), month: currentMonth, }) @@ -153,6 +168,7 @@ export async function loadAdminUsageData( selectedUser: selectedUser ? await buildSelectedUser({ db: env.APP_DB, + cache: rollupCache, user: selectedUser, currentMonth, now, @@ -226,6 +242,7 @@ async function buildUserSummary(input: { async function buildSelectedUser(input: { db: D1Database + cache: Cache | null user: UserWithUsageId currentMonth: string now: Date @@ -236,6 +253,7 @@ async function buildSelectedUser(input: { user: input.user, usage: await loadCurrentMonthRollups({ db: input.db, + cache: input.cache, userIds: [input.user.usageUserId], month: input.currentMonth, }), @@ -243,7 +261,9 @@ async function buildSelectedUser(input: { }), loadUserMonthRollups({ db: input.db, + cache: input.cache, userId: input.user.usageUserId, + currentMonth: input.currentMonth, }), readEntitlementConsumption({ db: input.db, @@ -335,12 +355,38 @@ async function readEntitlementConsumption(input: { ) } +async function hashCacheKeyPart(value: string) { + const digest = await crypto.subtle.digest( + 'SHA-256', + new TextEncoder().encode(value), + ) + return toHex(new Uint8Array(digest)) +} + async function loadCurrentMonthRollups(input: { db: D1Database + cache: Cache | null userIds: Array month: string }) { if (input.userIds.length === 0) return [] + if (!input.cache) return await queryCurrentMonthRollups(input) + // The user id list can be a full admin page (up to 100 ids), so the key + // carries a digest of the ids rather than the ids themselves. + const usersKey = await hashCacheKeyPart(input.userIds.join('\n')) + return await cachified({ + key: `usage-rollups:month:${input.month}:users:${usersKey}`, + cache: input.cache, + ttl: rollupCacheTtlMs, + getFreshValue: () => queryCurrentMonthRollups(input), + }) +} + +async function queryCurrentMonthRollups(input: { + db: D1Database + userIds: Array + month: string +}) { const placeholders = input.userIds.map(() => '?').join(', ') const result = await input.db .prepare( @@ -354,7 +400,27 @@ async function loadCurrentMonthRollups(input: { return result.results ?? [] } -async function loadUserMonthRollups(input: { db: D1Database; userId: string }) { +async function loadUserMonthRollups(input: { + db: D1Database + cache: Cache | null + userId: string + currentMonth: string +}) { + if (!input.cache) return await queryUserMonthRollups(input) + return await cachified({ + // Keyed by the current month so cached history rolls over on UTC + // month boundaries without waiting for the TTL. + key: `usage-rollups:user:${input.userId}:asof:${input.currentMonth}`, + cache: input.cache, + ttl: rollupCacheTtlMs, + getFreshValue: () => queryUserMonthRollups(input), + }) +} + +async function queryUserMonthRollups(input: { + db: D1Database + userId: string +}) { const result = await input.db .prepare( `SELECT user_id, metric, month, event_count, error_count, diff --git a/packages/worker/src/entitlements/service.ts b/packages/worker/src/entitlements/service.ts index a0008ad313..91cbdb018d 100644 --- a/packages/worker/src/entitlements/service.ts +++ b/packages/worker/src/entitlements/service.ts @@ -63,11 +63,14 @@ export type StableUserAccount = { /** * Reverse-resolve a stable MCP userId back to the account email, plan, and - * verified-email state. Stored stable ids are queried directly; legacy rows - * without one still fall back to the normalized-email hash. - * Only call this on paths that genuinely have no caller context email - * (for example inbound email routing); interactive surfaces already carry - * the email. + * verified-email state. Stored stable ids are one indexed point read; legacy + * rows without one fall back to the scan in `findUserRowByStableUserId`, + * which hashes each email and writes the computed id back so the scan is + * paid at most once per row. + * Only call this on paths that genuinely have no caller context email (for + * example package-runtime contexts acting with only the hashed userId); + * inbound email routing resolves accounts via the indexed username lookup + * and interactive surfaces already carry the email. */ export async function findUserAccountByStableUserId( db: D1Database, diff --git a/packages/worker/src/kv-cachified.ts b/packages/worker/src/kv-cachified.ts new file mode 100644 index 0000000000..3bf6f6bf65 --- /dev/null +++ b/packages/worker/src/kv-cachified.ts @@ -0,0 +1,53 @@ +/** + * KV-backed cache adapter for `@epic-web/cachified`, used for short-TTL + * derived read models (for example the admin usage rollup page). Entries are + * stored as JSON in `BUNDLE_ARTIFACTS_KV` under the documented key prefix + * `derived-cache:v1:` and expire via KV `expirationTtl`, so no eviction code + * is needed. + * + * All KV failures are swallowed: a broken cache must degrade to the direct + * query, never break the read path. Callers own key uniqueness below the + * prefix and must scope every key by user id(s) — cross-user data sharing is + * a bug. + */ + +import { totalTtl, type Cache, type CacheEntry } from '@epic-web/cachified' + +export const derivedCacheKeyPrefix = 'derived-cache:v1:' + +/** KV rejects expirationTtl values below 60 seconds. */ +const minimumKvExpirationTtlSeconds = 60 + +export function createKvCachifiedCache(kv: KVNamespace): Cache { + return { + name: 'kv-derived-cache', + async get(key) { + try { + return await kv.get(derivedCacheKeyPrefix + key, 'json') + } catch (error) { + console.debug('kv-cachified-get-failed', key, error) + return null + } + }, + async set(key, entry) { + try { + const ttlMs = totalTtl(entry.metadata) + const expirationTtl = Number.isFinite(ttlMs) + ? Math.max(Math.ceil(ttlMs / 1000), minimumKvExpirationTtlSeconds) + : undefined + await kv.put(derivedCacheKeyPrefix + key, JSON.stringify(entry), { + expirationTtl, + }) + } catch (error) { + console.debug('kv-cachified-set-failed', key, error) + } + }, + async delete(key) { + try { + await kv.delete(derivedCacheKeyPrefix + key) + } catch (error) { + console.debug('kv-cachified-delete-failed', key, error) + } + }, + } +} diff --git a/packages/worker/src/maintenance-handler.node.test.ts b/packages/worker/src/maintenance-handler.node.test.ts index 8e54aaba72..2dacf040cd 100644 --- a/packages/worker/src/maintenance-handler.node.test.ts +++ b/packages/worker/src/maintenance-handler.node.test.ts @@ -1,10 +1,15 @@ import { expect, test } from 'vitest' -import { handleSecretMaintenanceRequest } from './maintenance-handler.ts' +import { createStableUserIdFromEmail } from '#worker/user-id.ts' +import { + backfillStableUserIds, + handleSecretMaintenanceRequest, + handleStableUserIdBackfillRequest, +} from './maintenance-handler.ts' function createRequest( - input: { method?: string; authorization?: string } = {}, + input: { method?: string; authorization?: string; path?: string } = {}, ) { - return new Request('http://localhost/__maintenance/test', { + return new Request(`http://localhost${input.path ?? '/__maintenance/test'}`, { method: input.method ?? 'POST', headers: input.authorization === undefined @@ -83,3 +88,137 @@ test('handleSecretMaintenanceRequest enforces auth and reports maintenance resul error: 'boom', }) }) + +type BackfillUserRow = { + id: number + email: string + stable_user_id: string | null +} + +function createBackfillTestDb(input: { + users: Array + failingIds?: ReadonlySet +}) { + const users = input.users.map((user) => ({ ...user })) + const selectPages: Array = [] + const db = { + prepare(query: string) { + const normalized = query.replace(/\s+/g, ' ').trim().toLowerCase() + return { + bind(...params: Array) { + return { + async all() { + if (!normalized.startsWith('select')) { + throw new Error(`Unsupported all query: ${query}`) + } + const afterId = Number(params[0]) + const limit = Number(params[1]) + const page = users + .filter((user) => user.id > afterId) + .sort((left, right) => left.id - right.id) + .slice(0, limit) + .map((user) => ({ ...user })) + selectPages.push(page.length) + return { results: page } + }, + async run() { + if (!normalized.startsWith('update users set stable_user_id')) { + throw new Error(`Unsupported run query: ${query}`) + } + const id = Number(params[1]) + if (input.failingIds?.has(id)) { + throw new Error(`UNIQUE constraint failed for user ${id}`) + } + const user = users.find( + (candidate) => + candidate.id === id && candidate.stable_user_id === null, + ) + if (user) user.stable_user_id = String(params[0]) + return { meta: { changes: user ? 1 : 0 } } + }, + } + }, + } + }, + } as unknown as D1Database + return { db, users, selectPages } +} + +test('backfillStableUserIds pages with a keyset and only touches NULL rows', async () => { + const alreadyStored = 'stored-stable-id' + const rows: Array = Array.from( + { length: 250 }, + (_, index) => ({ + id: index + 1, + email: `backfill-${index + 1}@example.com`, + stable_user_id: index % 5 === 0 ? alreadyStored : null, + }), + ) + const { db, users, selectPages } = createBackfillTestDb({ users: rows }) + + const result = await backfillStableUserIds({ db }) + + expect(result).toEqual({ scanned: 250, backfilled: 200, failed: 0 }) + // Keyset paging in batches of 100: two full pages and a final short one. + expect(selectPages).toEqual([100, 100, 50]) + expect(users[0]?.stable_user_id).toBe(alreadyStored) + expect(users[1]?.stable_user_id).toBe( + await createStableUserIdFromEmail('backfill-2@example.com'), + ) + expect(users.every((user) => user.stable_user_id !== null)).toBe(true) + + // Idempotent: a second run finds nothing left to backfill. + expect(await backfillStableUserIds({ db })).toEqual({ + scanned: 250, + backfilled: 0, + failed: 0, + }) +}) + +test('backfillStableUserIds counts per-row failures without aborting', async () => { + const { db, users } = createBackfillTestDb({ + users: [ + { id: 1, email: 'ok@example.com', stable_user_id: null }, + { id: 2, email: 'collides@example.com', stable_user_id: null }, + { id: 3, email: 'also-ok@example.com', stable_user_id: null }, + ], + failingIds: new Set([2]), + }) + + const result = await backfillStableUserIds({ db }) + + expect(result).toEqual({ scanned: 3, backfilled: 2, failed: 1 }) + expect(users[1]?.stable_user_id).toBeNull() + expect(users[2]?.stable_user_id).not.toBeNull() +}) + +test('handleStableUserIdBackfillRequest enforces the maintenance auth pattern', async () => { + const { db } = createBackfillTestDb({ + users: [{ id: 1, email: 'route@example.com', stable_user_id: null }], + }) + const path = '/__maintenance/backfill-stable-user-ids' + + const notConfiguredResponse = await handleStableUserIdBackfillRequest( + createRequest({ path, authorization: 'Bearer secret' }), + { APP_DB: db, CAPABILITY_REINDEX_SECRET: undefined }, + ) + expect(notConfiguredResponse.status).toBe(503) + + const unauthorizedResponse = await handleStableUserIdBackfillRequest( + createRequest({ path, authorization: 'Bearer wrong' }), + { APP_DB: db, CAPABILITY_REINDEX_SECRET: 'secret' }, + ) + expect(unauthorizedResponse.status).toBe(401) + + const successResponse = await handleStableUserIdBackfillRequest( + createRequest({ path, authorization: 'Bearer secret' }), + { APP_DB: db, CAPABILITY_REINDEX_SECRET: 'secret' }, + ) + expect(successResponse.status).toBe(200) + await expect(successResponse.json()).resolves.toEqual({ + ok: true, + scanned: 1, + backfilled: 1, + failed: 0, + }) +}) diff --git a/packages/worker/src/maintenance-handler.ts b/packages/worker/src/maintenance-handler.ts index 7ed52dae70..524903609d 100644 --- a/packages/worker/src/maintenance-handler.ts +++ b/packages/worker/src/maintenance-handler.ts @@ -1,4 +1,8 @@ import { getErrorMessage } from '#mcp/capabilities/error-message.ts' +import { + createStableUserIdFromEmail, + normalizeStableUserId, +} from '#worker/user-id.ts' type MaintenanceResult = Record & { ok?: never } @@ -58,3 +62,71 @@ export async function handleSecretMaintenanceRequest( ) } } + +export const stableUserIdBackfillBatchSize = 100 + +/** + * Persist the stable user id (SHA-256 of the normalized email) for every + * legacy users row that still has a NULL `stable_user_id`, so lookups hit the + * partial unique index (migration 0052) instead of the scan fallback. Pages + * with a keyset on `id` in batches of 100 to stay within D1 limits on large + * tables. Per-row failures (for example two rows whose emails normalize to + * the same hash colliding on the unique index) are counted, not fatal. + */ +export async function backfillStableUserIds(input: { + db: D1Database + batchSize?: number +}) { + const batchSize = input.batchSize ?? stableUserIdBackfillBatchSize + let lastId = 0 + let scanned = 0 + let backfilled = 0 + let failed = 0 + for (;;) { + const { results } = await input.db + .prepare( + `SELECT id, email, stable_user_id FROM users + WHERE id > ? ORDER BY id ASC LIMIT ?`, + ) + .bind(lastId, batchSize) + .all<{ id: number; email: string; stable_user_id: string | null }>() + const rows = results ?? [] + for (const row of rows) { + lastId = row.id + scanned += 1 + if (normalizeStableUserId(row.stable_user_id)) continue + try { + const result = await input.db + .prepare( + `UPDATE users SET stable_user_id = ? + WHERE id = ? AND stable_user_id IS NULL`, + ) + .bind(await createStableUserIdFromEmail(row.email), row.id) + .run() + backfilled += result.meta.changes ?? 0 + } catch (error) { + failed += 1 + console.warn('stable-user-id-backfill-row-failed', row.id, error) + } + } + if (rows.length < batchSize) break + } + return { scanned, backfilled, failed } +} + +/** + * `POST /__maintenance/backfill-stable-user-ids`: one-off/idempotent backfill + * of `users.stable_user_id` for legacy rows, authenticated like the reindex + * endpoints (Bearer `CAPABILITY_REINDEX_SECRET`). + */ +export async function handleStableUserIdBackfillRequest( + request: Request, + env: Pick, +): Promise { + return handleSecretMaintenanceRequest({ + request, + secret: env.CAPABILITY_REINDEX_SECRET, + notConfiguredMessage: 'Stable user id backfill is not configured', + run: () => backfillStableUserIds({ db: env.APP_DB }), + }) +} diff --git a/packages/worker/src/mcp/executor.node.test.ts b/packages/worker/src/mcp/executor.node.test.ts index 40a9ead953..c0cfa6e7e0 100644 --- a/packages/worker/src/mcp/executor.node.test.ts +++ b/packages/worker/src/mcp/executor.node.test.ts @@ -454,13 +454,9 @@ test('createExecuteExecutor records one usage event per sandbox run with duratio 'success', ]) expect(dataPoints[0]?.doubles?.[0]).toBeGreaterThanOrEqual(0) - expect(rollupWrites).toHaveLength(1) - expect(rollupWrites[0]?.slice(0, 4)).toEqual([ - 'usage-user-1', - 'execute', - String(rollupWrites[0]?.[7]).slice(0, 7), - 0, - ]) + // With USAGE_EVENTS present, rollups are derived from Analytics Engine + // by the scheduled aggregation instead of a per-event D1 upsert. + expect(rollupWrites).toHaveLength(0) // Sandbox run returning an error result: one error event. const errorLoader = { @@ -488,8 +484,7 @@ test('createExecuteExecutor records one usage event per sandbox run with duratio expect(errorResult.error).toBe('boom') expect(dataPoints).toHaveLength(2) expect(dataPoints[1]?.blobs?.[3]).toBe('error') - expect(rollupWrites).toHaveLength(2) - expect(rollupWrites[1]?.[3]).toBe(1) + expect(rollupWrites).toHaveLength(0) // Loader throwing: error event recorded, original error rethrown. const throwingLoader = { @@ -521,7 +516,7 @@ test('createExecuteExecutor records one usage event per sandbox run with duratio gatewayProps: { ...createGatewayProps('usage-user-1'), userId: null }, }).execute('async () => "ok"', providers) expect(dataPoints).toHaveLength(3) - expect(rollupWrites).toHaveLength(3) + expect(rollupWrites).toHaveLength(0) // Provider validation failure never reaches the sandbox: nothing recorded. const validationLoader = createFakeWorkerLoader() @@ -535,7 +530,7 @@ test('createExecuteExecutor records one usage event per sandbox run with duratio }).execute('async () => "ok"', [{ name: 'class', fns: {} }]) expect(validationResult.error).toContain('reserved') expect(dataPoints).toHaveLength(3) - expect(rollupWrites).toHaveLength(3) + expect(rollupWrites).toHaveLength(0) }) test('createExecuteExecutor rejects reserved JavaScript provider names before loading a worker', async () => { diff --git a/packages/worker/src/usage/aggregate-rollups.node.test.ts b/packages/worker/src/usage/aggregate-rollups.node.test.ts new file mode 100644 index 0000000000..4c9d734b78 --- /dev/null +++ b/packages/worker/src/usage/aggregate-rollups.node.test.ts @@ -0,0 +1,264 @@ +import { expect, test, vi } from 'vitest' +import { + aggregateUsageRollups, + resolveUsageEventsDataset, + shouldRunUsageAggregationCron, +} from './aggregate-rollups.ts' + +type BoundStatement = { sql: string; params: Array } + +function createFakeDb() { + const batches: Array> = [] + const db = { + prepare(sql: string) { + return { + bind(...params: Array) { + return { sql, params } + }, + } + }, + async batch(statements: Array) { + batches.push(statements) + return [] + }, + } as unknown as D1Database + return { db, batches } +} + +function createAggregationEnv(db: D1Database) { + return { + USAGE_EVENTS: { writeDataPoint() {} }, + APP_DB: db, + CLOUDFLARE_ACCOUNT_ID: 'account-1', + CLOUDFLARE_API_TOKEN: 'token-1', + } +} + +function stubFetchResponse(input: { status?: number; body: unknown }) { + const fetchMock = vi.fn(async () => { + return new Response( + typeof input.body === 'string' ? input.body : JSON.stringify(input.body), + { status: input.status ?? 200 }, + ) + }) + vi.stubGlobal('fetch', fetchMock) + return Object.assign(fetchMock, { + [Symbol.dispose]() { + vi.unstubAllGlobals() + }, + }) +} + +test('shouldRunUsageAggregationCron gates to the top of each hour', () => { + expect( + shouldRunUsageAggregationCron(new Date('2026-07-05T10:00:30.000Z')), + ).toBe(true) + expect( + shouldRunUsageAggregationCron(new Date('2026-07-05T10:30:00.000Z')), + ).toBe(false) + expect( + shouldRunUsageAggregationCron(new Date('2026-07-05T10:59:00.000Z')), + ).toBe(false) +}) + +test('resolveUsageEventsDataset picks the preview dataset only for preview', () => { + expect(resolveUsageEventsDataset({})).toBe('kody_usage_events') + expect(resolveUsageEventsDataset({ SENTRY_ENVIRONMENT: 'production' })).toBe( + 'kody_usage_events', + ) + expect(resolveUsageEventsDataset({ SENTRY_ENVIRONMENT: 'preview' })).toBe( + 'kody_usage_events_preview', + ) +}) + +test('aggregateUsageRollups no-ops when the binding or credentials are missing', async () => { + using fetchMock = stubFetchResponse({ body: { data: [] } }) + const { db, batches } = createFakeDb() + + for (const env of [ + { APP_DB: db }, + { APP_DB: db, USAGE_EVENTS: { writeDataPoint() {} } }, + { + APP_DB: db, + USAGE_EVENTS: { writeDataPoint() {} }, + CLOUDFLARE_ACCOUNT_ID: 'account-1', + }, + { + APP_DB: db, + USAGE_EVENTS: { writeDataPoint() {} }, + CLOUDFLARE_API_TOKEN: 'token-1', + }, + ]) { + await expect(aggregateUsageRollups(env, new Date())).resolves.toEqual({ + skipped: true, + reason: 'missing-analytics-config', + }) + } + expect(fetchMock).not.toHaveBeenCalled() + expect(batches).toHaveLength(0) +}) + +test('aggregateUsageRollups queries Analytics Engine month-to-date and upserts absolute values', async () => { + using fetchMock = stubFetchResponse({ + body: { + data: [ + { + user_id: 'user-a', + metric: 'execute', + event_count: 12, + error_count: 2, + total_duration_ms: 3456.7, + total_cpu_ms: 0, + total_bytes: 1024, + }, + { + user_id: 'user-b', + metric: 'email_send', + // The SQL API may serialize aggregates as strings. + event_count: '3', + error_count: '0', + total_duration_ms: '0', + total_cpu_ms: '0', + total_bytes: '2048', + }, + // Rows without an owning user or metric are never upserted. + { + user_id: '', + metric: 'execute', + event_count: 1, + error_count: 0, + total_duration_ms: 0, + total_cpu_ms: 0, + total_bytes: 0, + }, + ], + }, + }) + const { db, batches } = createFakeDb() + const now = new Date('2026-07-15T10:00:00.000Z') + + const result = await aggregateUsageRollups(createAggregationEnv(db), now) + + expect(result).toEqual({ + skipped: false, + month: '2026-07', + upsertedRows: 2, + users: 2, + }) + + expect(fetchMock).toHaveBeenCalledTimes(1) + const [url, init] = fetchMock.mock.calls[0] as unknown as [ + string, + RequestInit, + ] + expect(url).toBe( + 'https://api.cloudflare.com/client/v4/accounts/account-1/analytics_engine/sql', + ) + expect(init.method).toBe('POST') + expect((init.headers as Record)['authorization']).toBe( + 'Bearer token-1', + ) + const query = String(init.body) + expect(query).toContain('FROM kody_usage_events') + expect(query).toContain(`toDateTime('2026-07-01 00:00:00')`) + // Sampling-correct aggregates: counts and sums weight by _sample_interval. + expect(query).toContain('sum(_sample_interval) AS event_count') + expect(query).toContain( + `sum(if(blob4 = 'error', _sample_interval, 0)) AS error_count`, + ) + expect(query).toContain( + 'sum(double1 * _sample_interval) AS total_duration_ms', + ) + expect(query).toContain('GROUP BY blob1, blob2') + + expect(batches).toHaveLength(1) + const statements = batches[0] ?? [] + expect(statements).toHaveLength(2) + expect(statements[0]?.sql).toContain('ON CONFLICT (user_id, metric, month)') + expect(statements[0]?.sql).toContain('event_count = excluded.event_count') + expect(statements[0]?.sql).not.toContain('event_count + ') + expect(statements[0]?.params).toEqual([ + 'user-a', + 'execute', + '2026-07', + 12, + 2, + 3457, + 0, + 1024, + now.toISOString(), + ]) + expect(statements[1]?.params).toEqual([ + 'user-b', + 'email_send', + '2026-07', + 3, + 0, + 0, + 0, + 2048, + now.toISOString(), + ]) +}) + +test('aggregateUsageRollups honors CLOUDFLARE_API_BASE_URL and the preview dataset', async () => { + using fetchMock = stubFetchResponse({ body: { data: [] } }) + const { db, batches } = createFakeDb() + + const result = await aggregateUsageRollups( + { + ...createAggregationEnv(db), + CLOUDFLARE_API_BASE_URL: 'https://cloudflare-mock.local/', + SENTRY_ENVIRONMENT: 'preview', + }, + new Date('2026-07-15T10:00:00.000Z'), + ) + + expect(result).toEqual({ + skipped: false, + month: '2026-07', + upsertedRows: 0, + users: 0, + }) + const [url, init] = fetchMock.mock.calls[0] as unknown as [ + string, + RequestInit, + ] + expect(url).toBe( + 'https://cloudflare-mock.local/client/v4/accounts/account-1/analytics_engine/sql', + ) + expect(String(init.body)).toContain('FROM kody_usage_events_preview') + expect(batches).toHaveLength(0) +}) + +test('aggregateUsageRollups batches large result sets and throws on SQL API errors', async () => { + const manyRows = Array.from({ length: 120 }, (_, index) => ({ + user_id: `user-${index}`, + metric: 'execute', + event_count: 1, + error_count: 0, + total_duration_ms: 0, + total_cpu_ms: 0, + total_bytes: 0, + })) + using _manyRowsFetch = stubFetchResponse({ body: { data: manyRows } }) + const { db, batches } = createFakeDb() + + const result = await aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ) + expect(result).toMatchObject({ upsertedRows: 120, users: 120 }) + expect(batches.map((batch) => batch.length)).toEqual([50, 50, 20]) + + using _errorFetch = stubFetchResponse({ + status: 400, + body: 'query error: unknown table', + }) + await expect( + aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ), + ).rejects.toThrow('Analytics Engine SQL query failed (400)') +}) diff --git a/packages/worker/src/usage/aggregate-rollups.ts b/packages/worker/src/usage/aggregate-rollups.ts new file mode 100644 index 0000000000..f8c9dfde88 --- /dev/null +++ b/packages/worker/src/usage/aggregate-rollups.ts @@ -0,0 +1,202 @@ +/** + * Derived usage rollup aggregation. + * + * In production/preview, `recordUsage` writes usage events only to Workers + * Analytics Engine (see `record-usage.ts`); the D1 `usage_rollups` table is a + * derived aggregate. `aggregateUsageRollups` recomputes the current UTC + * month's rows from Analytics Engine via the SQL API and upserts them with + * absolute values — an idempotent recompute, not an increment — so the hourly + * cron can run any number of times without drift. Analytics Engine retention + * (~90 days) always covers a full month, so a month-to-date recompute is + * complete; rows for prior months already in D1 stay untouched. + * + * In local dev and tests (no `USAGE_EVENTS` binding or Cloudflare REST + * credentials) this is a no-op: `recordUsage` upserts `usage_rollups` + * directly there. + */ + +export const usageAggregationCronGateMinutes = 5 +export const usageAggregationCronIntervalMinutes = 60 + +/** + * Hourly gate for the scheduled handler, mirroring `shouldRunRetentionCron` + * in `packages/worker/src/app/retention.ts`. + */ +export function shouldRunUsageAggregationCron(now: Date) { + return ( + now.getUTCMinutes() < usageAggregationCronGateMinutes && + now.getUTCMinutes() % usageAggregationCronIntervalMinutes === 0 + ) +} + +export type UsageAggregationEnv = { + USAGE_EVENTS?: AnalyticsEngineDataset + APP_DB: D1Database + CLOUDFLARE_ACCOUNT_ID?: string + CLOUDFLARE_API_TOKEN?: string + CLOUDFLARE_API_BASE_URL?: string + SENTRY_ENVIRONMENT?: string +} + +export type UsageAggregationResult = + | { skipped: true; reason: string } + | { skipped: false; month: string; upsertedRows: number; users: number } + +/** + * The Analytics Engine SQL API dataset (= table name) written by the + * `USAGE_EVENTS` binding; see `packages/worker/wrangler.jsonc`. + */ +export function resolveUsageEventsDataset(env: { + SENTRY_ENVIRONMENT?: string +}) { + return env.SENTRY_ENVIRONMENT === 'preview' + ? 'kody_usage_events_preview' + : 'kody_usage_events' +} + +const upsertBatchSize = 50 + +const usageRollupAbsoluteUpsertStatement = ` +INSERT INTO usage_rollups ( + user_id, metric, month, + event_count, error_count, + total_duration_ms, total_cpu_ms, total_bytes, + updated_at +) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) +ON CONFLICT (user_id, metric, month) DO UPDATE SET + event_count = excluded.event_count, + error_count = excluded.error_count, + total_duration_ms = excluded.total_duration_ms, + total_cpu_ms = excluded.total_cpu_ms, + total_bytes = excluded.total_bytes, + updated_at = excluded.updated_at +`.trim() + +/** + * Analytics Engine samples data under load, so every aggregate must weight + * by `_sample_interval`: counts are `sum(_sample_interval)` and value sums + * are `sum(doubleN * _sample_interval)`. Blob/double positions match the + * data point layout in `record-usage.ts`. + */ +function buildMonthToDateAggregateQuery(dataset: string, monthStart: string) { + return ` +SELECT + blob1 AS user_id, + blob2 AS metric, + sum(_sample_interval) AS event_count, + sum(if(blob4 = 'error', _sample_interval, 0)) AS error_count, + sum(double1 * _sample_interval) AS total_duration_ms, + sum(double2 * _sample_interval) AS total_cpu_ms, + sum(double3 * _sample_interval) AS total_bytes +FROM ${dataset} +WHERE timestamp >= toDateTime('${monthStart}') +GROUP BY blob1, blob2 +FORMAT JSON +`.trim() +} + +type AnalyticsEngineSqlRow = { + user_id: string + metric: string + event_count: number | string + error_count: number | string + total_duration_ms: number | string + total_cpu_ms: number | string + total_bytes: number | string +} + +async function queryAnalyticsEngineSql(input: { + accountId: string + apiToken: string + baseUrl: string + query: string +}): Promise> { + const url = `${input.baseUrl.replace(/\/$/, '')}/client/v4/accounts/${input.accountId}/analytics_engine/sql` + const response = await fetch(url, { + method: 'POST', + headers: { + authorization: `Bearer ${input.apiToken}`, + }, + body: input.query, + }) + const text = await response.text() + if (!response.ok) { + throw new Error( + `Analytics Engine SQL query failed (${response.status}): ${text.slice(0, 500)}`, + ) + } + const parsed = JSON.parse(text) as { + data?: Array + } + return parsed.data ?? [] +} + +function toCount(value: number | string) { + const parsed = Number(value) + return Number.isFinite(parsed) ? Math.round(parsed) : 0 +} + +/** + * Recompute the current UTC month's `usage_rollups` rows from Analytics + * Engine and upsert them into D1 with absolute values. No-op (with a debug + * log) when the Analytics Engine binding or the Cloudflare REST credentials + * are unavailable. + */ +export async function aggregateUsageRollups( + env: UsageAggregationEnv, + now: Date = new Date(), +): Promise { + const accountId = env.CLOUDFLARE_ACCOUNT_ID?.trim() + const apiToken = env.CLOUDFLARE_API_TOKEN?.trim() + if (!env.USAGE_EVENTS || !accountId || !apiToken) { + console.debug( + 'usage-rollup-aggregation-skipped', + 'missing USAGE_EVENTS binding or Cloudflare REST credentials', + ) + return { skipped: true, reason: 'missing-analytics-config' } + } + + const month = now.toISOString().slice(0, 'YYYY-MM'.length) + const rows = await queryAnalyticsEngineSql({ + accountId, + apiToken, + baseUrl: + env.CLOUDFLARE_API_BASE_URL?.trim() || 'https://api.cloudflare.com', + query: buildMonthToDateAggregateQuery( + resolveUsageEventsDataset(env), + `${month}-01 00:00:00`, + ), + }) + + const updatedAt = now.toISOString() + const statements = rows + .filter((row) => row.user_id && row.metric) + .map((row) => + env.APP_DB.prepare(usageRollupAbsoluteUpsertStatement).bind( + row.user_id, + row.metric, + month, + toCount(row.event_count), + toCount(row.error_count), + toCount(row.total_duration_ms), + toCount(row.total_cpu_ms), + toCount(row.total_bytes), + updatedAt, + ), + ) + const users = new Set( + rows.filter((row) => row.user_id && row.metric).map((row) => row.user_id), + ).size + for (let index = 0; index < statements.length; index += upsertBatchSize) { + await env.APP_DB.batch(statements.slice(index, index + upsertBatchSize)) + } + + const result = { + skipped: false as const, + month, + upsertedRows: statements.length, + users, + } + console.info('usage-rollup-aggregation', JSON.stringify(result)) + return result +} diff --git a/packages/worker/src/usage/record-usage.ts b/packages/worker/src/usage/record-usage.ts index 60e4bce4ef..92b1aa45b3 100644 --- a/packages/worker/src/usage/record-usage.ts +++ b/packages/worker/src/usage/record-usage.ts @@ -4,11 +4,18 @@ * One event schema covers every metered chokepoint (execute runs, package * export invocations, job runs, workflow runs, package service runtime, * realtime websocket sessions, gateway fetches, email sends and receives). - * Events are written to two sinks: * - * - Workers Analytics Engine (`USAGE_EVENTS`) for high-cardinality analysis. - * - The D1 `usage_rollups` table (per-user, per-metric, per-month counters) - * which future quota enforcement can read cheaply. + * The write path depends on the environment: + * + * - When the Workers Analytics Engine binding (`USAGE_EVENTS`) is present + * (production/preview), the event is written **only** to Analytics Engine. + * The D1 `usage_rollups` table is then a derived aggregate, recomputed + * hourly from Analytics Engine by + * `packages/worker/src/usage/aggregate-rollups.ts` — a per-event D1 upsert + * would serialize every metered request on D1's single writer. + * - When `USAGE_EVENTS` is absent (local dev, tests), the event is upserted + * directly into `usage_rollups` so local admin pages and tests keep working + * without Analytics Engine access. * * `recordUsage` never throws and never rejects; metering must not break the * paths it observes. In local dev and tests where a binding is missing it @@ -71,14 +78,20 @@ ON CONFLICT (user_id, metric, month) DO UPDATE SET `.trim() /** - * Record one usage event to Analytics Engine and the D1 rollup table. + * Record one usage event. + * + * With `USAGE_EVENTS` bound this is a single non-blocking `writeDataPoint` + * call; the D1 rollup is derived later by the hourly aggregation cron. + * Without it (local dev, tests) the event is upserted into `usage_rollups` + * directly. * * Guarantees: * - Never throws and never rejects: failures are logged at debug level. * - Each sink degrades independently when its binding is unavailable. * - * Callers should `await` it (cheap: one `writeDataPoint` plus one D1 upsert) - * or hand the promise to `ctx.waitUntil(...)` inside Durable Objects. + * Callers should `await` it (cheap: one `writeDataPoint`, or one D1 upsert in + * local dev) or hand the promise to `ctx.waitUntil(...)` inside Durable + * Objects. */ export async function recordUsage( env: UsageEnv, @@ -90,7 +103,11 @@ export async function recordUsage( return } const timestamp = event.timestamp ?? new Date().toISOString() - writeUsageDataPoint(env, event, timestamp) + if (env.USAGE_EVENTS) { + writeUsageDataPoint(env, event, timestamp) + return + } + console.debug('usage-event-local', JSON.stringify({ ...event, timestamp })) await writeUsageRollup(env, event, timestamp) } catch (error) { console.warn('usage-event-record-failed', error) @@ -102,10 +119,7 @@ function writeUsageDataPoint( event: UsageEvent, timestamp: string, ) { - if (!env.USAGE_EVENTS) { - console.debug('usage-event-local', JSON.stringify({ ...event, timestamp })) - return - } + if (!env.USAGE_EVENTS) return try { env.USAGE_EVENTS.writeDataPoint({ indexes: [event.userId], diff --git a/packages/worker/src/usage/record-usage.workers.test.ts b/packages/worker/src/usage/record-usage.workers.test.ts index 860b86da38..87d0814537 100644 --- a/packages/worker/src/usage/record-usage.workers.test.ts +++ b/packages/worker/src/usage/record-usage.workers.test.ts @@ -16,7 +16,7 @@ async function listRollups(db: D1Database, userId: string) { return results } -test('recordUsage accumulates per-user monthly rollups and writes analytics data points', async () => { +test('recordUsage writes only Analytics Engine data points when USAGE_EVENTS is bound', async () => { await ensureUsageRollupsTestSchema(env.APP_DB) const userA = `usage-user-a-${crypto.randomUUID()}` const userB = `usage-user-b-${crypto.randomUUID()}` @@ -30,6 +30,55 @@ test('recordUsage accumulates per-user monthly rollups and writes analytics data }, } + await recordUsage(usageEnv, { + userId: userA, + eventType: 'execute', + durationMs: 120, + outcome: 'success', + timestamp: '2026-07-05T10:00:00.000Z', + }) + await recordUsage(usageEnv, { + userId: userA, + eventType: 'execute', + entityId: 'pkg-1', + durationMs: 80, + bytes: 512, + outcome: 'error', + timestamp: '2026-07-05T11:00:00.000Z', + }) + await recordUsage(usageEnv, { + userId: userB, + eventType: 'execute', + durationMs: 40, + outcome: 'success', + timestamp: '2026-07-05T12:00:00.000Z', + }) + + expect(dataPoints).toHaveLength(3) + expect(dataPoints[0]).toEqual({ + indexes: [userA], + blobs: [userA, 'execute', '', 'success', '2026-07-05T10:00:00.000Z'], + doubles: [120, 0, 0], + }) + expect(dataPoints[1]).toEqual({ + indexes: [userA], + blobs: [userA, 'execute', 'pkg-1', 'error', '2026-07-05T11:00:00.000Z'], + doubles: [80, 0, 512], + }) + expect(dataPoints[2]?.indexes).toEqual([userB]) + + // Production path: usage_rollups is a derived aggregate recomputed by the + // hourly aggregation cron, never written per event. + expect(await listRollups(env.APP_DB, userA)).toEqual([]) + expect(await listRollups(env.APP_DB, userB)).toEqual([]) +}) + +test('recordUsage accumulates per-user monthly rollups without USAGE_EVENTS (local dev)', async () => { + await ensureUsageRollupsTestSchema(env.APP_DB) + const userA = `usage-user-a-${crypto.randomUUID()}` + const userB = `usage-user-b-${crypto.randomUUID()}` + const usageEnv = { APP_DB: env.APP_DB } + await recordUsage(usageEnv, { userId: userA, eventType: 'execute', @@ -97,19 +146,6 @@ test('recordUsage accumulates per-user monthly rollups and writes analytics data total_bytes: 0, }, ]) - - expect(dataPoints).toHaveLength(4) - expect(dataPoints[0]).toEqual({ - indexes: [userA], - blobs: [userA, 'execute', '', 'success', '2026-07-05T10:00:00.000Z'], - doubles: [120, 0, 0], - }) - expect(dataPoints[1]).toEqual({ - indexes: [userA], - blobs: [userA, 'execute', 'pkg-1', 'error', '2026-07-05T11:00:00.000Z'], - doubles: [80, 0, 512], - }) - expect(dataPoints[3]?.indexes).toEqual([userB]) }) test('recordUsage never throws when bindings are missing, sinks fail, or userId is empty', async () => { @@ -121,7 +157,8 @@ test('recordUsage never throws when bindings are missing, sinks fail, or userId recordUsage({}, { userId, eventType: 'execute', outcome: 'success' }), ).resolves.toBeUndefined() - // Analytics Engine sink throws; rollup must still be written. + // Analytics Engine sink throws: degrade, don't throw, and never fall back + // to the per-event D1 upsert (production must not write rollups inline). await expect( recordUsage( { @@ -142,18 +179,7 @@ test('recordUsage never throws when bindings are missing, sinks fail, or userId }, ), ).resolves.toBeUndefined() - expect(await listRollups(env.APP_DB, userId)).toEqual([ - { - user_id: userId, - metric: 'job_run', - month: '2026-07', - event_count: 1, - error_count: 0, - total_duration_ms: 10, - total_cpu_ms: 0, - total_bytes: 0, - }, - ]) + expect(await listRollups(env.APP_DB, userId)).toEqual([]) // Missing userId: skipped entirely, no row written. await expect( diff --git a/packages/worker/src/user-id.ts b/packages/worker/src/user-id.ts index 5b92ada98a..4a50b47c1f 100644 --- a/packages/worker/src/user-id.ts +++ b/packages/worker/src/user-id.ts @@ -76,6 +76,51 @@ function withoutStableUserIdColumn(select: string) { return select.replace(/,\s*stable_user_id\b/g, '') } +/** + * Best-effort write-back for the scan fallback: persist the computed stable + * id on a legacy row that still has a NULL `stable_user_id`, so each such row + * pays the full-table scan at most once. The `stable_user_id IS NULL` guard + * means a concurrently-materialized id is never overwritten, and errors + * (missing column on pre-migration databases, read-only test doubles) are + * swallowed: healing must never break the lookup. + */ +async function persistHealedStableUserId(input: { + db: D1Database + stableUserId: string + row: { id?: unknown; email?: unknown } +}) { + try { + if (typeof input.row.id === 'number') { + await input.db + .prepare( + `UPDATE users SET stable_user_id = ? + WHERE id = ? AND stable_user_id IS NULL`, + ) + .bind(input.stableUserId, input.row.id) + .run() + return + } + if (typeof input.row.email === 'string' && input.row.email) { + await input.db + .prepare( + `UPDATE users SET stable_user_id = ? + WHERE email = ? AND stable_user_id IS NULL`, + ) + .bind(input.stableUserId, input.row.email) + .run() + } + } catch (error) { + console.debug('stable-user-id-write-back-failed', error) + } +} + +/** + * Resolve a users row from a stable user id: one indexed point read on + * `stable_user_id`, falling back to a full scan that hashes each email for + * legacy rows without a materialized id. When the scan matches a row with a + * NULL `stable_user_id` (and the select includes `id` or `email`), the + * computed id is written back so the row is found via the index next time. + */ export async function findUserRowByStableUserId< T extends UserStableIdRow & { id: number }, >(input: { @@ -108,7 +153,11 @@ export async function findUserRowByStableUserId< .all() } for (const row of rows.results ?? []) { - if ((await resolveUserStableId(row)) === stableUserId) return row + if ((await resolveUserStableId(row)) !== stableUserId) continue + if (!normalizeStableUserId(row.stable_user_id)) { + await persistHealedStableUserId({ db: input.db, stableUserId, row }) + } + return row } return null } diff --git a/packages/worker/src/user-id.workers.test.ts b/packages/worker/src/user-id.workers.test.ts new file mode 100644 index 0000000000..2bd6f7b84a --- /dev/null +++ b/packages/worker/src/user-id.workers.test.ts @@ -0,0 +1,139 @@ +import { env } from 'cloudflare:workers' +import { expect, test } from 'vitest' +import { + createStableUserIdFromEmail, + findUserRowByStableUserId, +} from './user-id.ts' + +const usersSelect = `SELECT id, email, stable_user_id FROM users` + +async function recreateUsersTable(input: { + db: D1Database + withStableUserIdColumn: boolean +}) { + await input.db.prepare(`DROP TABLE IF EXISTS users`).run() + await input.db + .prepare( + `CREATE TABLE users ( + id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, + username TEXT NOT NULL UNIQUE, + email TEXT NOT NULL UNIQUE, + password_hash TEXT NOT NULL${ + input.withStableUserIdColumn ? ',\n\tstable_user_id TEXT' : '' + } +)`, + ) + .run() + if (input.withStableUserIdColumn) { + // Mirrors migration 0052. + await input.db + .prepare( + `CREATE UNIQUE INDEX idx_users_stable_user_id + ON users(stable_user_id) + WHERE stable_user_id IS NOT NULL`, + ) + .run() + } +} + +async function seedUser(input: { + db: D1Database + email: string + stableUserId?: string | null +}) { + await input.db + .prepare( + `INSERT INTO users (username, email, password_hash, stable_user_id) + VALUES (?, ?, 'test-password-hash', ?)`, + ) + .bind( + `user-id-${crypto.randomUUID().slice(0, 8)}`, + input.email, + input.stableUserId ?? null, + ) + .run() +} + +async function readStoredStableUserId(db: D1Database, email: string) { + const row = await db + .prepare(`SELECT stable_user_id FROM users WHERE email = ?`) + .bind(email) + .first<{ stable_user_id: string | null }>() + return row?.stable_user_id ?? null +} + +test('scan fallback heals NULL stable_user_id rows with a write-back', async () => { + await recreateUsersTable({ db: env.APP_DB, withStableUserIdColumn: true }) + const emailA = `legacy-a-${crypto.randomUUID()}@example.com` + const emailB = `legacy-b-${crypto.randomUUID()}@example.com` + await seedUser({ db: env.APP_DB, email: emailA }) + await seedUser({ db: env.APP_DB, email: emailB }) + const stableUserIdB = await createStableUserIdFromEmail(emailB) + + const row = await findUserRowByStableUserId({ + db: env.APP_DB, + stableUserId: stableUserIdB, + select: usersSelect, + }) + expect(row?.email).toBe(emailB) + + // The matched legacy row got its computed stable id persisted; the + // non-matching row was left alone. + expect(await readStoredStableUserId(env.APP_DB, emailB)).toBe(stableUserIdB) + expect(await readStoredStableUserId(env.APP_DB, emailA)).toBeNull() + + // The healed row is now found via the indexed point read. + const healedRow = await findUserRowByStableUserId({ + db: env.APP_DB, + stableUserId: stableUserIdB, + select: usersSelect, + }) + expect(healedRow?.email).toBe(emailB) + expect(healedRow?.stable_user_id).toBe(stableUserIdB) +}) + +test('stored stable ids resolve via the index and are never overwritten', async () => { + await recreateUsersTable({ db: env.APP_DB, withStableUserIdColumn: true }) + // A stored id that differs from the email hash (the email changed after + // signup) must stay authoritative. + const email = `changed-${crypto.randomUUID()}@example.com` + const storedStableUserId = crypto.randomUUID().replaceAll('-', '') + await seedUser({ db: env.APP_DB, email, stableUserId: storedStableUserId }) + + const row = await findUserRowByStableUserId({ + db: env.APP_DB, + stableUserId: storedStableUserId, + select: usersSelect, + }) + expect(row?.email).toBe(email) + expect(await readStoredStableUserId(env.APP_DB, email)).toBe( + storedStableUserId, + ) + + expect( + await findUserRowByStableUserId({ + db: env.APP_DB, + stableUserId: `missing-${storedStableUserId}`, + select: usersSelect, + }), + ).toBeNull() +}) + +test('pre-migration databases without the column still resolve via the scan', async () => { + await recreateUsersTable({ db: env.APP_DB, withStableUserIdColumn: false }) + const email = `pre-migration-${crypto.randomUUID()}@example.com` + await env.APP_DB.prepare( + `INSERT INTO users (username, email, password_hash) + VALUES (?, ?, 'test-password-hash')`, + ) + .bind(`user-id-${crypto.randomUUID().slice(0, 8)}`, email) + .run() + + // The write-back UPDATE fails on the missing column and is swallowed. + const row = await findUserRowByStableUserId({ + db: env.APP_DB, + stableUserId: await createStableUserIdFromEmail(email), + select: usersSelect, + }) + expect(row?.email).toBe(email) +}) From d650a4106ba9ba9e72bb255b265b5367f245ab14 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:09 +0000 Subject: [PATCH 03/14] Retention: loop to a time budget and cover growth tables Hourly retention now drains batches round-robin under a 20s budget, adds policies for email_messages (365d, R2-blob aware), entitlement daily counters (400d) and usage_rollups (24mo), replaces the whole-table ROW_NUMBER prune with a bounded cap-pair pass, prunes stale source snapshot KV keys, batches system-email deletes, and stops writing the legacy retriever scope index blob. Adds supporting indexes. Co-authored-by: Kent C. Dodds --- .../migrations/0055-retention-indexes.sql | 43 ++ .../worker/src/app/retention.node.test.ts | 416 ++++++++++- packages/worker/src/app/retention.ts | 658 +++++++++++++++--- packages/worker/src/email/system-email.ts | 133 +++- .../src/email/system-email.workers.test.ts | 66 +- .../manifest-cache.node.test.ts | 161 ++++- .../src/package-retrievers/manifest-cache.ts | 104 +-- 7 files changed, 1383 insertions(+), 198 deletions(-) create mode 100644 packages/worker/migrations/0055-retention-indexes.sql diff --git a/packages/worker/migrations/0055-retention-indexes.sql b/packages/worker/migrations/0055-retention-indexes.sql new file mode 100644 index 0000000000..5ca161f476 --- /dev/null +++ b/packages/worker/migrations/0055-retention-indexes.sql @@ -0,0 +1,43 @@ +-- Additive indexes supporting the global retention prunes in +-- packages/worker/src/app/retention.ts. The retention queries order by a +-- time column across all users, which the existing (user_id, ...) composite +-- indexes cannot serve. + +-- Age prune: WHERE created_at < ? ORDER BY created_at ASC. +CREATE INDEX IF NOT EXISTS idx_package_invocations_created_at +ON package_invocations(created_at); + +-- Terminal workflow prune filters on status before ordering on the +-- completed_at/updated_at/created_at coalesce. +CREATE INDEX IF NOT EXISTS idx_workflow_runs_status_updated_at +ON workflow_runs(status, updated_at); + +-- Daily entitlement counter prune: WHERE day < ? ORDER BY day ASC. +CREATE INDEX IF NOT EXISTS idx_entitlement_daily_counters_day +ON entitlement_daily_counters(day); + +-- Usage rollup prune: WHERE month < ? ORDER BY month ASC. +CREATE INDEX IF NOT EXISTS idx_usage_rollups_month +ON usage_rollups(month); + +-- Email message age prune crosses users: WHERE user_id != 'system:email' +-- AND created_at < ? ORDER BY created_at ASC. +CREATE INDEX IF NOT EXISTS idx_email_messages_created_at +ON email_messages(created_at); + +-- Email delivery event prune also orders on created_at across users. +CREATE INDEX IF NOT EXISTS idx_email_delivery_events_created_at +ON email_delivery_events(created_at); + +-- Runtime run age prune orders on started_at across users; the per-package +-- cap prune keeps using idx_package_runtime_runs_package_started. +CREATE INDEX IF NOT EXISTS idx_package_runtime_runs_started_at +ON package_runtime_runs(started_at); + +-- Orphan runtime log prune orders on created_at across users. +CREATE INDEX IF NOT EXISTS idx_package_runtime_logs_created_at +ON package_runtime_logs(created_at); + +-- Published bundle artifact prune orders on created_at across users. +CREATE INDEX IF NOT EXISTS idx_published_bundle_artifacts_created_at +ON published_bundle_artifacts(created_at); diff --git a/packages/worker/src/app/retention.node.test.ts b/packages/worker/src/app/retention.node.test.ts index 2c4ab6df49..dd70727050 100644 --- a/packages/worker/src/app/retention.node.test.ts +++ b/packages/worker/src/app/retention.node.test.ts @@ -4,16 +4,22 @@ import { expect, test, vi } from 'vitest' import { auditEventRetentionDays, emailDeliveryEventRetentionDays, + emailMessageRetentionDays, + entitlementDailyCounterRetentionDays, getRetentionPolicyCoverage, memorySuppressionRetentionDays, packageInvocationRetentionDays, packageRuntimeRunRetentionDays, pruneAuditEventsForRetention, pruneEmailDeliveryEventsForRetention, + pruneEntitlementDailyCountersForRetention, pruneMemorySuppressionsForRetention, prunePackageInvocationsForRetention, prunePackageRuntimeRetention, prunePublishedBundleArtifactsForRetention, + pruneRetention, + pruneUsageRollupsForRetention, + pruneUserEmailMessagesForRetention, pruneWorkflowRunsForRetention, publishedBundleArtifactRetentionDays, retentionPolicies, @@ -214,6 +220,59 @@ function createRetentionDb() { detail_json TEXT, created_at TEXT NOT NULL ); + CREATE TABLE email_threads ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + inbox_id TEXT, + subject_normalized TEXT NOT NULL DEFAULT '', + root_message_id_header TEXT, + last_message_at TEXT NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE TABLE email_messages ( + id TEXT PRIMARY KEY, + direction TEXT NOT NULL, + user_id TEXT NOT NULL, + inbox_id TEXT, + thread_id TEXT, + from_address TEXT NOT NULL, + subject TEXT NOT NULL DEFAULT '', + processing_status TEXT NOT NULL, + raw_mime_key TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE TABLE email_attachments ( + id TEXT PRIMARY KEY, + message_id TEXT NOT NULL, + filename TEXT, + content_type TEXT NOT NULL, + size INTEGER NOT NULL DEFAULT 0, + storage_kind TEXT NOT NULL, + 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, + month TEXT NOT NULL, + event_count INTEGER NOT NULL DEFAULT 0, + error_count INTEGER NOT NULL DEFAULT 0, + total_duration_ms INTEGER NOT NULL DEFAULT 0, + total_cpu_ms INTEGER NOT NULL DEFAULT 0, + total_bytes INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP), + PRIMARY KEY (user_id, metric, month) + ); CREATE TABLE audit_events ( id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, category TEXT NOT NULL, @@ -283,6 +342,66 @@ function insertRuntimeRun( ) } +function insertEmailMessage( + db: DatabaseSync, + input: { + id: string + userId?: string + threadId?: string | null + rawMimeKey?: string | null + createdAt: string + }, +) { + db.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, thread_id, from_address, subject, + processing_status, raw_mime_key, created_at, updated_at + ) VALUES (?, 'inbound', ?, ?, 'sender@example.net', 'Subject', 'stored', ?, ?, ?)`, + ).run( + input.id, + input.userId ?? 'user-1', + input.threadId ?? null, + input.rawMimeKey ?? null, + input.createdAt, + input.createdAt, + ) +} + +function insertEmailThread( + db: DatabaseSync, + input: { id: string; userId?: string; createdAt: string }, +) { + db.prepare( + `INSERT INTO email_threads ( + id, user_id, last_message_at, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?)`, + ).run( + input.id, + input.userId ?? 'user-1', + input.createdAt, + input.createdAt, + input.createdAt, + ) +} + +function insertPackageInvocation( + db: DatabaseSync, + input: { id: string; status?: string; createdAt: string }, +) { + db.prepare( + `INSERT INTO package_invocations ( + id, user_id, token_id, package_id, package_kody_id, export_name, + idempotency_key, request_hash, status, created_at, updated_at + ) VALUES (?, 'user-1', 'token', 'pkg-1', 'pkg-1', '.', ?, 'hash', ?, ?, ?)`, + ).run( + input.id, + `key-${input.id}`, + input.status ?? 'completed', + input.createdAt, + input.createdAt, + ) +} + function idsForTable(db: DatabaseSync, table: string) { return ( db.prepare(`SELECT id FROM ${table} ORDER BY id ASC`).all() as Array<{ @@ -294,12 +413,17 @@ function idsForTable(db: DatabaseSync, table: string) { test('retention manifest covers the requested growth tables and cron is hourly gated', () => { expect(retentionPolicies.map((policy) => policy.table).sort()).toEqual([ 'audit_events', + 'email_attachments', 'email_delivery_events', + 'email_messages', + 'email_threads', + 'entitlement_daily_counters', 'mcp_memory_conversation_suppressions', 'package_invocations', 'package_runtime_logs', 'package_runtime_runs', 'published_bundle_artifacts', + 'usage_rollups', 'workflow_runs', ]) expect(shouldRunRetentionCron(new Date('2026-07-07T03:00:00.000Z'))).toBe( @@ -386,6 +510,7 @@ test('runtime retention prunes old runs, caps per package, and keeps active refe deletedRuns: 3, deletedLogs: 3, deletedOrphanLogs: 1, + hasMore: false, }) expect(idsForTable(sqlite, 'package_runtime_runs')).toContain('boundary-keep') expect(idsForTable(sqlite, 'package_runtime_runs')).toContain('running-keep') @@ -405,6 +530,55 @@ test('runtime retention prunes old runs, caps per package, and keeps active refe ) }) +test('runtime cap retention ranks bounded candidate pairs largest first', async () => { + const { sqlite, db } = createRetentionDb() + for (const [packageId, runCount] of [ + ['pkg-large', 503], + ['pkg-small', 502], + ] as const) { + for (let index = 0; index < runCount; index += 1) { + insertRuntimeRun(sqlite, { + id: `${packageId}-${String(index).padStart(3, '0')}`, + packageId, + startedAt: new Date(now.getTime() - index * 1000).toISOString(), + }) + } + } + + const first = await prunePackageRuntimeRetention({ + db, + now, + batchSize: 50, + maxCapPairs: 1, + }) + expect(first.deletedRuns).toBe(3) + expect(first.hasMore).toBe(true) + expect(idsForTable(sqlite, 'package_runtime_runs')).not.toContain( + 'pkg-large-502', + ) + expect(idsForTable(sqlite, 'package_runtime_runs')).toContain('pkg-small-501') + + const second = await prunePackageRuntimeRetention({ + db, + now, + batchSize: 50, + maxCapPairs: 1, + }) + expect(second.deletedRuns).toBe(2) + expect(idsForTable(sqlite, 'package_runtime_runs')).not.toContain( + 'pkg-small-501', + ) + + const third = await prunePackageRuntimeRetention({ + db, + now, + batchSize: 50, + maxCapPairs: 1, + }) + expect(third.deletedRuns).toBe(0) + expect(third.hasMore).toBe(false) +}) + test('package invocation and workflow retention keeps boundary and active idempotency rows', async () => { const { sqlite, db } = createRetentionDb() for (const [id, status, createdAt] of [ @@ -541,7 +715,169 @@ test('memory suppression, email delivery, and audit retention respect boundaries ]) }) -test('published bundle artifact retention deletes only stale unreferenced rows and KV blobs', async () => { +test('email message retention deletes old user rows, R2 blobs, attachments, and orphan threads', async () => { + const { sqlite, db } = createRetentionDb() + insertEmailThread(sqlite, { + id: 'thread-orphan', + createdAt: daysAgo(emailMessageRetentionDays + 2), + }) + insertEmailThread(sqlite, { + id: 'thread-mixed', + createdAt: daysAgo(emailMessageRetentionDays + 2), + }) + insertEmailMessage(sqlite, { + id: 'msg-old-blob', + threadId: 'thread-orphan', + rawMimeKey: 'raw-mime/user-1/msg-old-blob', + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + insertEmailMessage(sqlite, { + id: 'msg-old-plain', + threadId: 'thread-mixed', + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + insertEmailMessage(sqlite, { + id: 'msg-fresh', + threadId: 'thread-mixed', + createdAt: daysAgo(1), + }) + insertEmailMessage(sqlite, { + id: 'msg-boundary', + createdAt: daysAgo(emailMessageRetentionDays), + }) + insertEmailMessage(sqlite, { + id: 'msg-old-system', + userId: systemEmailOwnerId, + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + sqlite + .prepare( + `INSERT INTO email_attachments ( + id, message_id, content_type, size, storage_kind, created_at + ) VALUES ('att-old', 'msg-old-blob', 'text/plain', 1, 'raw-mime', ?)`, + ) + .run(daysAgo(emailMessageRetentionDays + 1)) + const blobDelete = vi.fn(async () => undefined) + + const result = await pruneUserEmailMessagesForRetention({ + db, + blobs: { delete: blobDelete } as unknown as Pick, + now, + }) + + expect(result).toEqual({ + deletedMessages: 2, + deletedAttachments: 1, + deletedThreads: 1, + deletedRawMimeBlobs: 1, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 0, + hasMore: false, + }) + expect(blobDelete).toHaveBeenCalledWith(['raw-mime/user-1/msg-old-blob']) + expect(idsForTable(sqlite, 'email_messages')).toEqual([ + 'msg-boundary', + 'msg-fresh', + 'msg-old-system', + ]) + expect(idsForTable(sqlite, 'email_attachments')).toEqual([]) + expect(idsForTable(sqlite, 'email_threads')).toEqual(['thread-mixed']) +}) + +test('email message retention never deletes rows whose blob cannot be deleted first', async () => { + const { sqlite, db } = createRetentionDb() + insertEmailMessage(sqlite, { + id: 'msg-old-blob', + rawMimeKey: 'raw-mime/user-1/msg-old-blob', + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + insertEmailMessage(sqlite, { + id: 'msg-old-plain', + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + + const missingBinding = await pruneUserEmailMessagesForRetention({ db, now }) + expect(missingBinding).toMatchObject({ + deletedMessages: 1, + deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 1, + blobDeleteErrors: 0, + }) + expect(idsForTable(sqlite, 'email_messages')).toEqual(['msg-old-blob']) + + const failingDelete = vi.fn(async () => { + throw new Error('r2 unavailable') + }) + const failedDelete = await pruneUserEmailMessagesForRetention({ + db, + blobs: { delete: failingDelete } as unknown as Pick, + now, + }) + expect(failedDelete).toMatchObject({ + deletedMessages: 0, + deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 1, + }) + expect(idsForTable(sqlite, 'email_messages')).toEqual(['msg-old-blob']) + + const workingDelete = vi.fn(async () => undefined) + const succeeded = await pruneUserEmailMessagesForRetention({ + db, + blobs: { delete: workingDelete } as unknown as Pick, + now, + }) + expect(succeeded).toMatchObject({ + deletedMessages: 1, + deletedRawMimeBlobs: 1, + }) + expect(idsForTable(sqlite, 'email_messages')).toEqual([]) +}) + +test('entitlement counter and usage rollup retention respect 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 + .prepare( + `INSERT INTO usage_rollups ( + user_id, metric, month, event_count, updated_at + ) VALUES ('user-1', 'mcp_tool_call', ?, 1, ?)`, + ) + .run(month, now.toISOString()) + } + + expect(await pruneEntitlementDailyCountersForRetention({ db, now })).toBe(1) + expect(await pruneUsageRollupsForRetention({ db, now })).toBe(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 }> + expect(months.map((row) => row.month)).toEqual(['2024-07', '2026-06']) +}) + +test('published bundle artifact retention deletes stale rows, KV blobs, and source snapshots', async () => { const { sqlite, db } = createRetentionDb() const kvDelete = vi.fn(async () => undefined) const env = { @@ -623,9 +959,17 @@ test('published bundle artifact retention deletes only stale unreferenced rows a expect(result).toEqual({ deletedRows: 1, deletedKvKeys: 1, + deletedSnapshotKvKeys: 2, kvDeleteErrors: 0, + hasMore: false, }) expect(kvDelete).toHaveBeenCalledWith('kv:artifact-delete') + expect(kvDelete).toHaveBeenCalledWith( + 'source-snapshot:v1:source-stale:commit-old', + ) + expect(kvDelete).toHaveBeenCalledWith( + 'source-manifest-snapshot:v1:source-stale:commit-old', + ) expect(idsForTable(sqlite, 'published_bundle_artifacts')).toEqual([ 'artifact-current', 'artifact-fresh', @@ -709,7 +1053,9 @@ test('published bundle artifact retention rechecks staleness before deleting sel expect(result).toEqual({ deletedRows: 0, deletedKvKeys: 0, + deletedSnapshotKvKeys: 0, kvDeleteErrors: 0, + hasMore: false, }) expect(kvDelete).not.toHaveBeenCalled() expect(idsForTable(sqlite, 'published_bundle_artifacts')).toEqual([ @@ -761,6 +1107,69 @@ test('retention row deletes chunk ids to stay within the D1 binding limit', asyn expect(idsForTable(sqlite, 'package_invocations')).toEqual([]) }) +test('retention run loops batches per table until backlogs are drained', async () => { + const { sqlite, db } = createRetentionDb() + for (let index = 0; index < 501; index += 1) { + insertPackageInvocation(sqlite, { + id: `inv-${String(index).padStart(3, '0')}`, + createdAt: daysAgo(120), + }) + } + for (let index = 0; index < 300; index += 1) { + sqlite + .prepare( + `INSERT INTO audit_events ( + category, action, result, timestamp + ) VALUES ('auth', 'login', 'success', ?)`, + ) + .run(daysAgo(auditEventRetentionDays + 1)) + } + const env = { + APP_DB: db, + BUNDLE_ARTIFACTS_KV: { delete: vi.fn(async () => undefined) }, + } as unknown as Pick + + const result = await pruneRetention({ env, now }) + + expect(result.packageInvocations).toBe(501) + expect(result.auditEvents).toBe(300) + expect(result.batchesPerTable['package_invocations']).toBe(3) + expect(result.batchesPerTable['audit_events']).toBe(2) + expect(result.timeBudgetExhausted).toBe(false) + expect(idsForTable(sqlite, 'package_invocations')).toEqual([]) +}) + +test('retention run gives every table one batch and stops when the budget is exhausted', async () => { + const { sqlite, db } = createRetentionDb() + for (let index = 0; index < 300; index += 1) { + insertPackageInvocation(sqlite, { + id: `inv-${String(index).padStart(3, '0')}`, + createdAt: daysAgo(120), + }) + } + sqlite + .prepare( + `INSERT INTO audit_events ( + category, action, result, timestamp + ) VALUES ('auth', 'login', 'success', ?)`, + ) + .run(daysAgo(auditEventRetentionDays + 1)) + const env = { + APP_DB: db, + BUNDLE_ARTIFACTS_KV: { delete: vi.fn(async () => undefined) }, + } as unknown as Pick + + const result = await pruneRetention({ env, now, timeBudgetMs: 0 }) + + // The first round-robin pass always completes so a hot table cannot starve + // the others, then the exhausted budget stops further passes. + expect(result.packageInvocations).toBe(250) + expect(result.auditEvents).toBe(1) + expect(result.batchesPerTable['package_invocations']).toBe(1) + expect(result.timeBudgetExhausted).toBe(true) + expect(idsForTable(sqlite, 'package_invocations')).toHaveLength(50) +}) + test('retention coverage includes every live growth-pattern table or documented exemption', () => { const db = new DatabaseSync(':memory:') applyMigrations(db) @@ -775,7 +1184,8 @@ test('retention coverage includes every live growth-pattern table or documented .all() as Array<{ name: string }> const candidateTables = new Set() const growthPattern = - /(?:_runs|_logs|_events|_invocations|_suppressions|_artifacts)$/u + /(?:_runs|_logs|_events|_invocations|_suppressions|_artifacts|_messages|_threads|_attachments|_rollups|_counters|_memories)$/u + const growthTimeColumns = ['created_at', 'day', 'month'] for (const table of tables) { const columns = db .prepare(`PRAGMA table_info(${quoteSqlIdentifier(table.name)})`) @@ -783,7 +1193,7 @@ test('retention coverage includes every live growth-pattern table or documented const columnNames = new Set(columns.map((column) => column.name)) const hasUserCreatedGrowthShape = columnNames.has('user_id') && - columnNames.has('created_at') && + growthTimeColumns.some((column) => columnNames.has(column)) && growthPattern.test(table.name) const hasGlobalAuditShape = table.name === 'audit_events' && columnNames.has('timestamp') diff --git a/packages/worker/src/app/retention.ts b/packages/worker/src/app/retention.ts index 9baf25cd0e..51bba96d32 100644 --- a/packages/worker/src/app/retention.ts +++ b/packages/worker/src/app/retention.ts @@ -1,4 +1,8 @@ import { systemEmailOwnerId } from '#worker/email/system-email.ts' +import { + buildPublishedSourceManifestSnapshotKvKey, + buildPublishedSourceSnapshotKvKey, +} from '#worker/package-runtime/published-runtime-artifacts.ts' import { terminalWorkflowStatusValues } from '#worker/package-runtime/workflow-statuses.ts' type RetentionPolicy = { @@ -22,6 +26,19 @@ export const retentionCronIntervalMinutes = 60 export const retentionDefaultBatchSize = 250 export const retentionDeleteIdsMaxParameters = 100 export const publishedBundleArtifactRetentionBatchSize = 100 +/** + * Total wall-clock budget for one retention cron run. Each hourly run keeps + * deleting batches (round-robin across tables) until every table is drained + * or this budget is exhausted, so backlogs shrink instead of only losing one + * batch per table per hour. + */ +export const retentionRunTimeBudgetMs = 20_000 +/** + * Upper bound on distinct (user_id, package_id) pairs examined per runtime-run + * cap batch, largest pairs first, so the cap prune never ranks the whole + * table in one query. + */ +export const packageRuntimeCapPairsPerBatch = 20 export const packageRuntimeRunRetentionDays = 30 export const packageRuntimeMaxRunsPerPackage = 500 @@ -30,6 +47,9 @@ export const memorySuppressionRetentionDays = 90 export const workflowRunRetentionDays = 90 export const publishedBundleArtifactRetentionDays = 30 export const emailDeliveryEventRetentionDays = 90 +export const emailMessageRetentionDays = 365 +export const entitlementDailyCounterRetentionDays = 400 +export const usageRollupRetentionMonths = 24 export const auditEventRetentionDays = 180 const millisecondsPerDay = 24 * 60 * 60 * 1000 @@ -89,7 +109,7 @@ export const retentionPolicies: ReadonlyArray = [ retentionDays: publishedBundleArtifactRetentionDays, batchSize: publishedBundleArtifactRetentionBatchSize, description: - 'Published bundle rows and KV blobs are deleted only when older than 30 days, not current for any source, and not tied to an active repo session.', + 'Published bundle rows, KV blobs, and the matching source snapshot KV keys are deleted only when older than 30 days, not current for any source, and not tied to an active repo session.', }, { table: 'email_delivery_events', @@ -99,6 +119,43 @@ export const retentionPolicies: ReadonlyArray = [ description: 'User email delivery events keep 90 days; operator system email remains governed by system-email retention.', }, + { + table: 'email_messages', + scope: 'per-user', + retentionDays: emailMessageRetentionDays, + batchSize: retentionDefaultBatchSize, + description: + 'User email messages keep 365 days, oldest first. R2 raw MIME blobs are deleted before their rows; rows with blobs are skipped when the EMAIL_BLOBS binding is missing so blobs are never orphaned. Operator system:email messages remain governed by system-email retention.', + }, + { + table: 'email_attachments', + scope: 'per-user', + batchSize: retentionDefaultBatchSize, + description: + 'Attachment rows follow their parent email messages and are deleted before the messages.', + }, + { + table: 'email_threads', + scope: 'per-user', + batchSize: retentionDefaultBatchSize, + 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', + batchSize: retentionDefaultBatchSize, + description: + 'Per user/metric/month usage rollups keep 24 months; Analytics Engine retains the raw event stream separately.', + }, { table: 'audit_events', scope: 'global', @@ -116,6 +173,11 @@ export const retentionPolicyExemptions: ReadonlyArray reason: 'Archived job artifact rows are bounded by retain_until and cleaned by the job artifact cleanup path.', }, + { + table: 'mcp_memories', + reason: + 'Memories are durable user-curated content removed by explicit user action or account deletion, not by time-based retention.', + }, ] as const export type RetentionPruneResult = { @@ -130,10 +192,23 @@ export type RetentionPruneResult = { publishedBundleArtifacts: { deletedRows: number deletedKvKeys: number + deletedSnapshotKvKeys: number kvDeleteErrors: number } emailDeliveryEvents: number + emailMessages: { + deletedMessages: number + deletedAttachments: number + deletedThreads: number + deletedRawMimeBlobs: number + skippedMissingBlobBinding: number + blobDeleteErrors: number + } + entitlementDailyCounters: number + usageRollups: number auditEvents: number + batchesPerTable: Record + timeBudgetExhausted: boolean } export function shouldRunRetentionCron(now: Date) { @@ -158,6 +233,12 @@ function cutoffIso(now: Date, days: number) { return new Date(now.getTime() - days * millisecondsPerDay).toISOString() } +function cutoffMonth(now: Date, months: number) { + return new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - months, 1)) + .toISOString() + .slice(0, 'YYYY-MM'.length) +} + function placeholders(values: ReadonlyArray) { return values.map(() => '?').join(', ') } @@ -234,81 +315,128 @@ async function deletePublishedBundleArtifactRowIfStillStale(input: { return result.meta.changes ?? 0 } +/** + * Keep-conditions shared by the age prune and the per-package cap prune. + * Expects the outer query to alias package_runtime_runs as `run`. + */ +const runtimeRunKeepConditions = `run.status != 'running' + AND NOT EXISTS ( + SELECT 1 + FROM package_invocations AS invocation + WHERE invocation.user_id = run.user_id + AND invocation.id = run.invocation_id + AND invocation.status = 'in_progress' + ) + AND NOT EXISTS ( + SELECT 1 + FROM workflow_runs AS workflow + WHERE workflow.user_id = run.user_id + AND workflow.id = run.workflow_id + AND ( + workflow.status IS NULL + OR workflow.status NOT IN (${terminalWorkflowStatusList}) + ) + ) + AND NOT EXISTS ( + SELECT 1 + FROM repo_sessions AS session + WHERE session.user_id = run.user_id + AND session.id = run.session_id + AND session.status = 'active' + ) + AND NOT EXISTS ( + SELECT 1 + FROM package_runtime_runs AS child_run + WHERE child_run.user_id = run.user_id + AND child_run.parent_run_id = run.id + AND child_run.status = 'running' + )` + export async function prunePackageRuntimeRetention(input: { db: D1Database now?: Date batchSize?: number + maxCapPairs?: number }) { const now = input.now ?? new Date() const cutoff = cutoffIso(now, packageRuntimeRunRetentionDays) const batchSize = input.batchSize ?? retentionDefaultBatchSize - const runIds = await selectIds({ + const maxCapPairs = input.maxCapPairs ?? packageRuntimeCapPairsPerBatch + const ageRunIds = await selectIds({ db: input.db, - bindings: [cutoff, packageRuntimeMaxRunsPerPackage, batchSize], - sql: `WITH ranked_runs AS ( - SELECT - id, - user_id, - started_at, - status, - invocation_id, - workflow_id, - session_id, - ROW_NUMBER() OVER ( - PARTITION BY user_id, package_id - ORDER BY started_at DESC, id DESC - ) AS package_rank - FROM package_runtime_runs - ) - SELECT id - FROM ranked_runs AS run - WHERE (run.started_at < ? OR run.package_rank > ?) - AND run.status != 'running' - AND NOT EXISTS ( - SELECT 1 - FROM package_invocations AS invocation - WHERE invocation.user_id = run.user_id - AND invocation.id = run.invocation_id - AND invocation.status = 'in_progress' - ) - AND NOT EXISTS ( - SELECT 1 - FROM workflow_runs AS workflow - WHERE workflow.user_id = run.user_id - AND workflow.id = run.workflow_id - AND ( - workflow.status IS NULL - OR workflow.status NOT IN (${terminalWorkflowStatusList}) - ) - ) - AND NOT EXISTS ( - SELECT 1 - FROM repo_sessions AS session - WHERE session.user_id = run.user_id - AND session.id = run.session_id - AND session.status = 'active' - ) - AND NOT EXISTS ( - SELECT 1 - FROM package_runtime_runs AS child_run - WHERE child_run.user_id = run.user_id - AND child_run.parent_run_id = run.id - AND child_run.status = 'running' - ) + bindings: [cutoff, batchSize], + sql: `SELECT id + FROM package_runtime_runs AS run + WHERE run.started_at < ? + AND ${runtimeRunKeepConditions} ORDER BY run.started_at ASC, run.id ASC LIMIT ?`, }) + const runIds = new Set(ageRunIds) + // The per-(user, package) cap prune examines a bounded set of the largest + // pairs per batch instead of ranking the whole table with a window + // function; remaining pairs are picked up by later batches or runs. + let capPairCount = 0 + let capDeletedCandidates = 0 + if (runIds.size < batchSize) { + const { results: pairs } = await input.db + .prepare( + `SELECT user_id, package_id, COUNT(*) AS run_count + FROM package_runtime_runs + GROUP BY user_id, package_id + HAVING COUNT(*) > ? + ORDER BY run_count DESC + LIMIT ?`, + ) + .bind(packageRuntimeMaxRunsPerPackage, maxCapPairs) + .all<{ user_id: string; package_id: string; run_count: number }>() + capPairCount = (pairs ?? []).length + for (const pair of pairs ?? []) { + if (runIds.size >= batchSize) break + const capIds = await selectIds({ + db: input.db, + bindings: [ + pair.user_id, + pair.package_id, + pair.user_id, + pair.package_id, + packageRuntimeMaxRunsPerPackage, + batchSize - runIds.size, + ], + sql: `SELECT id + FROM package_runtime_runs AS run + WHERE run.user_id = ? + AND run.package_id = ? + AND run.id NOT IN ( + SELECT recent.id + FROM package_runtime_runs AS recent + WHERE recent.user_id = ? + AND recent.package_id = ? + ORDER BY recent.started_at DESC, recent.id DESC + LIMIT ? + ) + AND ${runtimeRunKeepConditions} + ORDER BY run.started_at ASC, run.id ASC + LIMIT ?`, + }) + for (const id of capIds) { + capDeletedCandidates += 1 + runIds.add(id) + } + } + } + const ids = [...runIds] const deletedLogs = await deleteByIds({ db: input.db, table: 'package_runtime_logs', idColumn: 'run_id', - ids: runIds, + ids, }) const deletedRuns = await deleteByIds({ db: input.db, table: 'package_runtime_runs', idColumn: 'id', - ids: runIds, + ids, }) const orphanLogIds = await selectIds({ db: input.db, @@ -329,7 +457,12 @@ export async function prunePackageRuntimeRetention(input: { idColumn: 'id', ids: orphanLogIds, }) - return { deletedRuns, deletedLogs, deletedOrphanLogs } + const hasMore = + ageRunIds.length >= batchSize || + ids.length >= batchSize || + orphanLogIds.length >= batchSize || + (capDeletedCandidates > 0 && capPairCount >= maxCapPairs) + return { deletedRuns, deletedLogs, deletedOrphanLogs, hasMore } } export async function prunePackageInvocationsForRetention(input: { @@ -422,8 +555,9 @@ export async function prunePublishedBundleArtifactsForRetention(input: { input.now ?? new Date(), publishedBundleArtifactRetentionDays, ) + const batchSize = input.batchSize ?? publishedBundleArtifactRetentionBatchSize const { results } = await input.env.APP_DB.prepare( - `SELECT artifact.id, artifact.kv_key + `SELECT artifact.id, artifact.kv_key, artifact.source_id, artifact.published_commit FROM published_bundle_artifacts AS artifact WHERE artifact.created_at < ? AND NOT EXISTS ( @@ -443,12 +577,22 @@ export async function prunePublishedBundleArtifactsForRetention(input: { ORDER BY artifact.created_at ASC, artifact.id ASC LIMIT ?`, ) - .bind(cutoff, input.batchSize ?? publishedBundleArtifactRetentionBatchSize) - .all<{ id: string; kv_key: string }>() + .bind(cutoff, batchSize) + .all<{ + id: string + kv_key: string + source_id: string + published_commit: string + }>() const rows = results ?? [] let deletedRows = 0 let deletedKvKeys = 0 + let deletedSnapshotKvKeys = 0 let kvDeleteErrors = 0 + const prunedSnapshotPairs = new Map< + string, + { sourceId: string; publishedCommit: string } + >() for (const row of rows) { const rowDeleted = await deletePublishedBundleArtifactRowIfStillStale({ db: input.env.APP_DB, @@ -458,6 +602,10 @@ export async function prunePublishedBundleArtifactsForRetention(input: { }) if (rowDeleted === 0) continue deletedRows += rowDeleted + prunedSnapshotPairs.set(`${row.source_id}:${row.published_commit}`, { + sourceId: row.source_id, + publishedCommit: row.published_commit, + }) try { await input.env.BUNDLE_ARTIFACTS_KV.delete(row.kv_key) deletedKvKeys += 1 @@ -470,7 +618,35 @@ export async function prunePublishedBundleArtifactsForRetention(input: { }) } } - return { deletedRows, deletedKvKeys, kvDeleteErrors } + // A pruned row already satisfied the safety conditions (commit not current + // for any entity_sources row, no active repo session), so the source + // snapshot KV keys for that commit are stale too. + for (const pair of prunedSnapshotPairs.values()) { + for (const kvKey of [ + buildPublishedSourceSnapshotKvKey(pair), + buildPublishedSourceManifestSnapshotKvKey(pair), + ]) { + try { + await input.env.BUNDLE_ARTIFACTS_KV.delete(kvKey) + deletedSnapshotKvKeys += 1 + } catch (error) { + kvDeleteErrors += 1 + console.warn('retention-source-snapshot-kv-delete-failed', { + sourceId: pair.sourceId, + publishedCommit: pair.publishedCommit, + kvKey, + error, + }) + } + } + } + return { + deletedRows, + deletedKvKeys, + deletedSnapshotKvKeys, + kvDeleteErrors, + hasMore: rows.length >= batchSize, + } } export async function pruneEmailDeliveryEventsForRetention(input: { @@ -504,6 +680,171 @@ export async function pruneEmailDeliveryEventsForRetention(input: { }) } +export type EmailMessageRetentionResult = { + deletedMessages: number + deletedAttachments: number + deletedThreads: number + deletedRawMimeBlobs: number + skippedMissingBlobBinding: number + blobDeleteErrors: number + hasMore: boolean +} + +export async function pruneUserEmailMessagesForRetention(input: { + db: D1Database + blobs?: Pick | undefined + now?: Date + batchSize?: number +}): Promise { + const cutoff = cutoffIso(input.now ?? new Date(), emailMessageRetentionDays) + const batchSize = input.batchSize ?? retentionDefaultBatchSize + const { results } = await input.db + .prepare( + `SELECT id, user_id, raw_mime_key + FROM email_messages + WHERE user_id != ? + AND created_at < ? + ORDER BY created_at ASC, id ASC + LIMIT ?`, + ) + .bind(systemEmailOwnerId, cutoff, batchSize) + .all<{ id: string; user_id: string; raw_mime_key: string | null }>() + const rows = results ?? [] + const result: EmailMessageRetentionResult = { + deletedMessages: 0, + deletedAttachments: 0, + deletedThreads: 0, + deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 0, + hasMore: false, + } + const rowsWithBlob = rows.filter((row) => row.raw_mime_key !== null) + let deletableRows = rows + if (rowsWithBlob.length > 0) { + if (!input.blobs) { + // Never delete a row whose R2 blob cannot be deleted first; skipping + // keeps blobs from being orphaned when the binding is missing. + result.skippedMissingBlobBinding = rowsWithBlob.length + deletableRows = rows.filter((row) => row.raw_mime_key === null) + console.warn('retention-email-raw-mime-binding-missing', { + skipped: rowsWithBlob.length, + }) + } else { + try { + await input.blobs.delete( + rowsWithBlob.map((row) => row.raw_mime_key as string), + ) + result.deletedRawMimeBlobs = rowsWithBlob.length + } catch (error) { + result.blobDeleteErrors = rowsWithBlob.length + deletableRows = rows.filter((row) => row.raw_mime_key === null) + console.warn('retention-email-raw-mime-delete-failed', { error }) + } + } + } + const messageIds = deletableRows.map((row) => row.id) + // email_attachments rows reference email_messages via FK, so delete the + // dependent attachment rows first. + result.deletedAttachments = await deleteByIds({ + db: input.db, + table: 'email_attachments', + idColumn: 'message_id', + ids: messageIds, + }) + result.deletedMessages = await deleteByIds({ + db: input.db, + table: 'email_messages', + idColumn: 'id', + ids: messageIds, + }) + const affectedUserIds = [...new Set(deletableRows.map((row) => row.user_id))] + for ( + let index = 0; + index < affectedUserIds.length; + index += retentionDeleteIdsMaxParameters + ) { + const userIds = affectedUserIds.slice( + index, + index + retentionDeleteIdsMaxParameters, + ) + const threadDelete = await input.db + .prepare( + `DELETE FROM email_threads + WHERE user_id IN (${placeholders(userIds)}) + AND NOT EXISTS ( + SELECT 1 + FROM email_messages + WHERE email_messages.thread_id = email_threads.id + )`, + ) + .bind(...userIds) + .run() + result.deletedThreads += threadDelete.meta.changes ?? 0 + } + // Skipped rows sit at the head of the oldest-first queue, so looping again + // in the same run would reselect them; wait for the next run instead. + result.hasMore = + rows.length >= batchSize && + result.skippedMissingBlobBinding === 0 && + result.blobDeleteErrors === 0 + 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) + const rowIds = await selectIds({ + 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 ?`, + }) + return deleteByIds({ + db: input.db, + table: 'entitlement_daily_counters', + idColumn: 'rowid', + ids: rowIds, + }) +} + +export async function pruneUsageRollupsForRetention(input: { + db: D1Database + now?: Date + batchSize?: number +}) { + const oldestKeptMonth = cutoffMonth( + input.now ?? new Date(), + usageRollupRetentionMonths, + ) + const rowIds = await selectIds({ + db: input.db, + column: 'rowid', + bindings: [oldestKeptMonth, input.batchSize ?? retentionDefaultBatchSize], + sql: `SELECT rowid + FROM usage_rollups + WHERE month < ? + ORDER BY month ASC, rowid ASC + LIMIT ?`, + }) + return deleteByIds({ + db: input.db, + table: 'usage_rollups', + idColumn: 'rowid', + ids: rowIds, + }) +} + export async function pruneAuditEventsForRetention(input: { db: D1Database now?: Date @@ -528,39 +869,178 @@ export async function pruneAuditEventsForRetention(input: { } export async function pruneRetention(input: { - env: Pick + // EMAIL_BLOBS stays optional so environments without the R2 binding still + // prune everything else; blob-backed email rows are then skipped. + env: Pick & { + EMAIL_BLOBS?: R2Bucket | undefined + } now?: Date + timeBudgetMs?: number }): Promise { const now = input.now ?? new Date() + const timeBudgetMs = input.timeBudgetMs ?? retentionRunTimeBudgetMs + const startedAtMs = Date.now() + const db = input.env.APP_DB + const emailBlobs = input.env.EMAIL_BLOBS const result: RetentionPruneResult = { - runtimeRuns: await prunePackageRuntimeRetention({ - db: input.env.APP_DB, - now, - }), - packageInvocations: await prunePackageInvocationsForRetention({ - db: input.env.APP_DB, - now, - }), - memorySuppressions: await pruneMemorySuppressionsForRetention({ - db: input.env.APP_DB, - now, - }), - workflowRuns: await pruneWorkflowRunsForRetention({ - db: input.env.APP_DB, - now, - }), - publishedBundleArtifacts: await prunePublishedBundleArtifactsForRetention({ - env: input.env, - now, - }), - emailDeliveryEvents: await pruneEmailDeliveryEventsForRetention({ - db: input.env.APP_DB, - now, - }), - auditEvents: await pruneAuditEventsForRetention({ - db: input.env.APP_DB, - now, - }), + runtimeRuns: { deletedRuns: 0, deletedLogs: 0, deletedOrphanLogs: 0 }, + packageInvocations: 0, + memorySuppressions: 0, + workflowRuns: 0, + publishedBundleArtifacts: { + deletedRows: 0, + deletedKvKeys: 0, + deletedSnapshotKvKeys: 0, + kvDeleteErrors: 0, + }, + emailDeliveryEvents: 0, + emailMessages: { + deletedMessages: 0, + deletedAttachments: 0, + deletedThreads: 0, + deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 0, + }, + entitlementDailyCounters: 0, + usageRollups: 0, + auditEvents: 0, + batchesPerTable: {}, + timeBudgetExhausted: false, + } + const countTask = ( + table: string, + prune: () => Promise, + assign: (count: number) => void, + ) => ({ + table, + done: false, + run: async () => { + const count = await prune() + assign(count) + return count >= retentionDefaultBatchSize + }, + }) + const tasks: Array<{ + table: string + done: boolean + run: () => Promise + }> = [ + { + table: 'package_runtime_runs', + done: false, + run: async () => { + const batch = await prunePackageRuntimeRetention({ db, now }) + result.runtimeRuns.deletedRuns += batch.deletedRuns + result.runtimeRuns.deletedLogs += batch.deletedLogs + result.runtimeRuns.deletedOrphanLogs += batch.deletedOrphanLogs + return batch.hasMore + }, + }, + countTask( + 'package_invocations', + () => prunePackageInvocationsForRetention({ db, now }), + (count) => { + result.packageInvocations += count + }, + ), + countTask( + 'mcp_memory_conversation_suppressions', + () => pruneMemorySuppressionsForRetention({ db, now }), + (count) => { + result.memorySuppressions += count + }, + ), + countTask( + 'workflow_runs', + () => pruneWorkflowRunsForRetention({ db, now }), + (count) => { + result.workflowRuns += count + }, + ), + { + table: 'published_bundle_artifacts', + done: false, + run: async () => { + const batch = await prunePublishedBundleArtifactsForRetention({ + env: input.env, + now, + }) + result.publishedBundleArtifacts.deletedRows += batch.deletedRows + result.publishedBundleArtifacts.deletedKvKeys += batch.deletedKvKeys + result.publishedBundleArtifacts.deletedSnapshotKvKeys += + batch.deletedSnapshotKvKeys + result.publishedBundleArtifacts.kvDeleteErrors += batch.kvDeleteErrors + return batch.hasMore + }, + }, + countTask( + 'email_delivery_events', + () => pruneEmailDeliveryEventsForRetention({ db, now }), + (count) => { + result.emailDeliveryEvents += count + }, + ), + { + table: 'email_messages', + done: false, + run: async () => { + const batch = await pruneUserEmailMessagesForRetention({ + db, + blobs: emailBlobs, + now, + }) + result.emailMessages.deletedMessages += batch.deletedMessages + result.emailMessages.deletedAttachments += batch.deletedAttachments + result.emailMessages.deletedThreads += batch.deletedThreads + result.emailMessages.deletedRawMimeBlobs += batch.deletedRawMimeBlobs + result.emailMessages.skippedMissingBlobBinding += + batch.skippedMissingBlobBinding + result.emailMessages.blobDeleteErrors += batch.blobDeleteErrors + return batch.hasMore + }, + }, + countTask( + 'entitlement_daily_counters', + () => pruneEntitlementDailyCountersForRetention({ db, now }), + (count) => { + result.entitlementDailyCounters += count + }, + ), + countTask( + 'usage_rollups', + () => pruneUsageRollupsForRetention({ db, now }), + (count) => { + result.usageRollups += count + }, + ), + countTask( + 'audit_events', + () => pruneAuditEventsForRetention({ db, now }), + (count) => { + result.auditEvents += count + }, + ), + ] + // Round-robin passes: every pending table gets one batch before any table + // gets a second one, so a hot table cannot starve the others. The first + // pass always completes so each table makes progress every run; later + // passes stop once the time budget is exhausted. + let pass = 0 + while (tasks.some((task) => !task.done)) { + if (pass > 0 && Date.now() - startedAtMs >= timeBudgetMs) { + result.timeBudgetExhausted = true + break + } + for (const task of tasks) { + if (task.done) continue + if (pass > 0 && Date.now() - startedAtMs >= timeBudgetMs) break + const hasMore = await task.run() + result.batchesPerTable[task.table] = + (result.batchesPerTable[task.table] ?? 0) + 1 + if (!hasMore) task.done = true + } + pass += 1 } console.info('retention-prune', JSON.stringify(result)) return result diff --git a/packages/worker/src/email/system-email.ts b/packages/worker/src/email/system-email.ts index 2d3c96385b..0c4cac7484 100644 --- a/packages/worker/src/email/system-email.ts +++ b/packages/worker/src/email/system-email.ts @@ -3,7 +3,6 @@ import { createEmailInbox, createEmailInboxAddress, deleteEmailInboxAddressById, - deleteEmailMessageById, getEmailInboxAddressByAddress, getEmailInboxById, getEmailInboxByName, @@ -32,6 +31,16 @@ export const systemEmailLimits = { pruneBatchSize: 500, } as const +/** + * Wall-clock budget for one system-email retention run. Deletes keep looping + * in pruneBatchSize batches until the backlog is drained or the budget is + * exhausted, so a large backlog never turns into one unbounded DELETE. + */ +export const systemEmailPruneTimeBudgetMs = 10_000 + +/** D1 caps bound parameters per statement, so IN (...) deletes are chunked. */ +const systemEmailDeleteIdsMaxParameters = 100 + const systemEmailLocalSet = new Set(systemEmailLocals) export type ProvisionedSystemEmailInbox = { @@ -203,28 +212,66 @@ async function listSystemEmailMessageIds(input: { async function deleteSystemEmailMessagesByIds(input: { db: D1Database + blobs?: R2Bucket | null messageIds: ReadonlyArray }) { - if (input.messageIds.length === 0) return 0 - const placeholders = input.messageIds.map(() => '?').join(', ') - await input.db - .prepare( - `DELETE FROM email_attachments - WHERE message_id IN (${placeholders})`, - ) - .bind(...input.messageIds) - .run() - await input.db - .prepare( - `DELETE FROM email_delivery_events - WHERE message_id IN (${placeholders})`, + let deletedRawMimeBlobs = 0 + for ( + let index = 0; + index < input.messageIds.length; + index += systemEmailDeleteIdsMaxParameters + ) { + const messageIds = input.messageIds.slice( + index, + index + systemEmailDeleteIdsMaxParameters, ) - .bind(...input.messageIds) - .run() - for (const messageId of input.messageIds) { - await deleteEmailMessageById({ db: input.db, messageId }) + const placeholders = messageIds.map(() => '?').join(', ') + // Capture blob keys before the rows disappear. Blob cleanup is + // best-effort and must never block the D1 delete. + let rawMimeKeys: Array = [] + if (input.blobs) { + const { results } = await input.db + .prepare( + `SELECT raw_mime_key + FROM email_messages + WHERE id IN (${placeholders}) + AND raw_mime_key IS NOT NULL`, + ) + .bind(...messageIds) + .all<{ raw_mime_key: string }>() + rawMimeKeys = (results ?? []).map((row) => row.raw_mime_key) + } + await input.db + .prepare( + `DELETE FROM email_attachments + WHERE message_id IN (${placeholders})`, + ) + .bind(...messageIds) + .run() + await input.db + .prepare( + `DELETE FROM email_delivery_events + WHERE message_id IN (${placeholders})`, + ) + .bind(...messageIds) + .run() + await input.db + .prepare( + `DELETE FROM email_messages + WHERE id IN (${placeholders})`, + ) + .bind(...messageIds) + .run() + if (input.blobs && rawMimeKeys.length > 0) { + try { + await input.blobs.delete(rawMimeKeys) + deletedRawMimeBlobs += rawMimeKeys.length + } catch (error) { + console.warn('system-email-raw-mime-blob-delete-failed', error) + } + } } - return input.messageIds.length + return { deletedMessages: input.messageIds.length, deletedRawMimeBlobs } } export type SystemEmailRetentionResult = { @@ -232,13 +279,19 @@ export type SystemEmailRetentionResult = { deletedDeliveryEvents: number deletedCounters: number deletedThreads: number + deletedRawMimeBlobs: number } export async function pruneSystemEmailRetention(input: { db: D1Database + /** EMAIL_BLOBS bucket; when present raw-MIME blobs are also deleted. */ + blobs?: R2Bucket | null now?: Date + timeBudgetMs?: number }) { const now = input.now ?? new Date() + const timeBudgetMs = input.timeBudgetMs ?? systemEmailPruneTimeBudgetMs + const deadlineMs = Date.now() + timeBudgetMs const cutoff = new Date( now.getTime() - systemEmailLimits.retentionDays * 24 * 60 * 60 * 1000, ) @@ -249,6 +302,7 @@ export async function pruneSystemEmailRetention(input: { deletedDeliveryEvents: 0, deletedCounters: 0, deletedThreads: 0, + deletedRawMimeBlobs: 0, } const oldMessageIds = await listSystemEmailMessageIds({ @@ -256,30 +310,49 @@ export async function pruneSystemEmailRetention(input: { before: cutoffIso, limit: systemEmailLimits.pruneBatchSize, }) - result.deletedMessages += await deleteSystemEmailMessagesByIds({ + const oldDelete = await deleteSystemEmailMessagesByIds({ db: input.db, + blobs: input.blobs, messageIds: oldMessageIds, }) + result.deletedMessages += oldDelete.deletedMessages + result.deletedRawMimeBlobs += oldDelete.deletedRawMimeBlobs const overCapMessageIds = await listSystemEmailMessageIds({ db: input.db, offset: systemEmailLimits.maxStoredMessages, limit: systemEmailLimits.pruneBatchSize, }) - result.deletedMessages += await deleteSystemEmailMessagesByIds({ + const overCapDelete = await deleteSystemEmailMessagesByIds({ db: input.db, + blobs: input.blobs, messageIds: overCapMessageIds, }) + result.deletedMessages += overCapDelete.deletedMessages + result.deletedRawMimeBlobs += overCapDelete.deletedRawMimeBlobs - const eventDelete = await input.db - .prepare( - `DELETE FROM email_delivery_events - WHERE user_id = ? - AND created_at < ?`, - ) - .bind(systemEmailOwnerId, cutoffIso) - .run() - result.deletedDeliveryEvents = eventDelete.meta.changes ?? 0 + let deletedEventsInBatch = 0 + do { + const eventDelete = await input.db + .prepare( + `DELETE FROM email_delivery_events + WHERE id IN ( + SELECT id + FROM email_delivery_events + WHERE user_id = ? + AND created_at < ? + ORDER BY created_at ASC, id ASC + LIMIT ? + )`, + ) + .bind(systemEmailOwnerId, cutoffIso, systemEmailLimits.pruneBatchSize) + .run() + deletedEventsInBatch = eventDelete.meta.changes ?? 0 + result.deletedDeliveryEvents += deletedEventsInBatch + } while ( + deletedEventsInBatch >= systemEmailLimits.pruneBatchSize && + Date.now() < deadlineMs + ) const counterDelete = await input.db .prepare(`DELETE FROM system_email_daily_counters WHERE day < ?`) diff --git a/packages/worker/src/email/system-email.workers.test.ts b/packages/worker/src/email/system-email.workers.test.ts index 1b4fb165d5..d77ea9739f 100644 --- a/packages/worker/src/email/system-email.workers.test.ts +++ b/packages/worker/src/email/system-email.workers.test.ts @@ -260,14 +260,24 @@ test('system email retention prunes old operator-owned messages and counters', a now.getTime() - (systemEmailLimits.retentionDays + 1) * 24 * 60 * 60 * 1000, ).toISOString() const fresh = now.toISOString() + const oldRawMimeKey = `email-raw:v1:${systemEmailOwnerId}/old-system-message` + await env.EMAIL_BLOBS.put(oldRawMimeKey, 'old raw mime') await env.APP_DB.prepare( `INSERT INTO email_messages ( - id, direction, user_id, from_address, subject, processing_status, created_at, updated_at + id, direction, user_id, from_address, subject, processing_status, raw_mime_key, created_at, updated_at ) VALUES - ('old-system-message', 'inbound', ?, 'old@example.net', 'Old', 'stored', ?, ?), - ('fresh-system-message', 'inbound', ?, 'fresh@example.net', 'Fresh', 'stored', ?, ?)`, + ('old-system-message', 'inbound', ?, 'old@example.net', 'Old', 'stored', ?, ?, ?), + ('fresh-system-message', 'inbound', ?, 'fresh@example.net', 'Fresh', 'stored', NULL, ?, ?)`, ) - .bind(systemEmailOwnerId, old, old, systemEmailOwnerId, fresh, fresh) + .bind( + systemEmailOwnerId, + oldRawMimeKey, + old, + old, + systemEmailOwnerId, + fresh, + fresh, + ) .run() await env.APP_DB.prepare( `INSERT INTO email_attachments ( @@ -292,10 +302,16 @@ test('system email retention prunes old operator-owned messages and counters', a .bind(old) .run() - const result = await pruneSystemEmailRetention({ db: env.APP_DB, now }) + const result = await pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + now, + }) expect(result.deletedMessages).toBe(1) expect(result.deletedCounters).toBe(1) + expect(result.deletedRawMimeBlobs).toBe(1) + expect(await env.EMAIL_BLOBS.get(oldRawMimeKey)).toBeNull() expect( await listEmailMessages({ db: env.APP_DB, @@ -319,3 +335,43 @@ test('system email retention prunes old operator-owned messages and counters', a ).all(), ).toMatchObject({ results: [{ id: 'fresh-event' }] }) }) + +test('system email retention drains delivery-event backlogs larger than one batch', async () => { + await ensureEmailTestSchema(env.APP_DB) + const now = new Date('2026-07-06T12:00:00.000Z') + const old = new Date( + now.getTime() - (systemEmailLimits.retentionDays + 1) * 24 * 60 * 60 * 1000, + ).toISOString() + const backlog = systemEmailLimits.pruneBatchSize + 5 + // 3 bindings per row must stay under D1's 100-variable statement limit. + const insertChunkSize = 30 + for (let start = 0; start < backlog; start += insertChunkSize) { + const count = Math.min(insertChunkSize, backlog - start) + const values = Array.from({ length: count }) + .map(() => `(?, NULL, ?, NULL, 'received', 'test', '{}', ?)`) + .join(', ') + const bindings = Array.from({ length: count }).flatMap((_, index) => [ + `backlog-event-${start + index}`, + systemEmailOwnerId, + old, + ]) + await env.APP_DB.prepare( + `INSERT INTO email_delivery_events ( + id, message_id, user_id, inbox_id, event_type, provider, detail_json, created_at + ) VALUES ${values}`, + ) + .bind(...bindings) + .run() + } + + const result = await pruneSystemEmailRetention({ db: env.APP_DB, now }) + + expect(result.deletedDeliveryEvents).toBeGreaterThanOrEqual(backlog) + const remaining = await env.APP_DB.prepare( + `SELECT COUNT(*) AS count FROM email_delivery_events + WHERE user_id = ? AND id LIKE 'backlog-event-%'`, + ) + .bind(systemEmailOwnerId) + .first<{ count: number }>() + expect(Number(remaining?.count ?? -1)).toBe(0) +}) diff --git a/packages/worker/src/package-retrievers/manifest-cache.node.test.ts b/packages/worker/src/package-retrievers/manifest-cache.node.test.ts index 5ad53a6638..9c1b812d4b 100644 --- a/packages/worker/src/package-retrievers/manifest-cache.node.test.ts +++ b/packages/worker/src/package-retrievers/manifest-cache.node.test.ts @@ -113,6 +113,23 @@ test('package retriever manifest cache refreshes, lists, removes entries, and pr const env = { BUNDLE_ARTIFACTS_KV: kv } as Env const manifest = createRetrieverManifest() + // A pre-existing legacy combined index blob is deleted on refresh; only + // per-entry keys are written going forward. + const legacySearchIndexKey = buildPackageRetrieverScopeIndexKey({ + userId: 'user-1', + scope: 'search', + }) + store.set( + legacySearchIndexKey, + JSON.stringify({ + version: 1, + userId: 'user-1', + scope: 'search', + retrievers: [], + updatedAt: '2026-04-20T00:00:00.000Z', + }), + ) + await refreshPackageRetrieverManifestCache({ env, userId: 'user-1', @@ -121,21 +138,12 @@ test('package retriever manifest cache refreshes, lists, removes entries, and pr manifest, }) - const searchIndex = JSON.parse( - store.get( - buildPackageRetrieverScopeIndexKey({ - userId: 'user-1', - scope: 'search', - }), - ) ?? '{}', - ) - expect(searchIndex.retrievers).toEqual([ - expect.objectContaining({ - packageId: 'package-1', - retrieverKey: 'notes-search', - revision: 'commit-1', - }), - ]) + expect(store.has(legacySearchIndexKey)).toBe(false) + expect( + store.has( + 'package-retriever-index-entry:v1:user-1:search:package-1:notes-search', + ), + ).toBe(true) expect(kv.put).toHaveBeenCalledWith( expect.stringContaining( 'package-retriever-manifest:v1:user-1:package-1:commit-1', @@ -213,17 +221,122 @@ test('package retriever manifest cache refreshes, lists, removes entries, and pr key.startsWith('package-retriever-manifest:v1:user-1:package-1:'), ), ).toEqual([]) - const clearedSearchIndex = JSON.parse( - store.get( - buildPackageRetrieverScopeIndexKey({ - userId: 'user-1', - scope: 'search', - }), - ) ?? '{}', + expect(store.has(legacySearchIndexKey)).toBe(false) +}) + +test('refreshing to a new revision deletes the stale manifest cache key', async () => { + const { kv, store } = createKv() + const env = { BUNDLE_ARTIFACTS_KV: kv } as Env + const manifest = createRetrieverManifest() + + await refreshPackageRetrieverManifestCache({ + env, + userId: 'user-1', + source: createSource({ published_commit: 'commit-1' }), + savedPackage: createSavedPackage(), + manifest, + }) + await refreshPackageRetrieverManifestCache({ + env, + userId: 'user-1', + source: createSource({ published_commit: 'commit-2' }), + savedPackage: createSavedPackage(), + manifest, + }) + + expect( + Array.from(store.keys()).filter((key) => + key.startsWith('package-retriever-manifest:v1:user-1:package-1:'), + ), + ).toEqual(['package-retriever-manifest:v1:user-1:package-1:commit-2']) + await expect( + listPackageRetrieversForScope({ + env, + userId: 'user-1', + scope: 'search', + }), + ).resolves.toEqual([ + expect.objectContaining({ + packageId: 'package-1', + revision: 'commit-2', + }), + ]) +}) + +test('readers without kv.list support still read the legacy index but nothing writes it', async () => { + const { kv, store } = createKv({ list: false }) + const env = { BUNDLE_ARTIFACTS_KV: kv } as Env + const legacyKey = buildPackageRetrieverScopeIndexKey({ + userId: 'user-1', + scope: 'search', + }) + store.set( + legacyKey, + JSON.stringify({ + version: 1, + userId: 'user-1', + scope: 'search', + retrievers: [ + { + userId: 'user-1', + packageId: 'package-legacy', + kodyId: 'legacy-inbox', + packageName: '@kentcdodds/legacy-inbox', + sourceId: 'source-legacy', + revision: 'commit-legacy', + retrieverKey: 'notes-search', + name: 'Notes Search', + description: 'Searches saved notes', + scopes: ['search'], + }, + ], + updatedAt: '2026-04-20T00:00:00.000Z', + }), ) - expect(clearedSearchIndex.retrievers).toEqual([ - expect.objectContaining({ packageId: 'package-2' }), + store.set( + 'package-retriever-manifest:v1:user-1:package-legacy:commit-legacy', + JSON.stringify({ + version: 1, + userId: 'user-1', + packageId: 'package-legacy', + kodyId: 'legacy-inbox', + packageName: '@kentcdodds/legacy-inbox', + sourceId: 'source-legacy', + revision: 'commit-legacy', + manifestHash: 'hash', + retrievers: [ + { + userId: 'user-1', + packageId: 'package-legacy', + kodyId: 'legacy-inbox', + packageName: '@kentcdodds/legacy-inbox', + sourceId: 'source-legacy', + revision: 'commit-legacy', + retrieverKey: 'notes-search', + exportName: './search-notes', + entryPoint: 'src/search-notes.ts', + name: 'Notes Search', + description: 'Searches saved notes', + scopes: ['search'], + }, + ], + cachedAt: '2026-04-20T00:00:00.000Z', + }), + ) + + await expect( + listPackageRetrieversForScope({ + env, + userId: 'user-1', + scope: 'search', + }), + ).resolves.toEqual([ + expect.objectContaining({ + packageId: 'package-legacy', + retrieverKey: 'notes-search', + }), ]) + expect(kv.put).not.toHaveBeenCalledWith(legacyKey, expect.any(String)) }) test('listPackageRetrieversForScope filters stale, malformed, and prefix-colliding cache rows', async () => { diff --git a/packages/worker/src/package-retrievers/manifest-cache.ts b/packages/worker/src/package-retrievers/manifest-cache.ts index 0f45b2206e..1ca3b903ea 100644 --- a/packages/worker/src/package-retrievers/manifest-cache.ts +++ b/packages/worker/src/package-retrievers/manifest-cache.ts @@ -225,6 +225,12 @@ async function listManifestCacheKeys(input: { return keys } +/** + * The combined per-user scope index blob is legacy: it is no longer written + * (per-entry keys are the source of truth) and is deleted best-effort on + * refresh/removal. Reading it remains only as the fallback for environments + * whose KV binding does not support list(). + */ async function readLegacyScopeIndex(input: { env: Env userId: string @@ -377,15 +383,16 @@ export async function refreshPackageRetrieverManifestCache(input: { manifestKey, JSON.stringify(manifestCache), ) + await deleteStaleManifestCacheKeys({ + env: input.env, + userId: input.userId, + packageId: input.savedPackage.id, + currentManifestKey: manifestKey, + }) for (const scope of [ 'search', 'context', ] satisfies Array) { - const existingEntries = await readScopeEntries({ - env: input.env, - userId: input.userId, - scope, - }) const existingScopeEntryKeys = await listScopeEntryKeys({ env: input.env, userId: input.userId, @@ -395,9 +402,6 @@ export async function refreshPackageRetrieverManifestCache(input: { const nextEntries = registeredRetrievers .filter((retriever) => retriever.scopes.includes(scope)) .map(toScopeIndexEntry) - const preservedEntries = existingEntries.filter( - (entry) => entry.packageId !== input.savedPackage.id, - ) const nextEntryKeys = new Set( nextEntries.map((entry) => buildPackageRetrieverScopeEntryKey({ @@ -422,16 +426,58 @@ export async function refreshPackageRetrieverManifestCache(input: { }), ), ]) - await writeLegacyScopeIndex({ + await deleteLegacyScopeIndex({ env: input.env, userId: input.userId, scope, - entries: [...preservedEntries, ...nextEntries], }) } return manifestCache } +async function deleteStaleManifestCacheKeys(input: { + env: Env + userId: string + packageId: string + currentManifestKey: string +}) { + try { + const manifestCacheKeys = await listManifestCacheKeys(input) + await Promise.all( + manifestCacheKeys + .filter((key) => key !== input.currentManifestKey) + .map(async (key) => await getRetrieverKv(input.env).delete(key)), + ) + } catch (error) { + console.warn('package-retriever-stale-manifest-cleanup-failed', { + userId: input.userId, + packageId: input.packageId, + error, + }) + } +} + +async function deleteLegacyScopeIndex(input: { + env: Env + userId: string + scope: PackageRetrieverScope +}) { + try { + await getRetrieverKv(input.env).delete( + buildPackageRetrieverScopeIndexKey({ + userId: input.userId, + scope: input.scope, + }), + ) + } catch (error) { + console.warn('package-retriever-legacy-index-delete-failed', { + userId: input.userId, + scope: input.scope, + error, + }) + } +} + export async function removePackageRetrieverManifestCacheEntries(input: { env: Env userId: string @@ -443,11 +489,6 @@ export async function removePackageRetrieverManifestCacheEntries(input: { 'search', 'context', ] satisfies Array) { - const existingEntries = await readScopeEntries({ - env: input.env, - userId: input.userId, - scope, - }) const existingScopeEntryKeys = await listScopeEntryKeys({ env: input.env, userId: input.userId, @@ -459,13 +500,10 @@ export async function removePackageRetrieverManifestCacheEntries(input: { async (key) => await getRetrieverKv(input.env).delete(key), ), ) - await writeLegacyScopeIndex({ + await deleteLegacyScopeIndex({ env: input.env, userId: input.userId, scope, - entries: existingEntries.filter( - (entry) => entry.packageId !== input.packageId, - ), }) } await Promise.all( @@ -520,34 +558,6 @@ export async function deleteAllPackageRetrieverCacheEntriesForUser(input: { return keys.size } -async function writeLegacyScopeIndex(input: { - env: Env - userId: string - scope: PackageRetrieverScope - entries?: Array -}) { - const entries = - input.entries ?? - (await readScopeEntries({ - env: input.env, - userId: input.userId, - scope: input.scope, - })) - await getRetrieverKv(input.env).put( - buildPackageRetrieverScopeIndexKey({ - userId: input.userId, - scope: input.scope, - }), - JSON.stringify({ - version: retrieverScopeIndexVersion, - userId: input.userId, - scope: input.scope, - retrievers: entries, - updatedAt: new Date().toISOString(), - } satisfies PackageRetrieverScopeIndex), - ) -} - export async function listPackageRetrieversForScope(input: { env: Env userId: string From 41855d23cd9dc04eabc62041d7c473c399308f8d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:18 +0000 Subject: [PATCH 04/14] Bound per-invocation Durable Object work Cap due jobs at 25 per JobManager alarm with a 1s backlog re-arm, cap runRepoChecks source walks at 2000 files / 15MiB with an actionable check failure, and give package service auto-start retries exponential backoff with jitter (5s doubling to 15min, persisted across eviction). Co-authored-by: Kent C. Dodds --- .../worker/src/jobs/manager-do.node.test.ts | 233 +++++++++++++----- packages/worker/src/jobs/manager-do.ts | 20 +- packages/worker/src/jobs/repo.ts | 30 ++- packages/worker/src/jobs/repo.workers.test.ts | 132 ++++++++++ .../package-service.node.test.ts | 188 +++++++++++++- .../src/package-runtime/package-service.ts | 53 +++- packages/worker/src/repo/checks.node.test.ts | 71 +++++- packages/worker/src/repo/checks.ts | 48 +++- 8 files changed, 698 insertions(+), 77 deletions(-) create mode 100644 packages/worker/src/jobs/repo.workers.test.ts diff --git a/packages/worker/src/jobs/manager-do.node.test.ts b/packages/worker/src/jobs/manager-do.node.test.ts index f39be17461..7b78351ba2 100644 --- a/packages/worker/src/jobs/manager-do.node.test.ts +++ b/packages/worker/src/jobs/manager-do.node.test.ts @@ -1,6 +1,7 @@ import type * as SchedulerLoggingType from './scheduler-logging.ts' import { expect, test, vi } from 'vitest' import { resolveJobManagerAlarmState } from './manager-state.ts' +import { maxDueJobsPerAlarm } from './repo.ts' const mockModule = vi.hoisted(() => ({ getNextRunnableJob: vi.fn(), @@ -136,81 +137,184 @@ test('resolveJobManagerAlarmState covers armed, out-of-sync, and idle states', ( test('syncAlarm arms, clears, and logs scheduler state transitions', async () => { resetMocks() - const nextRunAt = '2026-04-20T18:30:00.000Z' - mockModule.getNextRunnableJob.mockResolvedValue({ - id: 'job-123', - nextRunAt, - }) - const armedState = createState({ - currentAlarmAt: Date.parse('2026-04-20T18:00:00.000Z'), - }) - const armedManager = new JobManagerBase(armedState.state, {} as Env) + vi.useFakeTimers() + try { + vi.setSystemTime(new Date('2026-04-20T17:00:00.000Z')) + const nextRunAt = '2026-04-20T18:30:00.000Z' + mockModule.getNextRunnableJob.mockResolvedValue({ + id: 'job-123', + nextRunAt, + }) + const armedState = createState({ + currentAlarmAt: Date.parse('2026-04-20T18:00:00.000Z'), + }) + const armedManager = new JobManagerBase(armedState.state, {} as Env) - await expect(armedManager.syncAlarm({ userId: 'user-123' })).resolves.toEqual( - { + await expect( + armedManager.syncAlarm({ userId: 'user-123' }), + ).resolves.toEqual({ ok: true, userId: 'user-123', nextRunAt, - }, - ) + }) - expect(armedState.persistedEntries.get('user-id')).toBe('user-123') - expect(armedState.getAlarmAt()).toBe(Date.parse(nextRunAt)) - expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith({ - event: 'sync_alarm', - userId: 'user-123', - currentAlarmAt: '2026-04-20T18:00:00.000Z', - nextJobId: 'job-123', - nextRunAt, - reason: 'alarm_armed', - }) - expect(mockModule.logJobSchedulerError).not.toHaveBeenCalled() + expect(armedState.persistedEntries.get('user-id')).toBe('user-123') + expect(armedState.getAlarmAt()).toBe(Date.parse(nextRunAt)) + expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith({ + event: 'sync_alarm', + userId: 'user-123', + currentAlarmAt: '2026-04-20T18:00:00.000Z', + nextJobId: 'job-123', + nextRunAt, + reason: 'alarm_armed', + }) + expect(mockModule.logJobSchedulerError).not.toHaveBeenCalled() + resetMocks() + mockModule.getNextRunnableJob.mockResolvedValue(null) + const clearedState = createState({ + currentAlarmAt: Date.parse('2026-04-20T18:00:00.000Z'), + }) + const clearedManager = new JobManagerBase(clearedState.state, {} as Env) + + await expect( + clearedManager.syncAlarm({ userId: 'user-123' }), + ).resolves.toEqual({ + ok: true, + userId: 'user-123', + nextRunAt: null, + }) + + expect(clearedState.getAlarmAt()).toBeNull() + expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith({ + event: 'sync_alarm', + userId: 'user-123', + currentAlarmAt: '2026-04-20T18:00:00.000Z', + nextJobId: null, + nextRunAt: null, + reason: 'no_runnable_job', + }) + + resetMocks() + mockModule.getNextRunnableJob.mockRejectedValue( + new Error('next job lookup failed'), + ) + const lookupFailureState = createState() + const lookupFailureManager = new JobManagerBase( + lookupFailureState.state, + {} as Env, + ) + + await expect( + lookupFailureManager.syncAlarm({ userId: 'user-123', source: 'alarm' }), + ).rejects.toThrow('next job lookup failed') + expect(mockModule.logJobSchedulerError).toHaveBeenCalledWith({ + event: 'sync_alarm_failed', + userId: 'user-123', + source: 'alarm', + errorName: 'Error', + errorMessage: 'next job lookup failed', + }) + } finally { + vi.useRealTimers() + } +}) + +test('syncAlarm arms a near-immediate follow-up alarm when a due-job backlog remains', async () => { resetMocks() - mockModule.getNextRunnableJob.mockResolvedValue(null) - const clearedState = createState({ - currentAlarmAt: Date.parse('2026-04-20T18:00:00.000Z'), - }) - const clearedManager = new JobManagerBase(clearedState.state, {} as Env) + vi.useFakeTimers() + try { + vi.setSystemTime(new Date('2026-04-20T19:30:00.000Z')) + const overdueRunAt = '2026-04-20T18:30:00.000Z' + mockModule.getNextRunnableJob.mockResolvedValue({ + id: 'job-backlog', + nextRunAt: overdueRunAt, + }) + const backlogState = createState() + const backlogManager = new JobManagerBase(backlogState.state, {} as Env) - await expect( - clearedManager.syncAlarm({ userId: 'user-123' }), - ).resolves.toEqual({ - ok: true, - userId: 'user-123', - nextRunAt: null, - }) + await expect( + backlogManager.syncAlarm({ userId: 'user-123', source: 'alarm' }), + ).resolves.toEqual({ + ok: true, + userId: 'user-123', + nextRunAt: overdueRunAt, + }) - expect(clearedState.getAlarmAt()).toBeNull() - expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith({ - event: 'sync_alarm', - userId: 'user-123', - currentAlarmAt: '2026-04-20T18:00:00.000Z', - nextJobId: null, - nextRunAt: null, - reason: 'no_runnable_job', - }) + expect(backlogState.getAlarmAt()).toBe( + Date.parse('2026-04-20T19:30:01.000Z'), + ) + expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith({ + event: 'sync_alarm', + userId: 'user-123', + currentAlarmAt: null, + nextJobId: 'job-backlog', + nextRunAt: overdueRunAt, + reason: 'backlog_rearmed', + }) + } finally { + vi.useRealTimers() + } +}) +test('due-job backlogs drain across successive alarm invocations', async () => { resetMocks() - mockModule.getNextRunnableJob.mockRejectedValue( - new Error('next job lookup failed'), - ) - const lookupFailureState = createState() - const lookupFailureManager = new JobManagerBase( - lookupFailureState.state, - {} as Env, - ) + vi.useFakeTimers() + try { + vi.setSystemTime(new Date('2026-04-20T19:30:00.000Z')) + const overdueRunAt = '2026-04-20T18:00:00.000Z' + const futureRunAt = '2026-04-21T09:00:00.000Z' + // First alarm processes a full batch (the per-alarm limit) and leaves an + // overdue job behind; the second alarm drains the remainder. + mockModule.runDueJobsForUser + .mockResolvedValueOnce({ + dueJobCount: maxDueJobsPerAlarm, + successCount: maxDueJobsPerAlarm, + errorCount: 0, + jobOutcomes: [], + }) + .mockResolvedValueOnce({ + dueJobCount: 5, + successCount: 5, + errorCount: 0, + jobOutcomes: [], + }) + mockModule.getNextRunnableJob + .mockResolvedValueOnce({ + id: 'job-still-due', + nextRunAt: overdueRunAt, + }) + .mockResolvedValueOnce({ + id: 'job-future', + nextRunAt: futureRunAt, + }) + const { state, getAlarmAt } = createState() + const manager = new JobManagerBase(state, {} as Env) - await expect( - lookupFailureManager.syncAlarm({ userId: 'user-123', source: 'alarm' }), - ).rejects.toThrow('next job lookup failed') - expect(mockModule.logJobSchedulerError).toHaveBeenCalledWith({ - event: 'sync_alarm_failed', - userId: 'user-123', - source: 'alarm', - errorName: 'Error', - errorMessage: 'next job lookup failed', - }) + await expect(manager.alarm()).resolves.toBeUndefined() + expect(getAlarmAt()).toBe(Date.parse('2026-04-20T19:30:01.000Z')) + + vi.setSystemTime(new Date('2026-04-20T19:30:01.000Z')) + await expect(manager.alarm()).resolves.toBeUndefined() + expect(getAlarmAt()).toBe(Date.parse(futureRunAt)) + + expect(mockModule.runDueJobsForUser).toHaveBeenCalledTimes(2) + expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith( + expect.objectContaining({ + event: 'sync_alarm', + reason: 'backlog_rearmed', + }), + ) + expect(mockModule.logJobSchedulerEvent).toHaveBeenCalledWith( + expect.objectContaining({ + event: 'sync_alarm', + nextJobId: 'job-future', + reason: 'alarm_armed', + }), + ) + } finally { + vi.useRealTimers() + } }) test('exportUser returns the scheduler state for account export', async () => { @@ -238,6 +342,8 @@ test('exportUser returns the scheduler state for account export', async () => { test('alarm logs firing, due-job outcomes, and resyncs the next alarm', async () => { resetMocks() + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-04-20T18:30:00.000Z')) mockModule.runDueJobsForUser.mockResolvedValue({ dueJobCount: 2, successCount: 1, @@ -355,4 +461,5 @@ test('alarm logs firing, due-job outcomes, and resyncs the next alarm', async () errorName: 'Error', errorMessage: 'run due jobs failed', }) + vi.useRealTimers() }) diff --git a/packages/worker/src/jobs/manager-do.ts b/packages/worker/src/jobs/manager-do.ts index c285a1d4e1..3b62fe3ab9 100644 --- a/packages/worker/src/jobs/manager-do.ts +++ b/packages/worker/src/jobs/manager-do.ts @@ -15,6 +15,14 @@ import { type JobManagerDebugState } from './manager-client.ts' const userIdStorageKey = 'user-id' +/** + * When more due jobs remain after an alarm run (each run processes at most + * `maxDueJobsPerAlarm` jobs), the resync arms a near-immediate follow-up + * alarm instead of one at the (already past) next_run_at so backlogs drain + * across multiple short invocations. + */ +const jobBacklogRearmDelayMs = 1_000 + export class JobManagerBase extends DurableObject { async purgeUser(input: { userId: string }) { const userId = input.userId.trim() @@ -70,7 +78,12 @@ export class JobManagerBase extends DurableObject { nextRunAt: null, } } - await this.ctx.storage.setAlarm(new Date(nextJob.nextRunAt)) + const nextRunAtMs = new Date(nextJob.nextRunAt).valueOf() + const backlogRemains = nextRunAtMs <= Date.now() + const alarmAtMs = backlogRemains + ? Date.now() + jobBacklogRearmDelayMs + : nextRunAtMs + await this.ctx.storage.setAlarm(alarmAtMs) logJobSchedulerEvent({ event: 'sync_alarm', userId, @@ -80,8 +93,9 @@ export class JobManagerBase extends DurableObject { : new Date(currentAlarmAt).toISOString(), nextJobId: nextJob.id, nextRunAt: nextJob.nextRunAt, - reason: - currentAlarmAt === new Date(nextJob.nextRunAt).valueOf() + reason: backlogRemains + ? 'backlog_rearmed' + : currentAlarmAt === nextRunAtMs ? 'alarm_unchanged' : 'alarm_armed', }) diff --git a/packages/worker/src/jobs/repo.ts b/packages/worker/src/jobs/repo.ts index 23b8e27cb2..7ad3e0be2e 100644 --- a/packages/worker/src/jobs/repo.ts +++ b/packages/worker/src/jobs/repo.ts @@ -276,6 +276,31 @@ export async function listJobRowsByUserId( return (results ?? []).map(mapRow) } +/** + * Caps how many due jobs a single JobManager alarm invocation picks up so a + * large backlog cannot exhaust the Durable Object CPU budget in one wake. The + * post-run alarm resync schedules a near-immediate follow-up wake whenever + * more due jobs remain, so backlogs drain across successive alarms. + */ +export const maxDueJobsPerAlarm = 25 + +// Keyset-paged listing across all users, used by maintenance reindex so a +// single query never loads the whole table. Pass the last row id of the +// previous page (or null for the first page). +export async function listJobRowsPage( + db: D1Database, + input: { + afterId: string | null + limit: number + }, +): Promise> { + const { results } = await db + .prepare(`SELECT * FROM jobs WHERE id > ? ORDER BY id LIMIT ?`) + .bind(input.afterId ?? '', input.limit) + .all>() + return (results ?? []).map(mapRow) +} + export async function listDueJobRows( db: D1Database, userId: string, @@ -288,9 +313,10 @@ export async function listDueJobRows( AND enabled = 1 AND kill_switch_enabled = 0 AND next_run_at <= ? - ORDER BY next_run_at ASC, name ASC`, + ORDER BY next_run_at ASC, name ASC + LIMIT ?`, ) - .bind(userId, nowIso) + .bind(userId, nowIso, maxDueJobsPerAlarm) .all>() return (results ?? []).map(mapRow) } diff --git a/packages/worker/src/jobs/repo.workers.test.ts b/packages/worker/src/jobs/repo.workers.test.ts new file mode 100644 index 0000000000..9356afcd7e --- /dev/null +++ b/packages/worker/src/jobs/repo.workers.test.ts @@ -0,0 +1,132 @@ +import { env } from 'cloudflare:workers' +import { expect, test } from 'vitest' +import { listDueJobRows, maxDueJobsPerAlarm } from './repo.ts' + +async function ensureJobsSchema() { + await env.APP_DB.prepare( + `CREATE TABLE IF NOT EXISTS jobs ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL, + name TEXT NOT NULL, + source_id TEXT NOT NULL, + published_commit TEXT, + repo_check_policy_json TEXT, + storage_id TEXT NOT NULL, + params_json TEXT, + schedule_json TEXT NOT NULL, + timezone TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1 CHECK (enabled IN (0, 1)), + kill_switch_enabled INTEGER NOT NULL DEFAULT 0 CHECK (kill_switch_enabled IN (0, 1)), + caller_context_json TEXT NOT NULL, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + last_run_at TEXT, + last_run_status TEXT, + last_run_error TEXT, + last_duration_ms INTEGER, + next_run_at TEXT NOT NULL, + run_count INTEGER NOT NULL DEFAULT 0, + success_count INTEGER NOT NULL DEFAULT 0, + error_count INTEGER NOT NULL DEFAULT 0, + run_history_json TEXT NOT NULL DEFAULT '[]' + )`, + ).run() + await env.APP_DB.prepare(`DELETE FROM jobs`).run() +} + +async function insertJob(input: { + id: string + userId: string + nextRunAt: string + enabled?: boolean + killSwitchEnabled?: boolean +}) { + const now = '2026-04-20T00:00:00.000Z' + await env.APP_DB.prepare( + `INSERT INTO jobs ( + id, user_id, name, source_id, storage_id, schedule_json, timezone, + enabled, kill_switch_enabled, caller_context_json, created_at, + updated_at, next_run_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'null', ?, ?, ?)`, + ) + .bind( + input.id, + input.userId, + input.id, + `source-${input.id}`, + `job:${input.id}`, + JSON.stringify({ type: 'once', runAt: input.nextRunAt }), + 'UTC', + input.enabled === false ? 0 : 1, + input.killSwitchEnabled === true ? 1 : 0, + now, + now, + input.nextRunAt, + ) + .run() +} + +test('listDueJobRows caps a due-job backlog at maxDueJobsPerAlarm, oldest first', async () => { + await ensureJobsSchema() + const userId = 'user-due-limit' + const nowIso = '2026-04-20T12:00:00.000Z' + const backlogSize = maxDueJobsPerAlarm + 5 + for (let index = 0; index < backlogSize; index += 1) { + await insertJob({ + id: `due-${String(index).padStart(3, '0')}`, + userId, + nextRunAt: new Date( + Date.parse('2026-04-20T00:00:00.000Z') + index * 60_000, + ).toISOString(), + }) + } + // Rows that must never be picked up: other user, disabled, kill-switched, + // and not-yet-due jobs. + await insertJob({ + id: 'other-user', + userId: 'user-other', + nextRunAt: '2026-04-20T00:00:00.000Z', + }) + await insertJob({ + id: 'disabled', + userId, + nextRunAt: '2026-04-20T00:00:00.000Z', + enabled: false, + }) + await insertJob({ + id: 'kill-switched', + userId, + nextRunAt: '2026-04-20T00:00:00.000Z', + killSwitchEnabled: true, + }) + await insertJob({ + id: 'future', + userId, + nextRunAt: '2026-04-21T00:00:00.000Z', + }) + + const firstBatch = await listDueJobRows(env.APP_DB, userId, nowIso) + expect(firstBatch).toHaveLength(maxDueJobsPerAlarm) + expect(firstBatch.map((row) => row.id)).toEqual( + Array.from( + { length: maxDueJobsPerAlarm }, + (_, index) => `due-${String(index).padStart(3, '0')}`, + ), + ) + + // Once the first batch has been rescheduled out of the due window, the next + // alarm invocation picks up the remainder of the backlog. + for (const row of firstBatch) { + await env.APP_DB.prepare(`UPDATE jobs SET next_run_at = ? WHERE id = ?`) + .bind('2026-04-22T00:00:00.000Z', row.id) + .run() + } + const secondBatch = await listDueJobRows(env.APP_DB, userId, nowIso) + expect(secondBatch.map((row) => row.id)).toEqual( + Array.from( + { length: backlogSize - maxDueJobsPerAlarm }, + (_, index) => + `due-${String(maxDueJobsPerAlarm + index).padStart(3, '0')}`, + ), + ) +}) 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 1f29890a60..cfeef10ce4 100644 --- a/packages/worker/src/package-runtime/package-service.node.test.ts +++ b/packages/worker/src/package-runtime/package-service.node.test.ts @@ -64,6 +64,7 @@ const { PackageServiceInstance, buildPackageServiceStorageId, buildServiceRuntimeUsageEvent, + computePackageServiceRetryDelayMs, } = await import('./package-service.ts') const serviceBinding = { @@ -81,7 +82,10 @@ const savedPackage = { kodyId: '@scope/package-1', } -function createPackageSource(mode: 'bounded' | 'persistent') { +function createPackageSource( + mode: 'bounded' | 'persistent', + options?: { autoStart?: boolean }, +) { return { manifest: { kody: { @@ -91,6 +95,7 @@ function createPackageSource(mode: 'bounded' | 'persistent') { entry: 'services/realtime-supervisor.ts', mode, timeoutMs: mode === 'bounded' ? 30_000 : null, + ...(options?.autoStart === true ? { autoStart: true } : {}), }, }, }, @@ -396,3 +401,184 @@ test('durable object restore meters an eviction-orphaned in-flight run without a restoreUsageSpy.mockRestore() } }) + +test('computePackageServiceRetryDelayMs doubles from the base delay to a 15-minute cap with jitter', () => { + const atJitterMax = () => 1 + const atJitterMin = () => 0 + expect( + computePackageServiceRetryDelayMs({ consecutiveFailureCount: 0 }), + ).toBe(5_000) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 1, + random: atJitterMax, + }), + ).toBe(5_000) + // Equal jitter spreads a retry between half and the full nominal delay. + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 1, + random: atJitterMin, + }), + ).toBe(2_500) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 2, + random: atJitterMax, + }), + ).toBe(10_000) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 3, + random: atJitterMax, + }), + ).toBe(20_000) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 8, + random: atJitterMax, + }), + ).toBe(640_000) + // 5s * 2^8 exceeds the cap, so the ninth consecutive failure pins to 15m. + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 9, + random: atJitterMax, + }), + ).toBe(900_000) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 50, + random: atJitterMax, + }), + ).toBe(900_000) + expect( + computePackageServiceRetryDelayMs({ + consecutiveFailureCount: 50, + random: atJitterMin, + }), + ).toBe(450_000) +}) + +test('auto-start crash loops back off exponentially and reset after a successful run', async () => { + resetMocks() + const usageSpy = vi + .spyOn(usageModule, 'recordUsage') + .mockResolvedValue(undefined) + const randomSpy = vi.spyOn(Math, 'random').mockReturnValue(1) + vi.useFakeTimers() + try { + const startAt = new Date('2026-07-05T12:00:00.000Z') + vi.setSystemTime(startAt) + mockModule.getSavedPackageById.mockResolvedValue(savedPackage) + mockModule.loadPackageSourceBySourceId.mockResolvedValue( + createPackageSource('bounded', { autoStart: true }), + ) + mockModule.loadPublishedBundleArtifactByIdentity.mockResolvedValue({ + artifact: { + mainModule: 'main', + modules: {}, + }, + }) + mockModule.runBundledModuleWithRegistry.mockRejectedValue( + new Error('service crashed'), + ) + + const created = createPackageServiceState() + created.persistedEntries.set('package-service-state', { + binding: serviceBinding, + autoStart: true, + mode: 'bounded', + timeoutMs: 30_000, + stopRequested: false, + currentRunId: null, + nextAlarmAt: startAt.toISOString(), + nextAlarmSource: 'auto-start', + lastStartedAt: null, + lastStoppedAt: null, + status: 'stopped', + lastError: null, + lastResult: null, + lastRunFinishedAt: null, + consecutiveFailureCount: 0, + }) + const instance = new PackageServiceInstance(created.state, { + APP_DB: {}, + } as Env) + await waitForRestoreState(created.state) + + const expectedDelays = [5_000, 10_000, 20_000] + for (const [index, expectedDelay] of expectedDelays.entries()) { + await instance.alarm() + await flushWaitUntilTasks(created.waitUntilTasks) + expect(created.getAlarmAt()).toBe(startAt.valueOf() + expectedDelay) + const persisted = created.persistedEntries.get( + 'package-service-state', + ) as { consecutiveFailureCount: number } + expect(persisted.consecutiveFailureCount).toBe(index + 1) + } + + mockModule.runBundledModuleWithRegistry.mockResolvedValue({ + result: { ok: true }, + error: null, + }) + await instance.alarm() + await flushWaitUntilTasks(created.waitUntilTasks) + // A successful run resets the crash-loop backoff to the base delay. + expect(created.getAlarmAt()).toBe(startAt.valueOf() + 5_000) + const persistedAfterSuccess = created.persistedEntries.get( + 'package-service-state', + ) as { consecutiveFailureCount: number } + expect(persistedAfterSuccess.consecutiveFailureCount).toBe(0) + } finally { + randomSpy.mockRestore() + usageSpy.mockRestore() + vi.useRealTimers() + } +}) + +test('durable object restore keeps the persisted crash-loop backoff after eviction', async () => { + resetMocks() + const usageSpy = vi + .spyOn(usageModule, 'recordUsage') + .mockResolvedValue(undefined) + const randomSpy = vi.spyOn(Math, 'random').mockReturnValue(1) + vi.useFakeTimers() + try { + const restoredAt = new Date('2026-07-05T13:00:00.000Z') + vi.setSystemTime(restoredAt) + const restored = createPackageServiceState() + restored.persistedEntries.set('package-service-state', { + binding: serviceBinding, + autoStart: true, + mode: 'bounded', + timeoutMs: 30_000, + stopRequested: false, + currentRunId: 'run-evicted', + nextAlarmAt: null, + nextAlarmSource: null, + lastStartedAt: '2026-07-05T12:59:00.000Z', + lastStoppedAt: null, + status: 'running', + lastError: 'service crashed', + lastResult: null, + lastRunFinishedAt: null, + consecutiveFailureCount: 4, + }) + new PackageServiceInstance(restored.state, { APP_DB: {} } as Env) + await waitForRestoreState(restored.state) + await flushWaitUntilTasks(restored.waitUntilTasks) + + // The retry after eviction honors the persisted failure count + // (5s * 2^3 = 40s) instead of resetting to the base delay. + expect(restored.getAlarmAt()).toBe(restoredAt.valueOf() + 40_000) + const persisted = restored.persistedEntries.get( + 'package-service-state', + ) as { consecutiveFailureCount: number } + expect(persisted.consecutiveFailureCount).toBe(4) + } finally { + randomSpy.mockRestore() + usageSpy.mockRestore() + vi.useRealTimers() + } +}) diff --git a/packages/worker/src/package-runtime/package-service.ts b/packages/worker/src/package-runtime/package-service.ts index 352711895a..4cb0bd4ab6 100644 --- a/packages/worker/src/package-runtime/package-service.ts +++ b/packages/worker/src/package-runtime/package-service.ts @@ -18,6 +18,7 @@ import { assertPublishedSourceCanRebuildWithoutInstallingDeps } from './publishe const serviceStateStorageKey = 'package-service-state' const packageServiceRetryDelayMs = 5_000 +const packageServiceRetryMaxDelayMs = 15 * 60 * 1000 export type PackageServiceBindingState = { userId: string @@ -43,6 +44,12 @@ type PackageServiceState = { lastError: string | null lastResult: unknown lastRunFinishedAt: string | null + /** + * Consecutive failed runs (or failed auto-start attempts) since the last + * successful run. Persisted so Durable Object eviction does not reset the + * crash-loop backoff. + */ + consecutiveFailureCount: number } type PackageServiceRunResult = { @@ -101,6 +108,7 @@ function createInitialPackageServiceState(): PackageServiceState { lastError: null, lastResult: null, lastRunFinishedAt: null, + consecutiveFailureCount: 0, } } @@ -229,8 +237,32 @@ async function readPackageServiceRpcResponse( } } -function buildPackageServiceRetryTime() { - return new Date(Date.now() + packageServiceRetryDelayMs) +/** + * Auto-start retry delay: the 5-second base doubles per consecutive failure up + * to a 15-minute cap, with equal jitter (a retry lands between half and the + * full nominal delay) so crash-looping services do not retry in lockstep. A + * successful run resets the failure count back to the base delay. + */ +export function computePackageServiceRetryDelayMs(input: { + consecutiveFailureCount: number + random?: () => number +}) { + if (input.consecutiveFailureCount <= 0) { + return packageServiceRetryDelayMs + } + const cappedExponent = Math.min(input.consecutiveFailureCount - 1, 16) + const nominalDelayMs = Math.min( + packageServiceRetryDelayMs * 2 ** cappedExponent, + packageServiceRetryMaxDelayMs, + ) + const random = input.random ?? Math.random + return Math.round(nominalDelayMs / 2 + random() * (nominalDelayMs / 2)) +} + +function buildPackageServiceRetryTime(consecutiveFailureCount: number) { + return new Date( + Date.now() + computePackageServiceRetryDelayMs({ consecutiveFailureCount }), + ) } class PackageServiceInstanceBase extends DurableObject { @@ -291,7 +323,9 @@ class PackageServiceInstanceBase extends DurableObject { this.stateSnapshot.binding ) { await this.scheduleAlarm({ - runAt: buildPackageServiceRetryTime(), + runAt: buildPackageServiceRetryTime( + this.stateSnapshot.consecutiveFailureCount, + ), source: 'auto-start', }) } @@ -429,6 +463,10 @@ class PackageServiceInstanceBase extends DurableObject { if ('lastError' in input) { this.stateSnapshot.lastError = input.lastError ?? null } + this.stateSnapshot.consecutiveFailureCount = + input.lastError == null + ? 0 + : this.stateSnapshot.consecutiveFailureCount + 1 this.stateSnapshot.lastRunFinishedAt = new Date(finishedAtMs).toISOString() this.stateSnapshot.lastStoppedAt = this.stateSnapshot.lastRunFinishedAt await this.persistState() @@ -441,7 +479,9 @@ class PackageServiceInstanceBase extends DurableObject { !this.stateSnapshot.nextAlarmAt ) { await this.scheduleAlarm({ - runAt: buildPackageServiceRetryTime(), + runAt: buildPackageServiceRetryTime( + this.stateSnapshot.consecutiveFailureCount, + ), source: 'auto-start', }) } @@ -844,6 +884,7 @@ class PackageServiceInstanceBase extends DurableObject { this.stateSnapshot.lastError = error instanceof Error ? error.message : String(error) this.stateSnapshot.status = 'error' + this.stateSnapshot.consecutiveFailureCount += 1 await this.persistState() if ( this.stateSnapshot.autoStart && @@ -851,7 +892,9 @@ class PackageServiceInstanceBase extends DurableObject { !this.stateSnapshot.nextAlarmAt ) { await this.scheduleAlarm({ - runAt: buildPackageServiceRetryTime(), + runAt: buildPackageServiceRetryTime( + this.stateSnapshot.consecutiveFailureCount, + ), source: alarmSource, }) } diff --git a/packages/worker/src/repo/checks.node.test.ts b/packages/worker/src/repo/checks.node.test.ts index 82289aba92..d4a662cc71 100644 --- a/packages/worker/src/repo/checks.node.test.ts +++ b/packages/worker/src/repo/checks.node.test.ts @@ -27,7 +27,11 @@ vi.mock('#worker/package-runtime/module-graph.ts', () => ({ mockModule.buildKodyModuleBundle(...args), })) -import { runRepoChecks } from './checks.ts' +import { + repoChecksSourceMaxFiles, + repoChecksSourceMaxTotalBytes, + runRepoChecks, +} from './checks.ts' type MockSnapshot = { read: ReturnType @@ -186,6 +190,71 @@ async function runPackageJobTypecheckChecks( return { result, typeScriptFileSystem, getSemanticDiagnostics } } +test('runRepoChecks fails with an actionable check result when the source root exceeds the file-count cap', async () => { + const files = new Map([ + [ + 'package.json', + createPackageManifest({ + packageName: '@kody/too-many-files', + kodyId: 'too-many-files', + description: 'Exceeds the publish check file cap', + }), + ], + ['src/index.ts', 'export const ready = true\n'], + ]) + for (let index = 0; index < repoChecksSourceMaxFiles; index += 1) { + files.set(`generated/file-${index}.txt`, 'x') + } + + const result = await runChecksOnWorkspaceFiles(files) + + expect(result.ok).toBe(false) + expect(result.sourceFiles).toEqual({}) + expect(result.results).toEqual([ + expect.objectContaining({ kind: 'manifest', ok: true }), + expect.objectContaining({ + kind: 'bundle', + ok: false, + message: expect.stringContaining( + `${repoChecksSourceMaxFiles}-file publish check limit`, + ), + }), + ]) +}) + +test('runRepoChecks fails with an actionable check result when the source root exceeds the total-byte cap', async () => { + const halfCapChunk = 'x'.repeat(repoChecksSourceMaxTotalBytes / 2) + const files = new Map([ + [ + 'package.json', + createPackageManifest({ + packageName: '@kody/too-many-bytes', + kodyId: 'too-many-bytes', + description: 'Exceeds the publish check byte cap', + }), + ], + ['src/index.ts', 'export const ready = true\n'], + ['assets/blob-1.bin', halfCapChunk], + ['assets/blob-2.bin', halfCapChunk], + ['assets/blob-3.bin', halfCapChunk], + ]) + + const result = await runChecksOnWorkspaceFiles(files) + + expect(result.ok).toBe(false) + expect(result.sourceFiles).toEqual({}) + expect(result.results).toEqual([ + expect.objectContaining({ kind: 'manifest', ok: true }), + expect.objectContaining({ + kind: 'bundle', + ok: false, + message: expect.stringContaining( + `${repoChecksSourceMaxTotalBytes}-byte (15 MiB) publish check limit`, + ), + }), + ]) +}) + test('runRepoChecks keeps non-code source files for publish snapshots while excluding git internals', async () => { setupDefaultBundleMocks() const files = new Map([ diff --git a/packages/worker/src/repo/checks.ts b/packages/worker/src/repo/checks.ts index b5832f6252..e1bd5e8f45 100644 --- a/packages/worker/src/repo/checks.ts +++ b/packages/worker/src/repo/checks.ts @@ -51,6 +51,16 @@ const repoChecksSyntheticTsconfigPath = 'tsconfig.json' const repoChecksSyntheticTsconfigExtendsPath = './.__kody_repo_tsconfig_base__.json' +/** + * Publish checks materialize the whole source root in memory before bundling + * and typechecking, so the walk is capped to keep a single check run within + * Durable Object CPU/memory limits. The caps are intentionally larger than + * the execution-path caps in `repo-kody-execution.ts` (250 files / 2 MiB) + * because published packages may ship assets alongside runtime code. + */ +export const repoChecksSourceMaxFiles = 2_000 +export const repoChecksSourceMaxTotalBytes = 15 * 1024 * 1024 + async function loadWorkerBundlerSnapshotTools() { // Keep the experimental bundler out of the Worker's top-level deploy graph. const { createFileSystemSnapshot } = @@ -885,16 +895,50 @@ export async function runRepoChecks(input: { /\/+$/, '', ) - const sourceFiles = await (async () => { + const sourceWalk = await (async () => { const collected: Record = {} + const encoder = new TextEncoder() + let fileCount = 0 + let totalBytes = 0 for await (const [path, content] of workspaceFilesForSnapshot({ workspace: input.workspace, root: sourceRoot, })) { + fileCount += 1 + if (fileCount > repoChecksSourceMaxFiles) { + return { + ok: false as const, + message: `Repo checks aborted: source root "${sourceRoot || '/'}" contains more than the ${repoChecksSourceMaxFiles}-file publish check limit. Remove files that should not be published (for example build output, vendored dependencies, or data files) and run the checks again.`, + } + } + totalBytes += encoder.encode(content).byteLength + if (totalBytes > repoChecksSourceMaxTotalBytes) { + return { + ok: false as const, + message: `Repo checks aborted: source root "${sourceRoot || '/'}" exceeds the ${repoChecksSourceMaxTotalBytes}-byte (15 MiB) publish check limit. Remove or shrink large files that should not be published (for example build output, vendored dependencies, or data files) and run the checks again.`, + } + } collected[path] = content } - return collected + return { + ok: true as const, + collected, + } })() + if (!sourceWalk.ok) { + results.push({ + kind: 'bundle', + ok: false, + message: sourceWalk.message, + }) + return { + ok: false, + results, + manifest, + sourceFiles: {}, + } + } + const sourceFiles = sourceWalk.collected const { createFileSystemSnapshot } = await loadWorkerBundlerSnapshotTools() const snapshot = await createFileSystemSnapshot( (async function* () { From 4724dcd92da9fea063355980027efb0e43865e19 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:31 +0000 Subject: [PATCH 05/14] Make search paths cheap: projections, vectors, bounded scans Package search ranks D1 projections fused with the existing package vectors and hydrates full source only for returned matches; memory search queries Vectorize first and hydrates top hits by id; community search pre-filters in SQL and bounds candidates at 500. Co-authored-by: Kent C. Dodds --- .../community/community-scoring.node.test.ts | 17 +- .../community/community-service.node.test.ts | 8 +- packages/worker/src/community/repo.ts | 146 +++++++--- packages/worker/src/community/service.ts | 83 +++--- .../worker/src/mcp/memory/memory-search.ts | 52 +++- packages/worker/src/mcp/memory/repo.ts | 41 ++- .../src/mcp/memory/service.node.test.ts | 133 ++++++++++ packages/worker/src/mcp/memory/service.ts | 72 ++++- .../worker/src/mcp/tools/search.node.test.ts | 216 ++++++++++----- packages/worker/src/mcp/tools/search.ts | 251 ++++++++++++++++-- 10 files changed, 833 insertions(+), 186 deletions(-) diff --git a/packages/worker/src/community/community-scoring.node.test.ts b/packages/worker/src/community/community-scoring.node.test.ts index 5ce335ef62..0e0e9ae472 100644 --- a/packages/worker/src/community/community-scoring.node.test.ts +++ b/packages/worker/src/community/community-scoring.node.test.ts @@ -2,14 +2,14 @@ import { expect, test, vi } from 'vitest' import { type CommunityListingRecord } from './types.ts' const mockModule = vi.hoisted(() => ({ - listAllCommunityListings: vi.fn(), + listCommunityListingCandidates: vi.fn(), getCommunityRatingAggregatesByListingIds: vi.fn(), countCommunityForksByListingIds: vi.fn(), })) vi.mock('./repo.ts', () => ({ - listAllCommunityListings: (...args: Array) => - mockModule.listAllCommunityListings(...args), + listCommunityListingCandidates: (...args: Array) => + mockModule.listCommunityListingCandidates(...args), getCommunityRatingAggregatesByListingIds: (...args: Array) => mockModule.getCommunityRatingAggregatesByListingIds(...args), countCommunityForksByListingIds: (...args: Array) => @@ -17,6 +17,7 @@ vi.mock('./repo.ts', () => ({ })) const { + COMMUNITY_SEARCH_CANDIDATE_LIMIT, buildCommunityListingSearchDocument, computeCommunityBayesianScore, isCommunityListingSearchMatch, @@ -73,7 +74,7 @@ function mealListing(): CommunityListingRecord { } function mockListings(listings: Array) { - mockModule.listAllCommunityListings.mockResolvedValue(listings) + mockModule.listCommunityListingCandidates.mockResolvedValue(listings) mockModule.getCommunityRatingAggregatesByListingIds.mockResolvedValue( Object.fromEntries( listings.map((listing) => [ @@ -143,6 +144,14 @@ test('community scoring and search rank listings and filter by query', async () expect(githubResults.map((listing) => listing.kodyId)).toEqual([ 'github-triage', ]) + expect(mockModule.listCommunityListingCandidates).toHaveBeenCalledWith( + expect.anything(), + { + includeDelisted: false, + limit: COMMUNITY_SEARCH_CANDIDATE_LIMIT, + query: 'github', + }, + ) const mealResults = await searchCommunityListings({ env: createEnv(), diff --git a/packages/worker/src/community/community-service.node.test.ts b/packages/worker/src/community/community-service.node.test.ts index 8182610f7f..aadee1deae 100644 --- a/packages/worker/src/community/community-service.node.test.ts +++ b/packages/worker/src/community/community-service.node.test.ts @@ -7,7 +7,7 @@ const mockModule = vi.hoisted(() => ({ getCommunityBan: vi.fn(), getCommunityListingByOwnerAndPackage: vi.fn(), getCommunityListingById: vi.fn(), - listAllCommunityListings: vi.fn(), + listCommunityListingCandidates: vi.fn(), getCommunityRatingAggregatesByListingIds: vi.fn(), countCommunityForksByListingIds: vi.fn(), writeCommunitySnapshot: vi.fn(), @@ -74,8 +74,8 @@ vi.mock('./repo.ts', () => ({ mockModule.getCommunityListingByOwnerAndPackage(...args), getCommunityListingById: (...args: Array) => mockModule.getCommunityListingById(...args), - listAllCommunityListings: (...args: Array) => - mockModule.listAllCommunityListings(...args), + listCommunityListingCandidates: (...args: Array) => + mockModule.listCommunityListingCandidates(...args), getCommunityRatingAggregatesByListingIds: (...args: Array) => mockModule.getCommunityRatingAggregatesByListingIds(...args), countCommunityForksByListingIds: (...args: Array) => @@ -385,7 +385,7 @@ test('searchCommunityListings empty query uses publishedAt tiebreaker', async () id: 'listing-newer', publishedAt: '2026-07-03T00:00:00.000Z', }) - mockModule.listAllCommunityListings.mockResolvedValue([ + mockModule.listCommunityListingCandidates.mockResolvedValue([ olderListing, newerListing, ]) diff --git a/packages/worker/src/community/repo.ts b/packages/worker/src/community/repo.ts index ecf3c227e3..e20248d51f 100644 --- a/packages/worker/src/community/repo.ts +++ b/packages/worker/src/community/repo.ts @@ -100,6 +100,37 @@ function listingStatusFilter(includeDelisted: boolean) { return includeDelisted ? '' : `WHERE status = 'active'` } +// D1 caps bound parameters per statement, so IN (...) lookups over large id +// sets are issued in chunks. +const maxSqlBindingsPerChunk = 90 + +function chunkValues(values: Array, size: number): Array> { + const chunks: Array> = [] + for (let offset = 0; offset < values.length; offset += size) { + chunks.push(values.slice(offset, offset + size)) + } + return chunks +} + +const communityListingSearchTextColumns = [ + 'name', + 'kody_id', + 'description', + 'search_text', + 'tags_json', + 'readme_content', +] as const + +const maxCommunityListingLikeTokens = 8 + +function extractCommunityListingLikeTokens(query: string): Array { + // Tokens are restricted to [a-z0-9]+ so they are safe to embed inside + // LIKE patterns without wildcard escaping. + return Array.from( + new Set(query.toLowerCase().match(/[a-z0-9]+/g) ?? []), + ).slice(0, maxCommunityListingLikeTokens) +} + export async function insertCommunityListing( db: D1Database, row: Omit< @@ -271,21 +302,50 @@ export async function listCommunityListings( return (rows.results ?? []).map(mapCommunityListingRow) } -export async function listAllCommunityListings( +/** + * Bounded candidate query for community search/browse. Applies SQL-level text + * pre-filtering (LIKE across name/kody_id/description/search_text/tags_json/ + * readme_content) when a query is provided, orders by recency, and never + * returns more than `limit` rows so scoring stays off the full table. + */ +export async function listCommunityListingCandidates( db: D1Database, input: { includeDelisted: boolean + limit: number + query?: string | null }, ): Promise> { + const conditions: Array = [] + const bindings: Array = [] + if (!input.includeDelisted) { + conditions.push(`status = 'active'`) + } + const tokens = extractCommunityListingLikeTokens(input.query ?? '') + if (tokens.length > 0) { + const tokenClauses = tokens.map((token) => { + const pattern = `%${token}%` + const columnClauses = communityListingSearchTextColumns.map((column) => { + bindings.push(pattern) + return `${column} LIKE ?` + }) + return `(${columnClauses.join(' OR ')})` + }) + conditions.push(`(${tokenClauses.join(' OR ')})`) + } + const whereClause = + conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : '' const rows = await db .prepare( `SELECT id, owner_user_id, package_id, source_id, kody_id, name, description, tags_json, search_text, readme_content, license, pinned_commit, status, created_at, updated_at, published_at FROM community_listings - ${listingStatusFilter(input.includeDelisted)} - ORDER BY published_at DESC`, + ${whereClause} + ORDER BY published_at DESC + LIMIT ?`, ) + .bind(...bindings, input.limit) .all>() return (rows.results ?? []).map(mapCommunityListingRow) } @@ -395,21 +455,23 @@ export async function countCommunityForksByListingIds( listingIds: Array, ): Promise> { if (listingIds.length === 0) return {} - const placeholders = listingIds.map(() => '?').join(', ') - const rows = await db - .prepare( - `SELECT listing_id, COUNT(*) AS fork_count - FROM community_forks - WHERE listing_id IN (${placeholders}) - GROUP BY listing_id`, - ) - .bind(...listingIds) - .all>() const counts: Record = Object.fromEntries( listingIds.map((listingId) => [listingId, 0]), ) - for (const row of rows.results ?? []) { - counts[String(row['listing_id'])] = Number(row['fork_count'] ?? 0) + for (const idChunk of chunkValues(listingIds, maxSqlBindingsPerChunk)) { + const placeholders = idChunk.map(() => '?').join(', ') + const rows = await db + .prepare( + `SELECT listing_id, COUNT(*) AS fork_count + FROM community_forks + WHERE listing_id IN (${placeholders}) + GROUP BY listing_id`, + ) + .bind(...idChunk) + .all>() + for (const row of rows.results ?? []) { + counts[String(row['listing_id'])] = Number(row['fork_count'] ?? 0) + } } return counts } @@ -489,20 +551,6 @@ export async function getCommunityRatingAggregatesByListingIds( listingIds: Array, ): Promise> { if (listingIds.length === 0) return {} - const placeholders = listingIds.map(() => '?').join(', ') - const rows = await db - .prepare( - `SELECT - listing_id, - COUNT(*) AS rating_count, - AVG(stars) AS average_stars, - AVG(adaptation_effort) AS average_adaptation_effort - FROM community_ratings - WHERE listing_id IN (${placeholders}) - GROUP BY listing_id`, - ) - .bind(...listingIds) - .all>() const aggregates: Record = Object.fromEntries( listingIds.map((listingId) => [ @@ -515,18 +563,34 @@ export async function getCommunityRatingAggregatesByListingIds( }, ]), ) - for (const row of rows.results ?? []) { - const listingId = String(row['listing_id']) - const ratingCount = Number(row['rating_count'] ?? 0) - aggregates[listingId] = { - listingId, - ratingCount, - averageStars: - ratingCount === 0 ? null : Number(row['average_stars'] ?? 0), - averageAdaptationEffort: - ratingCount === 0 - ? null - : Number(row['average_adaptation_effort'] ?? 0), + for (const idChunk of chunkValues(listingIds, maxSqlBindingsPerChunk)) { + const placeholders = idChunk.map(() => '?').join(', ') + const rows = await db + .prepare( + `SELECT + listing_id, + COUNT(*) AS rating_count, + AVG(stars) AS average_stars, + AVG(adaptation_effort) AS average_adaptation_effort + FROM community_ratings + WHERE listing_id IN (${placeholders}) + GROUP BY listing_id`, + ) + .bind(...idChunk) + .all>() + for (const row of rows.results ?? []) { + const listingId = String(row['listing_id']) + const ratingCount = Number(row['rating_count'] ?? 0) + aggregates[listingId] = { + listingId, + ratingCount, + averageStars: + ratingCount === 0 ? null : Number(row['average_stars'] ?? 0), + averageAdaptationEffort: + ratingCount === 0 + ? null + : Number(row['average_adaptation_effort'] ?? 0), + } } } return aggregates diff --git a/packages/worker/src/community/service.ts b/packages/worker/src/community/service.ts index eab50cba6c..daefcc0ac8 100644 --- a/packages/worker/src/community/service.ts +++ b/packages/worker/src/community/service.ts @@ -34,7 +34,7 @@ import { insertCommunityListing, insertCommunityBan, insertCommunityReport, - listAllCommunityListings, + listCommunityListingCandidates, listCommunityReports as listCommunityReportsFromDb, resolveCommunityReportRow, setCommunityListingStatus, @@ -69,6 +69,10 @@ const intentHeadingPattern = /^##\s+intent\b/im // semantic matches on listing documents tend to land above ~0.12. export const COMMUNITY_SEARCH_VECTOR_MATCH_THRESHOLD = 0.12 +// Community search/browse never scores more than this many listings in +// memory; candidates are pre-filtered and recency-ordered in SQL. +export const COMMUNITY_SEARCH_CANDIDATE_LIMIT = 500 + export function isCommunityListingSearchMatch(input: { query: string document: string @@ -470,8 +474,9 @@ export async function listCommunityListingsWithAggregates(input: { limit: number offset: number }): Promise> { - const listings = await listAllCommunityListings(input.env.APP_DB, { + const listings = await listCommunityListingCandidates(input.env.APP_DB, { includeDelisted: input.includeDelisted, + limit: COMMUNITY_SEARCH_CANDIDATE_LIMIT, }) const withAggregates = await attachListingAggregatesBatch( input.env.APP_DB, @@ -487,48 +492,58 @@ export async function searchCommunityListings(input: { query: string limit: number }): Promise> { - const listings = await listAllCommunityListings(input.env.APP_DB, { + const trimmedQuery = input.query.trim() + let listings = await listCommunityListingCandidates(input.env.APP_DB, { includeDelisted: false, + limit: COMMUNITY_SEARCH_CANDIDATE_LIMIT, + query: trimmedQuery || null, }) - const withAggregates = await attachListingAggregatesBatch( - input.env.APP_DB, - listings, - ) - const trimmedQuery = input.query.trim() if (!trimmedQuery) { + const withAggregates = await attachListingAggregatesBatch( + input.env.APP_DB, + listings, + ) return withAggregates .sort(compareCommunityListingsByBayesianAndPublishedAt) .slice(0, input.limit) } + if (listings.length === 0) { + // The SQL LIKE pre-filter found nothing; fall back to a bounded recent + // candidate set so vector-threshold matches keep working. + listings = await listCommunityListingCandidates(input.env.APP_DB, { + includeDelisted: false, + limit: COMMUNITY_SEARCH_CANDIDATE_LIMIT, + }) + } + const matched = listings.filter((listing) => + isCommunityListingSearchMatch({ + query: trimmedQuery, + document: buildListingSearchDocument(listing), + }), + ) + const withAggregates = await attachListingAggregatesBatch( + input.env.APP_DB, + matched, + ) const queryEmbedding = deterministicEmbedding(trimmedQuery) - const scored = withAggregates - .map((listing) => { - const document = buildListingSearchDocument(listing) - if ( - !isCommunityListingSearchMatch({ - query: trimmedQuery, - document, - }) - ) { - return null - } - const lexical = lexicalScore(trimmedQuery, document) - const vector = cosineSimilarity( - queryEmbedding, - deterministicEmbedding(document), - ) - const blended = blendLexicalAndVectorScore(lexical, vector) - const bayesian = computeCommunityBayesianScore({ - averageStars: listing.averageStars, - ratingCount: listing.ratingCount, - }) - return { - listing, - score: blended * bayesian, - } + const scored = withAggregates.map((listing) => { + const document = buildListingSearchDocument(listing) + const lexical = lexicalScore(trimmedQuery, document) + const vector = cosineSimilarity( + queryEmbedding, + deterministicEmbedding(document), + ) + const blended = blendLexicalAndVectorScore(lexical, vector) + const bayesian = computeCommunityBayesianScore({ + averageStars: listing.averageStars, + ratingCount: listing.ratingCount, }) - .filter((entry): entry is NonNullable => entry != null) + return { + listing, + score: blended * bayesian, + } + }) return scored .sort((left, right) => right.score - left.score) diff --git a/packages/worker/src/mcp/memory/memory-search.ts b/packages/worker/src/mcp/memory/memory-search.ts index 475d1bef70..325dba05ea 100644 --- a/packages/worker/src/mcp/memory/memory-search.ts +++ b/packages/worker/src/mcp/memory/memory-search.ts @@ -14,11 +14,57 @@ import { parseJsonStringArray } from './json-string-array.ts' import { buildMemoryEmbedTextFromRow } from './memory-embed.ts' import { type McpMemoryRow, type MemorySearchMatch } from './types.ts' +/** + * Queries Vectorize for the user's memory vectors (metadata filter on kind, + * userId, status, and optional category) and returns memory ids ordered by + * similarity. Returns an empty list when the vector index is unavailable. + */ +export async function queryMemoryVectorIds(input: { + env: Env + query: string + userId: string + statuses: ReadonlyArray + category?: string | null + topK: number +}): Promise> { + const index = getCapabilityVectorIndex(input.env) + if (!index) return [] + const queryVector = await embedTextForVectorize(input.env, input.query) + const vectorMatches = await index.query(queryVector, { + topK: input.topK, + returnMetadata: 'none', + filter: { + kind: { $eq: 'memory' }, + userId: { $eq: input.userId }, + status: { $in: [...input.statuses] }, + ...(input.category ? { category: { $eq: input.category } } : {}), + }, + }) + const memoryIds: Array = [] + const seen = new Set() + for (const match of vectorMatches.matches) { + if (typeof match.id !== 'string') continue + const memoryId = getRawIdFromPassthroughVectorId({ + prefix: 'memory', + vectorId: match.id, + }) + if (!memoryId || seen.has(memoryId)) continue + seen.add(memoryId) + memoryIds.push(memoryId) + } + return memoryIds +} + export async function searchMemories(input: { env: Env query: string limit: number rows: Array + /** + * Vector-ranked memory ids already fetched from Vectorize (online path). + * When provided, no additional Vectorize query is issued here. + */ + vectorRankedIds?: ReadonlyArray }): Promise<{ matches: Array; offline: boolean }> { const query = input.query.trim() const rowsById = new Map(input.rows.map((row) => [row.id, row] as const)) @@ -39,7 +85,11 @@ export async function searchMemories(input: { ) let vectorOrder: Array - if (offline) { + if (input.vectorRankedIds) { + const fromIndex = input.vectorRankedIds.filter((id) => rowsById.has(id)) + const fromIndexSet = new Set(fromIndex) + vectorOrder = [...fromIndex, ...ids.filter((id) => !fromIndexSet.has(id))] + } else if (offline) { const queryVector = deterministicEmbedding(query) const similarityById = Object.fromEntries( ids.map((id) => { diff --git a/packages/worker/src/mcp/memory/repo.ts b/packages/worker/src/mcp/memory/repo.ts index 97085feecd..4af41724e1 100644 --- a/packages/worker/src/mcp/memory/repo.ts +++ b/packages/worker/src/mcp/memory/repo.ts @@ -181,16 +181,53 @@ export async function listMemoriesByUserId( return (results ?? []).map(mapMemoryRow) } -export async function listAllMemories(input: { +export async function listMemoriesByIds( + db: D1Database, + userId: string, + options: { + ids: Array + statuses?: Array + }, +): Promise> { + if (options.ids.length === 0) return [] + const statuses = + options.statuses && options.statuses.length > 0 ? options.statuses : null + const statusClause = statuses + ? `AND status IN (${statuses.map(() => '?').join(', ')})` + : '' + const idPlaceholders = options.ids.map(() => '?').join(', ') + const { results } = await db + .prepare( + `SELECT id, user_id, category, status, subject, summary, details, tags_json, + source_uris_json, dedupe_key, created_at, updated_at, last_accessed_at, deleted_at + FROM mcp_memories + WHERE user_id = ? + AND id IN (${idPlaceholders}) + ${statusClause}`, + ) + .bind(userId, ...options.ids, ...(statuses ?? [])) + .all>() + return (results ?? []).map(mapMemoryRow) +} + +// Keyset-paged listing across all users, used by maintenance reindex so a +// single query never loads the whole table. Pass the last row id of the +// previous page (or null for the first page). +export async function listMemoriesPage(input: { db: D1Database + afterId: string | null + limit: number }): Promise> { const { results } = await input.db .prepare( `SELECT id, user_id, category, status, subject, summary, details, tags_json, source_uris_json, dedupe_key, created_at, updated_at, last_accessed_at, deleted_at FROM mcp_memories - ORDER BY updated_at DESC`, + WHERE id > ? + ORDER BY id + LIMIT ?`, ) + .bind(input.afterId ?? '', input.limit) .all>() return (results ?? []).map(mapMemoryRow) } diff --git a/packages/worker/src/mcp/memory/service.node.test.ts b/packages/worker/src/mcp/memory/service.node.test.ts index 34280689db..6f1df3d4f5 100644 --- a/packages/worker/src/mcp/memory/service.node.test.ts +++ b/packages/worker/src/mcp/memory/service.node.test.ts @@ -1,4 +1,8 @@ import { expect, test } from 'vitest' +import { + CAPABILITY_EMBEDDING_DIMENSIONS, + deterministicEmbedding, +} from '#mcp/capabilities/capability-search.ts' import { deleteMemory, getMemory, @@ -43,6 +47,25 @@ function createMemoryTestDb() { return null }, async all() { + if ( + normalizedQuery.includes('from mcp_memories') && + normalizedQuery.includes('and id in (') + ) { + const [userId, ...rest] = params as Array + const statusValues = new Set(['active', 'archived', 'deleted']) + const ids = rest.filter((value) => !statusValues.has(value)) + const statuses = rest.filter((value) => statusValues.has(value)) + const rows = [...memories.values()] + .filter((row) => row.user_id === userId) + .filter((row) => ids.includes(row.id)) + .filter((row) => + statuses.length === 0 + ? true + : statuses.includes(row.status), + ) + .map((row) => ({ ...row })) + return { results: rows as Array, meta: { changes: 0 } } + } if ( normalizedQuery.includes('from mcp_memories') && normalizedQuery.includes('order by updated_at desc') @@ -244,6 +267,21 @@ const env = (db: D1Database) => APP_DB: db, }) satisfies Pick as Pick +function createDeterministicAiBinding(): Ai { + return { + async run(...args: Array) { + const input = args[1] as { text?: unknown } + const texts = Array.isArray(input.text) + ? input.text.map(String) + : [String(input.text ?? '')] + return { + data: texts.map((text) => deterministicEmbedding(text)), + shape: [texts.length, CAPABILITY_EMBEDDING_DIMENSIONS], + } + }, + } as unknown as Ai +} + test('memory service upserts, verifies, and soft deletes', async () => { const testDb = createMemoryTestDb() const runtimeEnv = env(testDb.db) @@ -484,3 +522,98 @@ test('memory service rejects invalid source uris and tolerates missing stored va expect(loaded?.sourceUris).toEqual([]) }) + +test('memory search online queries Vectorize first and hydrates vector hits by id', async () => { + const testDb = createMemoryTestDb() + const baseTime = Date.parse('2026-06-01T00:00:00.000Z') + function seedMemory(input: { + id: string + subject: string + summary: string + ageMinutes: number + }) { + const timestamp = new Date( + baseTime - input.ageMinutes * 60_000, + ).toISOString() + testDb.memories.set(input.id, { + id: input.id, + user_id: 'user-123', + category: null, + status: 'active', + subject: input.subject, + summary: input.summary, + details: '', + tags_json: '[]', + source_uris_json: '[]', + dedupe_key: null, + created_at: timestamp, + updated_at: timestamp, + last_accessed_at: null, + deleted_at: null, + }) + } + seedMemory({ + id: 'memory-recent', + subject: 'Deployment window', + summary: 'User prefers deployments after 4pm.', + ageMinutes: 0, + }) + // Enough newer filler rows that the old memory falls outside the bounded + // lexical candidate set (most recent 50 rows). + for (let index = 0; index < 60; index += 1) { + seedMemory({ + id: `filler-${index}`, + subject: `Filler note ${index}`, + summary: 'Unrelated grocery reminder.', + ageMinutes: index + 1, + }) + } + seedMemory({ + id: 'memory-old-vector-hit', + subject: 'Deployment window history', + summary: 'Old deployment window memory only reachable via Vectorize.', + ageMinutes: 10_000, + }) + + const vectorQueryCalls: Array<{ + topK: number + filter?: Record + }> = [] + const runtimeEnv = { + APP_DB: testDb.db, + SENTRY_ENVIRONMENT: 'production', + AI: createDeterministicAiBinding(), + CAPABILITY_VECTOR_INDEX: { + async query( + _values: Array, + options: { topK: number; filter?: Record }, + ) { + vectorQueryCalls.push(options) + return { + matches: [{ id: 'memory_memory-old-vector-hit', score: 0.92 }], + } + }, + }, + } as unknown as Pick + + const result = await searchMemoryRecords({ + env: runtimeEnv, + userId: 'user-123', + query: 'deployment window', + limit: 5, + }) + + // One Vectorize query with user/kind/status metadata filters; no second + // query happens inside the fusion step. + expect(vectorQueryCalls).toHaveLength(1) + expect(vectorQueryCalls[0]).toMatchObject({ + filter: { + kind: { $eq: 'memory' }, + userId: { $eq: 'user-123' }, + status: { $in: ['active', 'archived'] }, + }, + }) + const matchedIds = result.matches.map((match) => match.id) + expect(matchedIds).toContain('memory-old-vector-hit') + expect(matchedIds).toContain('memory-recent') +}) diff --git a/packages/worker/src/mcp/memory/service.ts b/packages/worker/src/mcp/memory/service.ts index b10848cd0a..ebc3aa3708 100644 --- a/packages/worker/src/mcp/memory/service.ts +++ b/packages/worker/src/mcp/memory/service.ts @@ -1,9 +1,11 @@ import { type StorageContext } from '#mcp/storage.ts' +import { isCapabilitySearchOffline } from '#mcp/capabilities/capability-search.ts' import { deleteMemory as deleteMemoryRow, getConversationSuppressions, getMemoryById, insertMemory, + listMemoriesByIds, listMemoriesByUserId, pruneExpiredConversationSuppressions, touchMemoryAccessedAt, @@ -13,7 +15,7 @@ import { import { parseJsonStringArray } from './json-string-array.ts' import { buildMemoryEmbedText } from './memory-embed.ts' import { deleteMemoryVector, upsertMemoryVector } from './memory-vectorize.ts' -import { searchMemories } from './memory-search.ts' +import { queryMemoryVectorIds, searchMemories } from './memory-search.ts' import { type McpMemoryRow, type MemoryRecord, @@ -30,6 +32,11 @@ const maxTagCount = 16 const maxSourceUriLength = 2_048 const maxSourceUriCount = 12 const defaultSuppressionTtlMs = 30 * 24 * 60 * 60 * 1_000 +/** Offline (no Vectorize): lexical + deterministic ranking over recent rows. */ +const memoryOfflineCandidateLimit = 200 +/** Online: bounded recent-row set fused with Vectorize top hits via RRF. */ +const memoryLexicalCandidateLimit = 50 +const memoryVectorTopKCap = 100 export const memorySearchMutationGuidance = 'This result is mutable for the signed-in user. Use this exact id as memory_id with meta_memory_upsert or meta_memory_delete after meta_memory_verify.' @@ -331,26 +338,65 @@ export async function searchMemoryRecords(input: MemorySearchInput): Promise<{ if (!query) { return { query: '', matches: [], suppressedCount: 0 } } - const rows = await listMemoriesByUserId(input.env.APP_DB, input.userId, { - statuses: input.includeDeleted - ? ['active', 'archived', 'deleted'] - : ['active', 'archived'], - limit: 200, - }) - const filteredRows = rows.filter((row) => { - if (!input.category) return true - return ( - row.category === - normalizeOptionalString(input.category, maxCategoryLength) + const statuses: Array = input.includeDeleted + ? ['active', 'archived', 'deleted'] + : ['active', 'archived'] + const targetLimit = normalizeLimit(input.limit) + const category = normalizeOptionalString(input.category, maxCategoryLength) + const offline = isCapabilitySearchOffline(input.env as Env) + + let rows: Array + let vectorRankedIds: Array | undefined + if (offline) { + rows = await listMemoriesByUserId(input.env.APP_DB, input.userId, { + statuses, + limit: memoryOfflineCandidateLimit, + }) + } else { + // Vectorize-first: query the index with metadata filters, hydrate only + // the top-hit rows from D1 by id, and merge with a bounded set of + // recent rows as the lexical candidate list. + const [recentRows, rankedIds] = await Promise.all([ + listMemoriesByUserId(input.env.APP_DB, input.userId, { + statuses, + limit: memoryLexicalCandidateLimit, + }), + queryMemoryVectorIds({ + env: input.env as Env, + query, + userId: input.userId, + statuses, + category, + topK: Math.min( + Math.max(memoryLexicalCandidateLimit, targetLimit * 5), + memoryVectorTopKCap, + ), + }), + ]) + vectorRankedIds = rankedIds + const recentIds = new Set(recentRows.map((row) => row.id)) + const missingIds = rankedIds.filter((id) => !recentIds.has(id)) + const hydratedRows = await listMemoriesByIds( + input.env.APP_DB, + input.userId, + { + ids: missingIds, + statuses, + }, ) + rows = [...recentRows, ...hydratedRows] + } + const filteredRows = rows.filter((row) => { + if (!category) return true + return row.category === category }) - const targetLimit = normalizeLimit(input.limit) const rankedLimit = Math.min(filteredRows.length, targetLimit * 5) const { matches } = await searchMemories({ env: input.env as Env, query, limit: rankedLimit, rows: filteredRows, + ...(vectorRankedIds ? { vectorRankedIds } : {}), }) const filtered = await filterSuppressedMatches({ db: input.env.APP_DB, diff --git a/packages/worker/src/mcp/tools/search.node.test.ts b/packages/worker/src/mcp/tools/search.node.test.ts index 63284b42ba..d40e1201b7 100644 --- a/packages/worker/src/mcp/tools/search.node.test.ts +++ b/packages/worker/src/mcp/tools/search.node.test.ts @@ -660,7 +660,7 @@ test('searchUnified annotates high-confidence package action matches', async () }) }) -test('buildSavedPackageSearchRows hydrates README and export JSDoc search signals', async () => { +test('buildSavedPackageSearchRows defers source loading and hydrates only top matches', async () => { const readmeBody = 'Package-first trace and debug workflow for failed processor service storage automation.' const exportDescription = 'Trace failed processor service storage writes.' @@ -689,6 +689,7 @@ test('buildSavedPackageSearchRows hydrates README and export JSDoc search signal export declare function traceProcessorFailure(messageId: string): Promise `, } + sourceMocks.loadPackageSourceBySourceId.mockClear() sourceMocks.loadPackageSourceBySourceId.mockResolvedValueOnce({ source: { id: 'source-trace' }, manifest, @@ -716,49 +717,22 @@ export declare function traceProcessorFailure(messageId: string): Promise ], }) + // Building rows only projects cheap D1 fields; no source snapshots load. expect(rows.warnings).toEqual([]) + expect(sourceMocks.loadPackageSourceBySourceId).not.toHaveBeenCalled() expect(rows.rows[0]).toMatchObject({ - readmeSnippet: { - path: 'README.md', - snippet: expect.stringContaining(readmeBody), - truncated: false, - }, - projection: { - exports: [ - expect.objectContaining({ - description: exportDescription, - }), - ], - }, + readmeSnippet: null, + projection: expect.objectContaining({ + kodyId: 'trace-package', + exports: [], + }), }) - const registry = buildCapabilityRegistry([ - { - name: 'storage', - description: 'Storage primitives', - capabilities: [ - { - name: 'storage_query', - domain: 'storage', - description: 'Low-level storage query for processor records.', - keywords: ['storage', 'service', 'processor', 'failed'], - readOnly: true, - idempotent: true, - destructive: false, - inputSchema: { - type: 'object', - properties: {}, - }, - handler: async () => null, - }, - ], - }, - ]) const result = await searchUnified({ env: {} as Env, - query: 'trace failed processor service storage', + query: 'trace package', limit: 3, - registry, + registry: buildCapabilityRegistry([]), optionalRows: { packageRows: rows.rows, userSecretRows: [], @@ -766,40 +740,162 @@ export declare function traceProcessorFailure(messageId: string): Promise }, }) - expect(result.matches).toEqual( + // The returned top match is hydrated with README and export metadata. + const packageMatch = result.matches.find((match) => match.type === 'package') + expect(packageMatch).toMatchObject({ + type: 'package', + kodyId: 'trace-package', + readmeSnippet: { + path: 'README.md', + snippet: expect.stringContaining(readmeBody), + truncated: false, + }, + }) + expect(packageMatch?.actionMatches).toEqual( expect.arrayContaining([ expect.objectContaining({ - type: 'package', - kodyId: 'trace-package', + subpath: './trace-processor', + functions: [ + expect.objectContaining({ + name: 'traceProcessorFailure', + }), + ], }), ]), ) + expect(sourceMocks.loadPackageSourceBySourceId).toHaveBeenCalledTimes(1) + // Hydration failures degrade to the lean match instead of failing search. sourceMocks.loadPackageSourceBySourceId.mockRejectedValueOnce( new Error('missing-source'), ) - await expect( - buildSavedPackageSearchRows({ - env: {} as Env, - baseUrl: 'http://localhost', - userId: 'user-123', - records: [ - { - id: 'package-123', - userId: 'user-123', - name: '@kody/observed', - kodyId: 'observed', - description: 'Observed package', - tags: ['observed'], - searchText: 'search text', - sourceId: 'missing-source', - hasApp: true, - createdAt: '2026-03-24T00:00:00.000Z', - updatedAt: '2026-03-24T00:00:00.000Z', - }, + const failedHydration = await buildSavedPackageSearchRows({ + env: {} as Env, + baseUrl: 'http://localhost', + userId: 'user-123', + records: [ + { + id: 'package-123', + userId: 'user-123', + name: '@kody/observed', + kodyId: 'observed', + description: 'Observed package', + tags: ['observed'], + searchText: 'search text', + sourceId: 'missing-source', + hasApp: true, + createdAt: '2026-03-24T00:00:00.000Z', + updatedAt: '2026-03-24T00:00:00.000Z', + }, + ], + }) + const degraded = await searchUnified({ + env: {} as Env, + query: 'observed package', + limit: 3, + registry: buildCapabilityRegistry([]), + optionalRows: { + packageRows: failedHydration.rows, + userSecretRows: [], + userValueRows: [], + }, + }) + expect(degraded.matches).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: 'package', + kodyId: 'observed', + readmeSnippet: null, + }), + ]), + ) +}) + +test('searchUnified ranks packages via user-filtered package vectors when Vectorize is online', async () => { + const capturedFilters: Array | undefined> = [] + const env = { + SENTRY_ENVIRONMENT: 'production', + AI: createDeterministicAiBinding(), + CAPABILITY_VECTOR_INDEX: { + async query( + _values: Array, + options: { filter?: Record }, + ) { + capturedFilters.push(options.filter) + const kind = (options.filter as { kind?: { $eq?: string } } | undefined) + ?.kind?.$eq + if (kind === 'package') { + return { matches: [{ id: 'package_pkg-vector', score: 0.91 }] } + } + return { matches: [] } + }, + }, + } as unknown as Env + function leanPackageRow(id: string, name: string): PackageSearchRow { + return { + record: { + id, + userId: 'user-1', + name, + kodyId: name, + description: 'automation helpers', + tags: [], + searchText: null, + sourceId: `source-${id}`, + hasApp: false, + createdAt: '2026-04-20T00:00:00.000Z', + updatedAt: '2026-04-20T00:00:00.000Z', + }, + projection: { + name, + kodyId: name, + description: 'automation helpers', + tags: [], + searchText: null, + hasApp: false, + appEntry: null, + exports: [], + jobs: [], + services: [], + subscriptions: [], + retrievers: [], + }, + readmeSnippet: null, + } + } + + const result = await searchUnified({ + env, + query: 'summarize inbox threads', + limit: 5, + registry: buildCapabilityRegistry([]), + optionalRows: { + packageRows: [ + leanPackageRow('pkg-vector', 'semantic-match'), + leanPackageRow('pkg-other', 'unrelated-package'), ], + userSecretRows: [], + userValueRows: [], + }, + }) + + expect(result.offline).toBe(false) + expect(capturedFilters).toContainEqual( + expect.objectContaining({ + kind: { $eq: 'package' }, + userId: { $eq: 'user-1' }, }), - ).rejects.toThrow('missing-source') + ) + expect( + result.matches.some( + (match) => match.type === 'package' && match.packageId === 'pkg-vector', + ), + ).toBe(true) + expect( + result.matches.some( + (match) => match.type === 'package' && match.packageId === 'pkg-other', + ), + ).toBe(false) }) test('down remote connector statuses surface only disconnected connectors for signed-in users', async () => { diff --git a/packages/worker/src/mcp/tools/search.ts b/packages/worker/src/mcp/tools/search.ts index f1875b792f..af6eda87ec 100644 --- a/packages/worker/src/mcp/tools/search.ts +++ b/packages/worker/src/mcp/tools/search.ts @@ -8,12 +8,17 @@ import { } from '#mcp/capabilities/integrations/integration-shared.ts' import { getCapabilityRegistryForContext } from '#mcp/capabilities/registry.ts' import { + CAPABILITY_SEARCH_RRF_K, blendLexicalAndVectorScore, cosineSimilarity, deterministicEmbedding, + embedTextForVectorize, + getCapabilityVectorIndex, isCapabilitySearchOffline, lexicalScore, + reciprocalRankFusion, searchCapabilities, + sortIdsByScore, } from '#mcp/capabilities/capability-search.ts' import { listUserSecretsForSearch } from '#mcp/secrets/service.ts' import { type SecretSearchRow } from '#mcp/secrets/types.ts' @@ -29,6 +34,7 @@ import { type ValueMetadata } from '#mcp/values/types.ts' import { getSavedPackageByKodyId, listSavedPackagesByUserId, + savedPackageVectorId, } from '#worker/package-registry/repo.ts' import { buildPackageSearchDocument, @@ -87,11 +93,22 @@ const maxTokens = 6_000 const maxChars = maxTokens * charsPerToken const defaultSearchLimit = 15 const defaultMaxResponseSize = 4_000 +/** Upper bound on package candidates kept after lexical/vector rank fusion. */ +const maxFusedPackageCandidates = 100 export type PackageSearchRow = { record: Awaited>[number] projection: PackageSearchProjection readmeSnippet?: PackageReadmeSnippet | null + /** + * Lazily loads full package source detail (export metadata + README). + * Only invoked for the bounded set of top-ranked package matches so broad + * searches never load every package's published source. + */ + hydrate?: () => Promise<{ + projection: PackageSearchProjection + readmeSnippet: PackageReadmeSnippet | null + }> } export type OptionalSearchRowsResult = { @@ -188,30 +205,65 @@ export type BuildSavedPackageSearchRowsResult = { warnings: Array } +/** + * Cheap projection built entirely from the `saved_packages` D1 row. Heavier + * fields (exports, jobs, README) stay empty until a row is hydrated for a + * top-ranked match. + */ +function buildLeanPackageSearchProjection( + record: Awaited>[number], +): PackageSearchProjection { + return { + name: record.name, + kodyId: record.kodyId, + description: record.description, + tags: record.tags, + searchText: record.searchText, + hasApp: record.hasApp, + appEntry: null, + exports: [], + jobs: [], + services: [], + subscriptions: [], + retrievers: [], + } +} + export async function buildSavedPackageSearchRows(input: { env: Env baseUrl: string userId: string records: Array>[number]> }): Promise { - const rows = await Promise.all( - input.records.map(async (record) => { - const loaded = await loadPackageSourceBySourceId({ - env: input.env, - baseUrl: input.baseUrl, - userId: input.userId, - sourceId: record.sourceId, - }) - return { - record, - projection: buildPackageSearchProjection(loaded.manifest, loaded.files), - readmeSnippet: buildPackageReadmeSnippet({ - files: loaded.files, - maxChars: 1_000, - }), - } - }), - ) + const rows = input.records.map((record) => { + let hydration: Promise<{ + projection: PackageSearchProjection + readmeSnippet: PackageReadmeSnippet | null + }> | null = null + return { + record, + projection: buildLeanPackageSearchProjection(record), + readmeSnippet: null, + hydrate: () => { + hydration ??= loadPackageSourceBySourceId({ + env: input.env, + baseUrl: input.baseUrl, + userId: input.userId, + sourceId: record.sourceId, + }).then((loaded) => ({ + projection: buildPackageSearchProjection( + loaded.manifest, + loaded.files, + ), + readmeSnippet: buildPackageReadmeSnippet({ + files: loaded.files, + maxChars: 1_000, + }), + })) + return hydration + }, + } satisfies PackageSearchRow + }) return { rows, warnings: [] } } @@ -874,13 +926,67 @@ async function buildCapabilityCandidates(input: { .filter((candidate) => candidate.scoreComponents.base > 0) } -function buildPackageCandidates(input: { +/** + * Queries Vectorize for this user's `package_{id}` vectors (written at publish + * time) and returns similarity scores keyed by saved package record id. + * Returns null when the vector index is unavailable so callers fall back to + * deterministic offline scoring. + */ +async function queryPackageVectorScores(input: { + env: Env + query: string + rows: Array + limit: number +}): Promise | null> { + const index = getCapabilityVectorIndex(input.env) + if (!index) return null + const userIds = Array.from( + new Set(input.rows.map((row) => row.record.userId).filter(Boolean)), + ) + const recordIdByVectorId = new Map( + input.rows.map( + (row) => [savedPackageVectorId(row.record.id), row.record.id] as const, + ), + ) + const queryVector = await embedTextForVectorize(input.env, input.query) + const topK = Math.min(Math.max(input.rows.length, input.limit * 5), 100) + const vectorMatches = await index.query(queryVector, { + topK, + returnMetadata: 'none', + filter: { + kind: { $eq: 'package' }, + ...(userIds.length === 1 ? { userId: { $eq: userIds[0] } } : {}), + }, + }) + const scores = new Map() + for (const match of vectorMatches.matches) { + if (typeof match.id !== 'string') continue + const recordId = recordIdByVectorId.get(match.id) + if (!recordId || scores.has(recordId)) continue + scores.set(recordId, match.score) + } + return scores +} + +async function buildPackageCandidates(input: { + env: Env query: string rows: Array queryEmbedding: ReadonlyArray -}): Array { + limit: number + offline: boolean +}): Promise> { + if (input.rows.length === 0) return [] const meaningfulTokens = extractMeaningfulSearchTokens(input.query) - return input.rows + const vectorScoresByRecordId = input.offline + ? null + : await queryPackageVectorScores({ + env: input.env, + query: input.query, + rows: input.rows, + limit: input.limit, + }) + const candidates = input.rows .map((entry) => { const exports = Array.isArray(entry.projection.exports) ? entry.projection.exports @@ -910,10 +1016,14 @@ function buildPackageCandidates(input: { .filter((value) => value.trim().length > 0) .join('\n') const lexical = lexicalScore(input.query, document) - const vector = cosineSimilarity( - input.queryEmbedding, - deterministicEmbedding(document), - ) + // Online, Vectorize similarity replaces the per-package deterministic + // embedding so broad searches never embed every package document. + const vector = vectorScoresByRecordId + ? (vectorScoresByRecordId.get(entry.record.id) ?? 0) + : cosineSimilarity( + input.queryEmbedding, + deterministicEmbedding(document), + ) return { match: { type: 'package' as const, @@ -973,6 +1083,42 @@ function buildPackageCandidates(input: { } }) .filter((candidate) => candidate.scoreComponents.base > 0) + if (!vectorScoresByRecordId) { + return candidates + } + // Online: fuse lexical and vector rankings (RRF, mirroring + // capability-search) to bound the package candidate set. + const candidateIds = candidates.map((candidate) => candidate.match.packageId) + const lexicalById = new Map( + candidates.map( + (candidate) => + [candidate.match.packageId, candidate.scoreComponents.lexical] as const, + ), + ) + const lexicalOrder = sortIdsByScore( + candidateIds, + (id) => lexicalById.get(id) ?? 0, + ) + const vectorOrder = sortIdsByScore( + candidateIds, + (id) => vectorScoresByRecordId.get(id) ?? 0, + ) + const fused = reciprocalRankFusion( + [lexicalOrder, vectorOrder], + CAPABILITY_SEARCH_RRF_K, + ) + const keptIds = new Set( + sortIdsByScore(candidateIds, (id) => fused.get(id) ?? 0).slice( + 0, + Math.min( + maxFusedPackageCandidates, + Math.max(input.limit * 5, input.limit), + ), + ), + ) + return candidates.filter((candidate) => + keptIds.has(candidate.match.packageId), + ) } function buildRetrieverResultCandidates(input: { @@ -1182,6 +1328,49 @@ function capabilityMatchToCandidate( } } +/** + * Loads README snippets and export metadata for the bounded set of package + * matches actually being returned (at most `limit` rows), instead of loading + * every saved package's full source up front. + */ +async function hydrateTopPackageMatches(input: { + query: string + matches: Array + rows: Array +}): Promise { + const rowsByRecordId = new Map( + input.rows.map((row) => [row.record.id, row] as const), + ) + const packageMatches = input.matches.flatMap((match) => + match.type === 'package' ? [match] : [], + ) + if (packageMatches.length === 0) return + const meaningfulTokens = extractMeaningfulSearchTokens(input.query) + await Promise.all( + packageMatches.map(async (match) => { + const row = rowsByRecordId.get(match.packageId) + if (!row?.hydrate) return + try { + const hydrated = await row.hydrate() + match.readmeSnippet = hydrated.readmeSnippet + match.actionMatches = buildPackageActionMatches({ + query: input.query, + meaningfulTokens, + exports: hydrated.projection.exports, + }) + } catch (error) { + console.warn( + JSON.stringify({ + message: 'package search match hydration failed', + packageId: match.packageId, + error: error instanceof Error ? error.message : String(error), + }), + ) + } + }), + ) +} + export async function searchUnified(input: { env: Env query: string @@ -1244,11 +1433,14 @@ export async function searchUnified(input: { env: input.env, registry: input.registry, })), - ...buildPackageCandidates({ + ...(await buildPackageCandidates({ + env: input.env, query: intent.normalizedQuery, rows: input.optionalRows.packageRows, queryEmbedding, - }), + limit, + offline, + })), ...buildValueCandidates({ query: intent.normalizedQuery, rows: input.optionalRows.userValueRows, @@ -1278,6 +1470,11 @@ export async function searchUnified(input: { limit, }) const matches = reranked.map((candidate) => candidate.match) + await hydrateTopPackageMatches({ + query: intent.normalizedQuery, + matches, + rows: input.optionalRows.packageRows, + }) const rerankingMs = Math.max( 0, Math.round(performance.now() - rerankingStart), From f1850a1fb144710fd8f1b7a4c4e8e777606317c0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:31 +0000 Subject: [PATCH 06/14] Paginate account export and maintenance reindex at the SQL level Account export reads keyset-paged rowid windows (500 rows per query) instead of whole tables, and the memory/job/package reindex drivers process keyset-paged batches of 200 with merged result reporting. Co-authored-by: Kent C. Dodds --- .../src/app/account-export.node.test.ts | 77 +++- packages/worker/src/app/account-export.ts | 356 ++++++++++-------- .../worker/src/jobs/job-reindex.node.test.ts | 107 ++++-- packages/worker/src/jobs/job-reindex.ts | 76 ++-- .../src/mcp/capabilities/reindex-batches.ts | 26 ++ .../mcp/memory/memory-reindex.node.test.ts | 107 ++++++ .../worker/src/mcp/memory/memory-reindex.ts | 57 +-- .../package-reindex.node.test.ts | 75 +++- .../src/package-registry/package-reindex.ts | 128 ++++--- packages/worker/src/package-registry/repo.ts | 14 +- 10 files changed, 721 insertions(+), 302 deletions(-) create mode 100644 packages/worker/src/mcp/memory/memory-reindex.node.test.ts diff --git a/packages/worker/src/app/account-export.node.test.ts b/packages/worker/src/app/account-export.node.test.ts index d868ac544d..b3d91f9631 100644 --- a/packages/worker/src/app/account-export.node.test.ts +++ b/packages/worker/src/app/account-export.node.test.ts @@ -20,7 +20,10 @@ function applyMigrations(db: DatabaseSync) { } } -function createD1FromSqlite(db: DatabaseSync) { +function createD1FromSqlite( + db: DatabaseSync, + options?: { onQueryRows?: (rowCount: number) => void }, +) { return { prepare(query: string) { return { @@ -29,6 +32,7 @@ function createD1FromSqlite(db: DatabaseSync) { async all() { const statement = db.prepare(query) const rows = statement.all(...params) as Array + options?.onQueryRows?.(rows.length) return { results: rows, meta: { changes: 0 } } }, async first() { @@ -45,6 +49,7 @@ function createD1FromSqlite(db: DatabaseSync) { async all() { const statement = db.prepare(query) const rows = statement.all() as Array + options?.onQueryRows?.(rows.length) return { results: rows, meta: { changes: 0 } } }, async first() { @@ -64,12 +69,14 @@ function createD1FromSqlite(db: DatabaseSync) { } as unknown as D1Database } -function createMigratedDb() { +function createMigratedDb(options?: { + onQueryRows?: (rowCount: number) => void +}) { const sqlite = new DatabaseSync(':memory:') applyMigrations(sqlite) return { sqlite, - db: createD1FromSqlite(sqlite), + db: createD1FromSqlite(sqlite, options), } } @@ -373,9 +380,9 @@ test('createAccountExport records partial-failure warnings and section paginatio table: 'value_entries', pageSize: 1, }) - expect(page.items).toHaveLength(1) + expect(page.items).toEqual([expect.objectContaining({ name: 'first' })]) expect(page.truncated).toBe(true) - expect(page.nextStartAfter).toBe('value-bucket-a:first') + expect(page.nextStartAfter).not.toBeNull() const nextPage = await readAccountExportSection({ env, dbUserId: 1, @@ -387,4 +394,64 @@ test('createAccountExport records partial-failure warnings and section paginatio }) expect(nextPage.items).toEqual([expect.objectContaining({ name: 'second' })]) expect(nextPage.truncated).toBe(false) + expect(nextPage.nextStartAfter).toBeNull() +}) + +test('D1 export reads large tables in bounded keyset pages', async () => { + const rowCounts: Array = [] + const { sqlite, db } = createMigratedDb({ + onQueryRows: (rowCount) => rowCounts.push(rowCount), + }) + sqlite.exec(` + INSERT INTO users (id, username, email, password_hash, created_at, updated_at, email_verified_at) + VALUES (1, 'user-a', 'a@example.com', 'password-hash-a', '2026-07-05', '2026-07-05', '2026-07-05'); + `) + const totalMessages = 1201 + const insert = sqlite.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, from_address, subject, processing_status, created_at, updated_at + ) VALUES (?, 'inbound', 'user-aaa', 'sender@example.net', ?, 'stored', '2026-07-05', '2026-07-05')`, + ) + for (let index = 0; index < totalMessages; index += 1) { + insert.run(`message-${String(index).padStart(4, '0')}`, `Mail ${index}`) + } + + const accountExport = await createAccountExport({ + env: { APP_DB: db } as Env, + dbUserId: 1, + mcpUserId: 'user-aaa', + generatedAt: '2026-07-05T00:00:00.000Z', + }) + expect(accountExport.d1.email_messages.rows).toHaveLength(totalMessages) + expect(accountExport.manifest.sections['d1.email_messages']?.count).toBe( + totalMessages, + ) + // The keyset page size is 500 (+1 lookahead row), so no single query may + // return the whole table. + expect(Math.max(...rowCounts)).toBeLessThanOrEqual(501) + + rowCounts.length = 0 + const seenIds = new Set() + let startAfter: string | undefined + let pages = 0 + while (true) { + const page = await readAccountExportSection({ + env: { APP_DB: db } as Env, + dbUserId: 1, + mcpUserId: 'user-aaa', + section: 'd1_table', + table: 'email_messages', + pageSize: 500, + startAfter, + }) + pages += 1 + for (const item of page.items as Array<{ id: string }>) { + seenIds.add(item.id) + } + if (!page.truncated) break + startAfter = page.nextStartAfter ?? undefined + } + expect(pages).toBe(3) + expect(seenIds.size).toBe(totalMessages) + expect(Math.max(...rowCounts)).toBeLessThanOrEqual(501) }) diff --git a/packages/worker/src/app/account-export.ts b/packages/worker/src/app/account-export.ts index f4dd5448cf..7ac15abb00 100644 --- a/packages/worker/src/app/account-export.ts +++ b/packages/worker/src/app/account-export.ts @@ -22,6 +22,12 @@ import { resolveUserStableId } from '#worker/user-id.ts' const accountExportSchemaVersion = 1 const defaultExportPageSize = 100 const maxExportPageSize = 500 +// Full exports stream each D1 table in keyset-paged batches of this size so +// memory stays bounded per query even for very large tables (e.g. mailboxes). +const d1ExportPageSize = 500 +// Internal alias for the SQLite rowid used as the keyset cursor. Stripped from +// exported rows so the export document schema is unchanged. +const exportRowidColumn = '__account_export_rowid' // Secret values and credential-equivalent hashes must never leave Kody through // account export. Secret entries are metadata-only; encrypted payloads stay @@ -248,100 +254,114 @@ function getErrorMessage(error: unknown) { return error instanceof Error ? error.message : String(error) } -function getTargetTableName(target: UserScopedDataTarget) { - return target.kind === 'mcp_memory_suppression' - ? 'mcp_memory_conversation_suppressions' - : target.table +type D1TableCondition = { + condition: string + params: Array } -function buildSelectForTarget(input: { +function buildConditionForTarget(input: { target: UserScopedDataTarget mcpUserId: string dbUserId: number -}) { +}): { table: string; condition: D1TableCondition } { const { target } = input switch (target.kind) { case 'user_id': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE user_id = ?`, - params: [input.mcpUserId], + condition: { + condition: `${target.table}.user_id = ?`, + params: [input.mcpUserId], + }, } } case 'db_user_id': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE user_id = ?`, - params: [input.dbUserId], + condition: { + condition: `${target.table}.user_id = ?`, + params: [input.dbUserId], + }, } } case 'db_user_target': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE target = ?`, - params: [String(input.dbUserId)], + condition: { + condition: `${target.table}.target = ?`, + params: [String(input.dbUserId)], + }, } } case 'user_columns': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE ${target.columns - .map((column) => `${column} = ?`) - .join(' OR ')}`, - params: target.columns.map(() => input.mcpUserId), + condition: { + condition: target.columns + .map((column) => `${target.table}.${column} = ?`) + .join(' OR '), + params: target.columns.map(() => input.mcpUserId), + }, } } case 'null_user_column': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE ${target.matchColumn} = ?`, - params: [input.mcpUserId], + condition: { + condition: `${target.table}.${target.matchColumn} = ?`, + params: [input.mcpUserId], + }, } } case 'replace_user_column': { return { table: target.table, - sql: `SELECT * FROM ${target.table} WHERE ${target.matchColumn} = ?`, - params: [input.mcpUserId], + condition: { + condition: `${target.table}.${target.matchColumn} = ?`, + params: [input.mcpUserId], + }, } } case 'bucket_parent': { return { table: target.table, - sql: `SELECT ${target.table}.* - FROM ${target.table} - INNER JOIN ${target.parentTable} - ON ${target.parentTable}.id = ${target.table}.bucket_id - WHERE ${target.parentTable}.user_id = ?`, - params: [input.mcpUserId], + condition: { + condition: `${target.table}.bucket_id IN ( + SELECT id FROM ${target.parentTable} WHERE user_id = ? + )`, + params: [input.mcpUserId], + }, } } case 'attachment_parent': { return { table: 'email_attachments', - sql: `SELECT email_attachments.* - FROM email_attachments - INNER JOIN email_messages - ON email_messages.id = email_attachments.message_id - WHERE email_messages.user_id = ?`, - params: [input.mcpUserId], + condition: { + condition: `email_attachments.message_id IN ( + SELECT id FROM email_messages WHERE user_id = ? + )`, + params: [input.mcpUserId], + }, } } case 'community_listing_child': { return { table: target.table, - sql: `SELECT * FROM ${target.table} - WHERE ${target.listingColumn} IN ( + condition: { + condition: `${target.table}.${target.listingColumn} IN ( SELECT id FROM community_listings WHERE owner_user_id = ? )`, - params: [input.mcpUserId], + params: [input.mcpUserId], + }, } } case 'mcp_memory_suppression': { return { table: 'mcp_memory_conversation_suppressions', - sql: `SELECT * FROM mcp_memory_conversation_suppressions WHERE user_id = ?`, - params: [input.mcpUserId], + condition: { + condition: `mcp_memory_conversation_suppressions.user_id = ?`, + params: [input.mcpUserId], + }, } } default: { @@ -353,6 +373,31 @@ function buildSelectForTarget(input: { } } +function buildD1TableConditions(input: { + mcpUserId: string + dbUserId: number +}) { + const conditionsByTable = new Map>() + function add(table: string, condition: D1TableCondition) { + const existing = conditionsByTable.get(table) + if (existing) { + existing.push(condition) + return + } + conditionsByTable.set(table, [condition]) + } + add('users', { condition: `users.id = ?`, params: [input.dbUserId] }) + for (const target of accountUserDataTargets) { + const built = buildConditionForTarget({ + target, + mcpUserId: input.mcpUserId, + dbUserId: input.dbUserId, + }) + add(built.table, built.condition) + } + return conditionsByTable +} + function sanitizeRow(table: string, row: Record) { const redactedColumns = redactedColumnsByTable[table] ?? [] const sanitized: Record = {} @@ -395,54 +440,82 @@ async function selectRows>( return result.results ?? [] } -function createD1ExportCollector(input: { +// Reads one keyset page of a table: rows are selected by ascending rowid +// strictly after the cursor, with a SQL LIMIT, so a single query never loads +// more than one page regardless of table size. Conditions for every export +// target of the table are OR-combined into one query, which also removes the +// need for cross-target row deduplication. +async function selectD1TablePage(input: { env: AccountExportEnv - warnings: Array + table: string + conditions: ReadonlyArray + afterRowid: number + limit: number }) { - const tables = new Map() - function getTable(table: string) { - const existing = tables.get(table) - if (existing) return existing - const created: AccountExportD1Table = { - table, - rows: [], - redactedColumns: [], - warnings: [], - } - tables.set(table, created) - return created + const where = input.conditions + .map((condition) => `(${condition.condition})`) + .join(' OR ') + const sql = `SELECT ${input.table}.rowid AS ${exportRowidColumn}, ${input.table}.* + FROM ${input.table} + WHERE (${where}) AND ${input.table}.rowid > ? + ORDER BY ${input.table}.rowid + LIMIT ?` + const params = [ + ...input.conditions.flatMap((condition) => condition.params), + input.afterRowid, + input.limit + 1, + ] + const rawRows = await selectRows(input.env, sql, params) + const truncated = rawRows.length > input.limit + const pageRows = truncated ? rawRows.slice(0, input.limit) : rawRows + let lastRowid = input.afterRowid + const rows: Array> = [] + for (const rawRow of pageRows) { + const { [exportRowidColumn]: rowid, ...columns } = rawRow + lastRowid = Number(rowid) + rows.push(sanitizeRow(input.table, columns)) } - async function collectQuery( - table: string, - sql: string, - params: Array, - ) { - const section = getTable(table) - try { - const rows = await selectRows(input.env, sql, params) - const seen = new Set(section.rows.map(rowDedupeKey)) - const redacted = new Set(section.redactedColumns) - for (const rawRow of rows) { - const sanitized = sanitizeRow(table, rawRow) - for (const column of sanitized.redactedColumns) redacted.add(column) - const key = rowDedupeKey(sanitized.row) - if (seen.has(key)) continue - seen.add(key) - section.rows.push(sanitized.row) + return { rows, lastRowid, truncated } +} + +async function collectD1TableRows(input: { + env: AccountExportEnv + table: string + conditions: ReadonlyArray + warnings: Array +}): Promise { + const section: AccountExportD1Table = { + table: input.table, + rows: [], + redactedColumns: [], + warnings: [], + } + const redacted = new Set() + let afterRowid = 0 + try { + while (true) { + const page = await selectD1TablePage({ + env: input.env, + table: input.table, + conditions: input.conditions, + afterRowid, + limit: d1ExportPageSize, + }) + for (const entry of page.rows) { + for (const column of entry.redactedColumns) redacted.add(column) + section.rows.push(entry.row) } - section.rows = sortRowsByDedupeKey(section.rows) - section.redactedColumns = Array.from(redacted).sort() - } catch (error) { - const warning = `D1 export failed for ${table}: ${getErrorMessage(error)}` - section.warnings.push(warning) - input.warnings.push(warning) + if (!page.truncated) break + afterRowid = page.lastRowid } + } catch (error) { + const warning = `D1 export failed for ${input.table}: ${getErrorMessage(error)}` + section.warnings.push(warning) + input.warnings.push(warning) } - return { - collectQuery, - getTable, - tables, - } + section.rows = sortRowsByDedupeKey(section.rows) + section.redactedColumns = Array.from(redacted).sort() + return section } async function listUserStorageIds(env: Env, userId: string) { @@ -735,80 +808,72 @@ async function collectD1Tables(input: { mcpUserId: string warnings: Array }) { - const collector = createD1ExportCollector(input) - - await collector.collectQuery('users', `SELECT * FROM users WHERE id = ?`, [ - input.dbUserId, - ]) - for (const target of accountUserDataTargets) { - const query = buildSelectForTarget({ - target, - mcpUserId: input.mcpUserId, - dbUserId: input.dbUserId, - }) - await collector.collectQuery( - getTargetTableName(target), - query.sql, - query.params, - ) + const conditionsByTable = buildD1TableConditions({ + mcpUserId: input.mcpUserId, + dbUserId: input.dbUserId, + }) + const tables: Array<[string, AccountExportD1Table]> = [] + for (const [table, conditions] of conditionsByTable) { + tables.push([ + table, + await collectD1TableRows({ + env: input.env, + table, + conditions, + warnings: input.warnings, + }), + ]) } return Object.fromEntries( - Array.from(collector.tables.entries()).sort(([left], [right]) => - left.localeCompare(right), - ), + tables.sort(([left], [right]) => left.localeCompare(right)), ) } -async function collectD1Table(input: { +function parseRowidCursor(startAfter: string | undefined) { + if (!startAfter) return 0 + const parsed = Number.parseInt(startAfter, 10) + return Number.isFinite(parsed) && parsed > 0 ? parsed : 0 +} + +async function readD1TableSectionPage(input: { env: AccountExportEnv dbUserId: number mcpUserId: string table: string - warnings: Array -}) { - const collector = createD1ExportCollector(input) - if (input.table === 'users') { - await collector.collectQuery('users', `SELECT * FROM users WHERE id = ?`, [ - input.dbUserId, - ]) - return collector.getTable('users') - } - const targets = accountUserDataTargets.filter( - (target) => getTargetTableName(target) === input.table, - ) - if (targets.length === 0) return null - for (const target of targets) { - const query = buildSelectForTarget({ - target, - mcpUserId: input.mcpUserId, - dbUserId: input.dbUserId, - }) - await collector.collectQuery(input.table, query.sql, query.params) - } - return collector.getTable(input.table) -} - -function paginateRowsByDedupeKey(input: { - rows: ReadonlyArray> pageSize: number | undefined startAfter: string | undefined + warnings: Array }) { + const conditions = buildD1TableConditions({ + mcpUserId: input.mcpUserId, + dbUserId: input.dbUserId, + }).get(input.table) + if (!conditions) return null const pageSize = normalizePageSize(input.pageSize) - const sortedRows = sortRowsByDedupeKey([...input.rows]) - const cursor = input.startAfter - const startIndex = cursor - ? sortedRows.findIndex((row) => rowDedupeKey(row) > cursor) - : 0 - const safeStart = startIndex < 0 ? sortedRows.length : startIndex - const page = sortedRows.slice(safeStart, safeStart + pageSize) - const nextRow = page.at(-1) - const nextIndex = safeStart + page.length - const truncated = nextIndex < sortedRows.length - return { - items: page, - truncated, - nextStartAfter: truncated && nextRow ? rowDedupeKey(nextRow) : null, - pageSize, + try { + const page = await selectD1TablePage({ + env: input.env, + table: input.table, + conditions, + afterRowid: parseRowidCursor(input.startAfter), + limit: pageSize, + }) + return { + items: page.rows.map((entry) => entry.row), + truncated: page.truncated, + nextStartAfter: page.truncated ? String(page.lastRowid) : null, + pageSize, + } + } catch (error) { + input.warnings.push( + `D1 export failed for ${input.table}: ${getErrorMessage(error)}`, + ) + return { + items: [] as Array>, + truncated: false, + nextStartAfter: null, + pageSize, + } } } @@ -1259,21 +1324,18 @@ export async function readAccountExportSection(input: { if (!input.table) { throw new Error('table is required when section is d1_table.') } - const table = await collectD1Table({ + const page = await readD1TableSectionPage({ env: input.env, dbUserId: input.dbUserId, mcpUserId: input.mcpUserId, table: input.table, + pageSize: input.pageSize, + startAfter: input.startAfter, warnings, }) - if (!table) { + if (!page) { throw new Error(`Table "${input.table}" is not part of account export.`) } - const page = paginateRowsByDedupeKey({ - rows: table.rows, - pageSize: input.pageSize, - startAfter: input.startAfter, - }) return { section: input.section, ...page, diff --git a/packages/worker/src/jobs/job-reindex.node.test.ts b/packages/worker/src/jobs/job-reindex.node.test.ts index ac3446c2d9..42183ef4a4 100644 --- a/packages/worker/src/jobs/job-reindex.node.test.ts +++ b/packages/worker/src/jobs/job-reindex.node.test.ts @@ -1,12 +1,13 @@ import { expect, test, vi } from 'vitest' import { jobVectorId } from '#mcp/jobs-vectorize.ts' +import { type JobRecord } from './types.ts' const mockModule = vi.hoisted(() => ({ embedTextForVectorize: vi.fn(), embedTextsForVectorize: vi.fn(), getCapabilityVectorIndex: vi.fn(), isCapabilitySearchOffline: vi.fn(), - listJobs: vi.fn(), + listJobRowsPage: vi.fn(), })) vi.mock('#mcp/capabilities/capability-search.ts', () => ({ @@ -20,8 +21,9 @@ vi.mock('#mcp/capabilities/capability-search.ts', () => ({ mockModule.isCapabilitySearchOffline(...args), })) -vi.mock('./service.ts', () => ({ - listJobs: (...args: Array) => mockModule.listJobs(...args), +vi.mock('./repo.ts', () => ({ + listJobRowsPage: (...args: Array) => + mockModule.listJobRowsPage(...args), })) const { reindexJobVectors } = await import('./job-reindex.ts') @@ -32,37 +34,51 @@ function byteLength(value: string) { return textEncoder.encode(value).length } +function resetMocks() { + mockModule.embedTextsForVectorize.mockReset() + mockModule.getCapabilityVectorIndex.mockReset() + mockModule.isCapabilitySearchOffline.mockReset() + mockModule.listJobRowsPage.mockReset() +} + +function buildJobRow(id: string, userId = 'user-1') { + const record: JobRecord = { + version: 1, + id, + userId, + name: `Job ${id}`, + sourceId: 'source-1', + publishedCommit: null, + storageId: `job:${id}`, + schedule: { type: 'once', runAt: '2026-01-01T00:00:00.000Z' }, + timezone: 'UTC', + enabled: true, + killSwitchEnabled: false, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + nextRunAt: '2026-01-01T00:00:00.000Z', + runCount: 0, + successCount: 0, + errorCount: 0, + runHistory: [], + } + return { id, user_id: userId, record } +} + test('job reindex keeps package job vector ids under the Vectorize limit', async () => { + resetMocks() const rawJobId = 'package-job:b2fda105-005a-4e2b-9f22-1513b6752da2:event-runner' const vectorId = jobVectorId(rawJobId) const upsert = vi.fn(async () => {}) - const env = { - APP_DB: { - prepare: () => ({ - all: async () => ({ results: [{ user_id: 'user-1' }] }), - }), - }, - } as unknown as Env - - mockModule.embedTextsForVectorize.mockReset() - mockModule.getCapabilityVectorIndex.mockReset() - mockModule.isCapabilitySearchOffline.mockReset() - mockModule.listJobs.mockReset() mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) mockModule.isCapabilitySearchOffline.mockReturnValue(false) mockModule.embedTextsForVectorize.mockResolvedValue([[0.4, 0.5, 0.6]]) - mockModule.listJobs.mockResolvedValue([ - { - id: rawJobId, - name: 'Event runner', - scheduleSummary: 'Every hour', - sourceId: 'source-1', - publishedCommit: null, - }, - ]) + mockModule.listJobRowsPage.mockResolvedValue([buildJobRow(rawJobId)]) - await expect(reindexJobVectors(env)).resolves.toEqual({ upserted: 1 }) + await expect(reindexJobVectors({ APP_DB: {} } as Env)).resolves.toEqual({ + upserted: 1, + }) expect(byteLength(`job_${rawJobId}`)).toBeGreaterThan(64) expect(byteLength(vectorId)).toBeLessThanOrEqual(64) @@ -74,3 +90,44 @@ test('job reindex keeps package job vector ids under the Vectorize limit', async }, ]) }) + +test('job reindex walks keyset pages and merges the page results', async () => { + resetMocks() + const upsert = vi.fn(async (_vectors: Array<{ id: string }>) => {}) + mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) + mockModule.isCapabilitySearchOffline.mockReturnValue(false) + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: unknown, texts: Array) => texts.map(() => [0.1]), + ) + // The first page fills the requested limit, forcing a second page fetch. + mockModule.listJobRowsPage.mockImplementationOnce( + async (_db: unknown, input: { afterId: string | null; limit: number }) => + Array.from({ length: input.limit }, (_, index) => + buildJobRow(`job-${String(index).padStart(4, '0')}`), + ), + ) + mockModule.listJobRowsPage.mockImplementationOnce(async () => [ + buildJobRow('job-last'), + ]) + + await expect(reindexJobVectors({ APP_DB: {} } as Env)).resolves.toEqual({ + upserted: 201, + }) + + expect(mockModule.listJobRowsPage).toHaveBeenCalledTimes(2) + expect(mockModule.listJobRowsPage).toHaveBeenNthCalledWith( + 1, + expect.anything(), + { afterId: null, limit: 200 }, + ) + expect(mockModule.listJobRowsPage).toHaveBeenNthCalledWith( + 2, + expect.anything(), + { afterId: 'job-0199', limit: 200 }, + ) + const upsertedIds = upsert.mock.calls.flatMap(([vectors]) => + vectors.map((vector) => vector.id), + ) + expect(upsertedIds).toHaveLength(201) + expect(new Set(upsertedIds).size).toBe(201) +}) diff --git a/packages/worker/src/jobs/job-reindex.ts b/packages/worker/src/jobs/job-reindex.ts index c02ac3c96a..1a3413b990 100644 --- a/packages/worker/src/jobs/job-reindex.ts +++ b/packages/worker/src/jobs/job-reindex.ts @@ -3,12 +3,18 @@ import { isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' import { + mergeVectorReindexResults, reindexVectorCandidates, type VectorReindexResult, } from '#mcp/capabilities/reindex-batches.ts' import { buildJobEmbedText } from '#mcp/jobs-embed.ts' import { jobVectorId } from '#mcp/jobs-vectorize.ts' -import { listJobs } from './service.ts' +import { listJobRowsPage } from './repo.ts' +import { toJobView } from './schedule.ts' + +// Jobs are reindexed in keyset-paged batches so memory stays bounded +// regardless of table size. +const reindexPageSize = 200 export async function reindexJobVectors( env: Env, @@ -21,37 +27,41 @@ export async function reindexJobVectors( return { upserted: 0 } } - const rows = await env.APP_DB.prepare( - `SELECT DISTINCT user_id FROM jobs`, - ).all<{ user_id: string }>() - const userIds = (rows.results ?? []).map((row) => row.user_id).filter(Boolean) - if (userIds.length === 0) return { upserted: 0 } - - const jobs = ( - await Promise.all( - userIds.map((userId) => - listJobs({ + const pageResults: Array = [] + let afterId: string | null = null + while (true) { + const rows = await listJobRowsPage(env.APP_DB, { + afterId, + limit: reindexPageSize, + }) + if (rows.length === 0) break + const candidates = rows + .filter((row) => row.user_id) + .map((row) => { + const view = toJobView(row.record) + return { + id: jobVectorId(row.id), + text: buildJobEmbedText({ + name: view.name, + scheduleSummary: view.scheduleSummary, + sourceId: view.sourceId, + publishedCommit: view.publishedCommit, + }), + metadata: { kind: 'job', userId: row.user_id }, + } + }) + if (candidates.length > 0) { + pageResults.push( + await reindexVectorCandidates({ env, - userId, - }).then((userJobs) => userJobs.map((job) => ({ userId, job }))), - ), - ) - ).flat() - if (jobs.length === 0) return { upserted: 0 } - - return reindexVectorCandidates({ - env, - index, - kind: 'job', - candidates: jobs.map(({ userId, job }) => ({ - id: jobVectorId(job.id), - text: buildJobEmbedText({ - name: job.name, - scheduleSummary: job.scheduleSummary, - sourceId: job.sourceId, - publishedCommit: job.publishedCommit, - }), - metadata: { kind: 'job', userId }, - })), - }) + index, + kind: 'job', + candidates, + }), + ) + } + if (rows.length < reindexPageSize) break + afterId = rows[rows.length - 1]!.id + } + return mergeVectorReindexResults('job', pageResults) } diff --git a/packages/worker/src/mcp/capabilities/reindex-batches.ts b/packages/worker/src/mcp/capabilities/reindex-batches.ts index db7e5a7813..6785437e51 100644 --- a/packages/worker/src/mcp/capabilities/reindex-batches.ts +++ b/packages/worker/src/mcp/capabilities/reindex-batches.ts @@ -26,6 +26,32 @@ export type VectorReindexResult = { const upsertBatchSize = 16 const maxReportedFailures = 20 +// Combines per-page reindex results from keyset-paged drivers into a single +// result with the same shape and messages as a one-shot reindex run. +export function mergeVectorReindexResults( + kind: string, + results: ReadonlyArray, +): VectorReindexResult { + let upserted = 0 + let failed = 0 + const failures: Array = [] + for (const result of results) { + upserted += result.upserted + failed += result.failed ?? 0 + for (const failure of result.failures ?? []) { + if (failures.length < maxReportedFailures) failures.push(failure) + } + } + if (failed === 0) return { upserted } + const message = `${failed} ${kind} vector(s) failed to reindex` + return { + upserted, + failed, + failures, + ...(upserted === 0 ? { error: message } : { warning: message }), + } +} + export async function reindexVectorCandidates(input: { env: Env index: VectorizeIndex diff --git a/packages/worker/src/mcp/memory/memory-reindex.node.test.ts b/packages/worker/src/mcp/memory/memory-reindex.node.test.ts new file mode 100644 index 0000000000..35b1c62641 --- /dev/null +++ b/packages/worker/src/mcp/memory/memory-reindex.node.test.ts @@ -0,0 +1,107 @@ +import { expect, test, vi } from 'vitest' +import { type McpMemoryRow } from './types.ts' + +const mockModule = vi.hoisted(() => ({ + embedTextForVectorize: vi.fn(), + embedTextsForVectorize: vi.fn(), + getCapabilityVectorIndex: vi.fn(), + isCapabilitySearchOffline: vi.fn(), + listMemoriesPage: vi.fn(), +})) + +vi.mock('#mcp/capabilities/capability-search.ts', () => ({ + embedTextForVectorize: (...args: Array) => + mockModule.embedTextForVectorize(...args), + embedTextsForVectorize: (...args: Array) => + mockModule.embedTextsForVectorize(...args), + getCapabilityVectorIndex: (...args: Array) => + mockModule.getCapabilityVectorIndex(...args), + isCapabilitySearchOffline: (...args: Array) => + mockModule.isCapabilitySearchOffline(...args), +})) + +vi.mock('./repo.ts', () => ({ + listMemoriesPage: (...args: Array) => + mockModule.listMemoriesPage(...args), +})) + +const { reindexMemoryVectors } = await import('./memory-reindex.ts') + +function resetMocks() { + mockModule.embedTextsForVectorize.mockReset() + mockModule.getCapabilityVectorIndex.mockReset() + mockModule.isCapabilitySearchOffline.mockReset() + mockModule.listMemoriesPage.mockReset() +} + +function buildMemoryRow(id: string): McpMemoryRow { + return { + id, + user_id: `user-${id}`, + category: null, + status: 'active', + subject: `Subject ${id}`, + summary: `Summary ${id}`, + details: '', + tags_json: '[]', + source_uris_json: '[]', + dedupe_key: null, + created_at: '2026-01-01T00:00:00.000Z', + updated_at: '2026-01-01T00:00:00.000Z', + last_accessed_at: null, + deleted_at: null, + } +} + +test('memory reindex walks keyset pages and merges the page results', async () => { + resetMocks() + const upsert = vi.fn(async (_vectors: Array<{ id: string }>) => {}) + mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) + mockModule.isCapabilitySearchOffline.mockReturnValue(false) + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: unknown, texts: Array) => texts.map(() => [0.1]), + ) + // The first page fills the requested limit, forcing a second page fetch. + mockModule.listMemoriesPage.mockImplementationOnce( + async (input: { afterId: string | null; limit: number }) => + Array.from({ length: input.limit }, (_, index) => + buildMemoryRow(`memory-${String(index).padStart(4, '0')}`), + ), + ) + mockModule.listMemoriesPage.mockImplementationOnce(async () => [ + buildMemoryRow('memory-last'), + ]) + + await expect(reindexMemoryVectors({ APP_DB: {} } as Env)).resolves.toEqual({ + upserted: 201, + }) + + expect(mockModule.listMemoriesPage).toHaveBeenCalledTimes(2) + expect(mockModule.listMemoriesPage).toHaveBeenNthCalledWith(1, { + db: expect.anything(), + afterId: null, + limit: 200, + }) + expect(mockModule.listMemoriesPage).toHaveBeenNthCalledWith(2, { + db: expect.anything(), + afterId: 'memory-0199', + limit: 200, + }) + const upsertedIds = upsert.mock.calls.flatMap(([vectors]) => + vectors.map((vector) => vector.id), + ) + expect(upsertedIds).toHaveLength(201) + expect(new Set(upsertedIds).size).toBe(201) +}) + +test('memory reindex returns zero upserts for an empty table', async () => { + resetMocks() + mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert: vi.fn() }) + mockModule.isCapabilitySearchOffline.mockReturnValue(false) + mockModule.listMemoriesPage.mockResolvedValue([]) + + await expect(reindexMemoryVectors({ APP_DB: {} } as Env)).resolves.toEqual({ + upserted: 0, + }) + expect(mockModule.listMemoriesPage).toHaveBeenCalledTimes(1) +}) diff --git a/packages/worker/src/mcp/memory/memory-reindex.ts b/packages/worker/src/mcp/memory/memory-reindex.ts index 657d2217b9..218b9fa1d2 100644 --- a/packages/worker/src/mcp/memory/memory-reindex.ts +++ b/packages/worker/src/mcp/memory/memory-reindex.ts @@ -3,13 +3,18 @@ import { isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' import { + mergeVectorReindexResults, reindexVectorCandidates, type VectorReindexResult, } from '#mcp/capabilities/reindex-batches.ts' import { buildMemoryEmbedTextFromRow } from './memory-embed.ts' -import { listAllMemories } from './repo.ts' +import { listMemoriesPage } from './repo.ts' import { memoryVectorId } from './memory-vectorize.ts' +// Memories are reindexed in keyset-paged batches so memory stays bounded +// regardless of table size. +const reindexPageSize = 200 + export async function reindexMemoryVectors( env: Env, ): Promise { @@ -21,26 +26,34 @@ export async function reindexMemoryVectors( return { upserted: 0 } } - const rows = await listAllMemories({ - db: env.APP_DB, - }) - if (rows.length === 0) { - return { upserted: 0 } - } - - return reindexVectorCandidates({ - env, - index, - kind: 'memory', - candidates: rows.map((row) => ({ - id: memoryVectorId(row.id), - text: buildMemoryEmbedTextFromRow(row), - metadata: { + const pageResults: Array = [] + let afterId: string | null = null + while (true) { + const rows = await listMemoriesPage({ + db: env.APP_DB, + afterId, + limit: reindexPageSize, + }) + if (rows.length === 0) break + pageResults.push( + await reindexVectorCandidates({ + env, + index, kind: 'memory', - userId: row.user_id, - status: row.status, - ...(row.category ? { category: row.category } : {}), - }, - })), - }) + candidates: rows.map((row) => ({ + id: memoryVectorId(row.id), + text: buildMemoryEmbedTextFromRow(row), + metadata: { + kind: 'memory', + userId: row.user_id, + status: row.status, + ...(row.category ? { category: row.category } : {}), + }, + })), + }), + ) + if (rows.length < reindexPageSize) break + afterId = rows[rows.length - 1]!.id + } + return mergeVectorReindexResults('memory', pageResults) } diff --git a/packages/worker/src/package-registry/package-reindex.node.test.ts b/packages/worker/src/package-registry/package-reindex.node.test.ts index e7e50adfc1..028bd891ff 100644 --- a/packages/worker/src/package-registry/package-reindex.node.test.ts +++ b/packages/worker/src/package-registry/package-reindex.node.test.ts @@ -5,7 +5,7 @@ const mockModule = vi.hoisted(() => ({ embedTextsForVectorize: vi.fn(), getCapabilityVectorIndex: vi.fn(), isCapabilitySearchOffline: vi.fn(), - listAllSavedPackages: vi.fn(), + listSavedPackagesPage: vi.fn(), loadPackageManifestBySourceId: vi.fn(), })) @@ -24,8 +24,8 @@ vi.mock('./embed.ts', () => ({ })) vi.mock('./repo.ts', () => ({ - listAllSavedPackages: (...args: Array) => - mockModule.listAllSavedPackages(...args), + listSavedPackagesPage: (...args: Array) => + mockModule.listSavedPackagesPage(...args), savedPackageVectorId: (packageId: string) => `package_${packageId}`, })) @@ -41,10 +41,26 @@ function resetMocks() { mockModule.embedTextsForVectorize.mockReset() mockModule.getCapabilityVectorIndex.mockReset() mockModule.isCapabilitySearchOffline.mockReset() - mockModule.listAllSavedPackages.mockReset() + mockModule.listSavedPackagesPage.mockReset() mockModule.loadPackageManifestBySourceId.mockReset() } +function buildSavedPackage(id: string) { + return { + id, + userId: 'user-1', + name: `@user/${id}`, + kodyId: id, + description: `Package ${id}`, + tags: [], + searchText: null, + sourceId: `source-${id}`, + hasApp: false, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + } +} + test('saved package reindex embeds full manifests with user-scoped metadata', async () => { resetMocks() const upsert = vi.fn() @@ -63,7 +79,7 @@ test('saved package reindex embeds full manifests with user-scoped metadata', as } mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) mockModule.isCapabilitySearchOffline.mockReturnValue(false) - mockModule.listAllSavedPackages.mockResolvedValue([ + mockModule.listSavedPackagesPage.mockResolvedValue([ { id: 'pkg-1', userId: 'user-1', @@ -129,7 +145,7 @@ test('saved package reindex skips failed manifest loads and continues the batch' } mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) mockModule.isCapabilitySearchOffline.mockReturnValue(false) - mockModule.listAllSavedPackages.mockResolvedValue([ + mockModule.listSavedPackagesPage.mockResolvedValue([ { id: 'pkg-bad', userId: 'user-1', @@ -203,3 +219,50 @@ test('saved package reindex skips failed manifest loads and continues the batch' consoleError.mockRestore() } }) + +test('saved package reindex walks keyset pages and merges the page results', async () => { + resetMocks() + const upsert = vi.fn(async (_vectors: Array<{ id: string }>) => {}) + mockModule.getCapabilityVectorIndex.mockReturnValue({ upsert }) + mockModule.isCapabilitySearchOffline.mockReturnValue(false) + mockModule.loadPackageManifestBySourceId.mockResolvedValue({ + manifest: { name: '@user/pkg' }, + }) + mockModule.buildSavedPackageEmbedText.mockReturnValue('manifest embed') + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: unknown, texts: Array) => texts.map(() => [0.1]), + ) + // The first page fills the requested limit, forcing a second page fetch. + mockModule.listSavedPackagesPage.mockImplementationOnce( + async (_db: unknown, input: { afterId: string | null; limit: number }) => + Array.from({ length: input.limit }, (_, index) => + buildSavedPackage(`pkg-${String(index).padStart(4, '0')}`), + ), + ) + mockModule.listSavedPackagesPage.mockImplementationOnce(async () => [ + buildSavedPackage('pkg-last'), + ]) + + await expect( + reindexSavedPackageVectors({ APP_DB: {} } as Env, { + baseUrl: 'https://kody.example.com', + }), + ).resolves.toEqual({ upserted: 201 }) + + expect(mockModule.listSavedPackagesPage).toHaveBeenCalledTimes(2) + expect(mockModule.listSavedPackagesPage).toHaveBeenNthCalledWith( + 1, + expect.anything(), + { afterId: null, limit: 200 }, + ) + expect(mockModule.listSavedPackagesPage).toHaveBeenNthCalledWith( + 2, + expect.anything(), + { afterId: 'pkg-0199', limit: 200 }, + ) + const upsertedIds = upsert.mock.calls.flatMap(([vectors]) => + vectors.map((vector) => vector.id), + ) + expect(upsertedIds).toHaveLength(201) + expect(new Set(upsertedIds).size).toBe(201) +}) diff --git a/packages/worker/src/package-registry/package-reindex.ts b/packages/worker/src/package-registry/package-reindex.ts index 3f7dc12746..f7456482da 100644 --- a/packages/worker/src/package-registry/package-reindex.ts +++ b/packages/worker/src/package-registry/package-reindex.ts @@ -3,25 +3,25 @@ import { isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' import { + mergeVectorReindexResults, reindexVectorCandidates, type VectorReindexCandidate, type VectorReindexFailure, + type VectorReindexResult, } from '#mcp/capabilities/reindex-batches.ts' import { getErrorMessage } from '#mcp/capabilities/error-message.ts' import { buildSavedPackageEmbedText } from './embed.ts' -import { listAllSavedPackages, savedPackageVectorId } from './repo.ts' +import { listSavedPackagesPage, savedPackageVectorId } from './repo.ts' import { loadPackageManifestBySourceId } from './source.ts' +// Saved packages are reindexed in keyset-paged batches so memory stays +// bounded regardless of table size (each row also loads its manifest). +const reindexPageSize = 200 + export async function reindexSavedPackageVectors( env: Env, input: { baseUrl: string }, -): Promise<{ - upserted: number - failed?: number - warning?: string - error?: string - failures?: Array -}> { +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') @@ -30,64 +30,68 @@ export async function reindexSavedPackageVectors( return { upserted: 0 } } - const rows = await listAllSavedPackages(env.APP_DB) - if (rows.length === 0) { - return { upserted: 0 } - } - - let failed = 0 - const failures: Array = [] - const candidates: Array = [] - for (const row of rows) { - try { - const { manifest } = await loadPackageManifestBySourceId({ - env, - baseUrl: input.baseUrl, - userId: row.userId, - sourceId: row.sourceId, - }) - candidates.push({ - id: savedPackageVectorId(row.id), - text: buildSavedPackageEmbedText(manifest), - metadata: { - kind: 'package', + const pageResults: Array = [] + let afterId: string | null = null + while (true) { + const rows = await listSavedPackagesPage(env.APP_DB, { + afterId, + limit: reindexPageSize, + }) + if (rows.length === 0) break + const loadFailures: Array = [] + const candidates: Array = [] + for (const row of rows) { + try { + const { manifest } = await loadPackageManifestBySourceId({ + env, + baseUrl: input.baseUrl, userId: row.userId, - }, - }) - } catch (error) { - failed += 1 - const failure = { - id: savedPackageVectorId(row.id), - phase: 'load' as const, - error: getErrorMessage(error), - } - failures.push(failure) - console.error( - JSON.stringify({ - message: 'saved package vector reindex skipped package', - packageId: row.id, sourceId: row.sourceId, - error: failure.error, + }) + candidates.push({ + id: savedPackageVectorId(row.id), + text: buildSavedPackageEmbedText(manifest), + metadata: { + kind: 'package', + userId: row.userId, + }, + }) + } catch (error) { + const failure = { + id: savedPackageVectorId(row.id), + phase: 'load' as const, + error: getErrorMessage(error), + } + loadFailures.push(failure) + console.error( + JSON.stringify({ + message: 'saved package vector reindex skipped package', + packageId: row.id, + sourceId: row.sourceId, + error: failure.error, + }), + ) + } + } + if (loadFailures.length > 0) { + pageResults.push({ + upserted: 0, + failed: loadFailures.length, + failures: loadFailures, + }) + } + if (candidates.length > 0) { + pageResults.push( + await reindexVectorCandidates({ + env, + index, + kind: 'saved package', + candidates, }), ) } + if (rows.length < reindexPageSize) break + afterId = rows[rows.length - 1]!.id } - - const result = await reindexVectorCandidates({ - env, - index, - kind: 'saved package', - candidates, - }) - const totalFailed = failed + (result.failed ?? 0) - const allFailures = [...failures, ...(result.failures ?? [])] - if (totalFailed === 0) return { upserted: result.upserted } - - const summary = `${totalFailed} saved package vector(s) failed to reindex` - return { - upserted: result.upserted, - failed: totalFailed, - failures: allFailures, - ...(result.upserted === 0 ? { error: summary } : { warning: summary }), - } + return mergeVectorReindexResults('saved package', pageResults) } diff --git a/packages/worker/src/package-registry/repo.ts b/packages/worker/src/package-registry/repo.ts index c82104ca2c..0f6cb390ee 100644 --- a/packages/worker/src/package-registry/repo.ts +++ b/packages/worker/src/package-registry/repo.ts @@ -213,16 +213,26 @@ export async function listSavedPackagesByUserId( return (rows.results ?? []).map(mapSavedPackageRow) } -export async function listAllSavedPackages( +// Keyset-paged listing across all users, used by maintenance reindex so a +// single query never loads the whole table. Pass the last row id of the +// previous page (or null for the first page). +export async function listSavedPackagesPage( db: D1Database, + input: { + afterId: string | null + limit: number + }, ): Promise> { const rows = await db .prepare( `SELECT id, user_id, name, kody_id, description, tags_json, search_text, source_id, has_app, created_at, updated_at FROM saved_packages - ORDER BY updated_at DESC`, + WHERE id > ? + ORDER BY id + LIMIT ?`, ) + .bind(input.afterId ?? '', input.limit) .all>() return (rows.results ?? []).map(mapSavedPackageRow) } From f7ca540faaa68ed05e844a6d6229363efbb732f1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:44 +0000 Subject: [PATCH 07/14] Isolate scheduled lanes, drain reconcile to a budget, trim cold start Scheduled lanes log and report to Sentry independently instead of rethrowing, usage aggregation joins the cron, artifact reconcile loops batches under a 60s budget with an unchanged-HEAD fast path, the static capability registry builds lazily, and OG image rendering (satori + resvg wasm) loads dynamically. Co-authored-by: Kent C. Dodds --- .../src/app/handlers/community-detail.tsx | 7 +- packages/worker/src/index.ts | 80 ++++-- packages/worker/src/index.workers.test.ts | 97 +++++++ .../reconcile-artifacts-pushes.node.test.ts | 126 +++++++++ .../src/jobs/reconcile-artifacts-pushes.ts | 241 +++++++++++------- .../mcp/capabilities/registry.node.test.ts | 15 +- .../worker/src/mcp/capabilities/registry.ts | 74 +++++- 7 files changed, 504 insertions(+), 136 deletions(-) create mode 100644 packages/worker/src/index.workers.test.ts diff --git a/packages/worker/src/app/handlers/community-detail.tsx b/packages/worker/src/app/handlers/community-detail.tsx index 7a70f91916..590637846b 100644 --- a/packages/worker/src/app/handlers/community-detail.tsx +++ b/packages/worker/src/app/handlers/community-detail.tsx @@ -17,7 +17,6 @@ import { type routes } from '#app/routes.ts' import { getRequestDataCacheLookup } from '#app/request-cache.ts' import { readAuthenticatedAppUser } from '#app/authenticated-user.ts' import { CommunityActionError } from '#worker/community/errors.ts' -import { renderCommunityOgImage } from '#worker/community/og-image.ts' import { getCommunityListingWithAggregates, reportCommunityListing, @@ -121,6 +120,12 @@ export function createCommunityDetailOgImageHandler(env: Env) { } const publicListing = toPublicCommunityListing(listing) + // Lazy import (sanctioned exception to the no-inline-imports rule): + // the OG renderer pulls in satori and @resvg/resvg-wasm plus two wasm + // binaries, which would otherwise bloat isolate cold starts for a + // route that is only hit by social-media crawlers. + const { renderCommunityOgImage } = + await import('#worker/community/og-image.ts') const png = await renderCommunityOgImage({ name: publicListing.name, description: publicListing.description, diff --git a/packages/worker/src/index.ts b/packages/worker/src/index.ts index 90cceda58b..0c79f5f47e 100644 --- a/packages/worker/src/index.ts +++ b/packages/worker/src/index.ts @@ -43,6 +43,7 @@ import { handleCapabilityReindexRequest } from './capability-maintenance.ts' import { handleExecuteSmokeRequest } from './execute-maintenance.ts' import { handleJobReindexRequest } from './job-maintenance.ts' import { handleMemoryReindexRequest } from './memory-maintenance.ts' +import { handleStableUserIdBackfillRequest } from './maintenance-handler.ts' import { reconcileArtifactsPushes } from './jobs/reconcile-artifacts-pushes.ts' import { cleanupRepoSessionBranches } from './repo/repo-session-cleanup.ts' import { KodyFetchGateway } from '#mcp/fetch-gateway.ts' @@ -59,6 +60,10 @@ import { handleInboundEmail } from '#worker/email/inbound.ts' import { pruneSystemEmailRetention } from '#worker/email/system-email.ts' import { findPublicUserIdentityByUsername } from '#app/user-lookup.ts' import { pruneRetention, shouldRunRetentionCron } from '#app/retention.ts' +import { + aggregateUsageRollups, + shouldRunUsageAggregationCron, +} from '#worker/usage/aggregate-rollups.ts' export { RepoSession, @@ -232,6 +237,10 @@ const appHandler = withCors({ return handleJobReindexRequest(request, env) } + if (url.pathname === '/__maintenance/backfill-stable-user-ids') { + return handleStableUserIdBackfillRequest(request, env) + } + if (url.pathname.startsWith('/__maintenance/')) { return Response.json( { error: 'Unknown maintenance endpoint.' }, @@ -542,33 +551,50 @@ const workerHandler = { ) { const baseUrl = env.APP_BASE_URL ?? 'https://kody.local' const scheduledAt = new Date(controller.scheduledTime) - const retentionPromise = shouldRunRetentionCron(scheduledAt) - ? pruneRetention({ - env, - now: scheduledAt, - }) - : Promise.resolve(null) - const [pushesResult, cleanupResult, systemEmailResult, retentionResult] = - await Promise.allSettled([ - reconcileArtifactsPushes({ - env, - baseUrl, - now: scheduledAt, - }), - cleanupRepoSessionBranches({ - env, - now: scheduledAt, - }), - pruneSystemEmailRetention({ - db: env.APP_DB, - now: scheduledAt, - }), - retentionPromise, - ]) - if (pushesResult.status === 'rejected') throw pushesResult.reason - if (cleanupResult.status === 'rejected') throw cleanupResult.reason - if (systemEmailResult.status === 'rejected') throw systemEmailResult.reason - if (retentionResult.status === 'rejected') throw retentionResult.reason + const lanes: Array<{ name: string; run: () => Promise }> = [ + { + name: 'reconcile_artifacts_pushes', + run: () => reconcileArtifactsPushes({ env, baseUrl, now: scheduledAt }), + }, + { + name: 'repo_session_cleanup', + run: () => cleanupRepoSessionBranches({ env, now: scheduledAt }), + }, + { + name: 'system_email_retention', + run: () => + pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + now: scheduledAt, + }), + }, + ] + if (shouldRunRetentionCron(scheduledAt)) { + lanes.push({ + name: 'retention', + run: () => pruneRetention({ env, now: scheduledAt }), + }) + } + if (shouldRunUsageAggregationCron(scheduledAt)) { + lanes.push({ + name: 'usage_aggregation', + run: () => aggregateUsageRollups(env, scheduledAt), + }) + } + // Lane failures are isolated: each rejection is logged and reported to + // Sentry explicitly (the withSentry wrapper flushes captured events via + // waitUntil), but the handler never throws, so one broken lane cannot + // fail the cron invocation or hide sibling lane failures. + const results = await Promise.allSettled(lanes.map((lane) => lane.run())) + for (const [index, result] of results.entries()) { + if (result.status !== 'rejected') continue + console.error( + `scheduled_lane_failed lane=${lanes[index]?.name ?? 'unknown'}`, + result.reason, + ) + Sentry.captureException(result.reason) + } }, } satisfies ExportedHandler diff --git a/packages/worker/src/index.workers.test.ts b/packages/worker/src/index.workers.test.ts new file mode 100644 index 0000000000..7d519a3f58 --- /dev/null +++ b/packages/worker/src/index.workers.test.ts @@ -0,0 +1,97 @@ +import { expect, test, vi } from 'vitest' +import { createExecutionContext, env } from 'cloudflare:test' + +const mocks = vi.hoisted(() => ({ + reconcileArtifactsPushes: vi.fn(async () => ({})), + cleanupRepoSessionBranches: vi.fn(async () => ({})), + pruneSystemEmailRetention: vi.fn(async () => ({})), + pruneRetention: vi.fn(async () => ({})), + shouldRunRetentionCron: vi.fn(() => false), + aggregateUsageRollups: vi.fn(async () => ({ skipped: true })), + shouldRunUsageAggregationCron: vi.fn(() => false), +})) + +vi.mock('./jobs/reconcile-artifacts-pushes.ts', () => ({ + reconcileArtifactsPushes: mocks.reconcileArtifactsPushes, +})) + +vi.mock('./repo/repo-session-cleanup.ts', () => ({ + cleanupRepoSessionBranches: mocks.cleanupRepoSessionBranches, +})) + +vi.mock('#worker/email/system-email.ts', () => ({ + pruneSystemEmailRetention: mocks.pruneSystemEmailRetention, +})) + +vi.mock('#app/retention.ts', () => ({ + pruneRetention: mocks.pruneRetention, + shouldRunRetentionCron: mocks.shouldRunRetentionCron, +})) + +vi.mock('#worker/usage/aggregate-rollups.ts', () => ({ + aggregateUsageRollups: mocks.aggregateUsageRollups, + shouldRunUsageAggregationCron: mocks.shouldRunUsageAggregationCron, +})) + +const worker = (await import('./index.ts')).default + +function createController(scheduledTime: number) { + return { + scheduledTime, + cron: '*/5 * * * *', + noRetry() {}, + } satisfies ScheduledController +} + +test('scheduled runs gated lanes and passes EMAIL_BLOBS to system-email retention', async () => { + mocks.shouldRunUsageAggregationCron.mockReturnValueOnce(true) + const scheduledTime = Date.parse('2026-07-05T10:00:30.000Z') + + await worker.scheduled?.( + createController(scheduledTime), + env, + createExecutionContext(), + ) + + expect(mocks.reconcileArtifactsPushes).toHaveBeenCalledWith( + expect.objectContaining({ now: new Date(scheduledTime) }), + ) + expect(mocks.cleanupRepoSessionBranches).toHaveBeenCalledTimes(1) + expect(mocks.pruneSystemEmailRetention).toHaveBeenCalledWith( + expect.objectContaining({ blobs: env.EMAIL_BLOBS }), + ) + expect(mocks.aggregateUsageRollups).toHaveBeenCalledWith( + env, + new Date(scheduledTime), + ) + expect(mocks.pruneRetention).not.toHaveBeenCalled() +}) + +test('scheduled isolates a failing lane: logs it and keeps siblings and the invocation alive', async () => { + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}) + mocks.reconcileArtifactsPushes.mockRejectedValueOnce( + new Error('reconcile exploded'), + ) + const scheduledTime = Date.parse('2026-07-05T10:10:00.000Z') + + try { + await expect( + worker.scheduled?.( + createController(scheduledTime), + env, + createExecutionContext(), + ), + ).resolves.toBeUndefined() + + expect(mocks.cleanupRepoSessionBranches).toHaveBeenCalled() + expect(mocks.pruneSystemEmailRetention).toHaveBeenCalled() + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'scheduled_lane_failed lane=reconcile_artifacts_pushes', + expect.objectContaining({ message: 'reconcile exploded' }), + ) + } finally { + consoleErrorSpy.mockRestore() + } +}) diff --git a/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts b/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts index 50a085999f..c9a319341d 100644 --- a/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts +++ b/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts @@ -98,6 +98,8 @@ test('publishes changed Artifacts HEADs and records reconcile checks', async () errors: 0, tokenCleanupErrors: 0, tokensRevoked: 0, + batches: 1, + budgetExhausted: false, }) expect(mockModule.publishFromExternalRef).toHaveBeenCalledWith( expect.objectContaining({ @@ -182,6 +184,8 @@ test('reconcile records source checks and continues the batch when a source or i errors: 1, tokenCleanupErrors: 0, tokensRevoked: 0, + batches: 1, + budgetExhausted: false, }) expect(mockModule.resolveArtifactSourceHead).toHaveBeenCalledTimes(2) }) @@ -241,3 +245,125 @@ test('reconcile runs token cleanup in the 03:00 UTC window without blocking publ ) expect(mockModule.publishFromExternalRef).toHaveBeenCalledTimes(1) }) + +test('reconcile drains multiple batches in one tick without DO work for unchanged heads', async () => { + mockModule.listEntitySourcesForExternalReconcile.mockReset() + mockModule.listEntitySourcesForExternalReconcile + .mockResolvedValueOnce([ + source({ id: 'source-1', published_commit: 'commit-current' }), + source({ + id: 'source-2', + repo_id: 'repo-2', + published_commit: 'commit-current', + }), + ]) + .mockResolvedValueOnce([ + source({ + id: 'source-3', + repo_id: 'repo-3', + published_commit: 'commit-current', + }), + ]) + mockModule.resolveArtifactSourceHead.mockReset() + mockModule.resolveArtifactSourceHead.mockResolvedValue({ + branch: 'main', + commit: 'commit-current', + }) + mockModule.updateEntitySource.mockClear() + mockModule.repoSessionRpc.mockClear() + + const result = await reconcileArtifactsPushes({ + env: { APP_DB: {} } as Env, + baseUrl: 'https://kody.test', + now: new Date('2026-05-04T02:00:00.000Z'), + batchSize: 2, + }) + + expect(result).toEqual( + expect.objectContaining({ + checked: 3, + alreadyPublished: 3, + batches: 2, + budgetExhausted: false, + }), + ) + expect( + mockModule.listEntitySourcesForExternalReconcile, + ).toHaveBeenCalledTimes(2) + // Unchanged HEADs must only advance the cursor, never touch the DO. + expect(mockModule.repoSessionRpc).not.toHaveBeenCalled() + expect(mockModule.updateEntitySource).toHaveBeenCalledTimes(3) +}) + +test('reconcile stops between items once the time budget is exhausted', async () => { + mockModule.listEntitySourcesForExternalReconcile.mockReset() + mockModule.listEntitySourcesForExternalReconcile.mockResolvedValue([ + source({ id: 'source-1', published_commit: 'commit-current' }), + source({ + id: 'source-2', + repo_id: 'repo-2', + published_commit: 'commit-current', + }), + ]) + mockModule.resolveArtifactSourceHead.mockReset() + mockModule.resolveArtifactSourceHead.mockResolvedValue({ + branch: 'main', + commit: 'commit-current', + }) + + const result = await reconcileArtifactsPushes({ + env: { APP_DB: {} } as Env, + baseUrl: 'https://kody.test', + now: new Date('2026-05-04T02:00:00.000Z'), + timeBudgetMs: 0, + }) + + expect(result).toEqual( + expect.objectContaining({ + checked: 1, + batches: 1, + budgetExhausted: true, + }), + ) + expect( + mockModule.listEntitySourcesForExternalReconcile, + ).toHaveBeenCalledTimes(1) +}) + +test('reconcile does not reprocess a source whose cursor write keeps failing', async () => { + mockModule.listEntitySourcesForExternalReconcile.mockReset() + mockModule.listEntitySourcesForExternalReconcile.mockResolvedValue([ + source({ id: 'source-1', published_commit: 'commit-current' }), + ]) + mockModule.resolveArtifactSourceHead.mockReset() + mockModule.resolveArtifactSourceHead.mockResolvedValue({ + branch: 'main', + commit: 'commit-current', + }) + mockModule.updateEntitySource.mockReset() + mockModule.updateEntitySource.mockRejectedValue(new Error('d1 unavailable')) + + const result = await reconcileArtifactsPushes({ + env: { APP_DB: {} } as Env, + baseUrl: 'https://kody.test', + now: new Date('2026-05-04T02:00:00.000Z'), + batchSize: 1, + }) + + // The stale source stays selectable (its cursor write failed) but must be + // processed at most once per tick instead of spinning until the budget. + expect(result).toEqual( + expect.objectContaining({ + checked: 1, + errors: 1, + batches: 1, + budgetExhausted: false, + }), + ) + expect( + mockModule.listEntitySourcesForExternalReconcile, + ).toHaveBeenCalledTimes(2) + + mockModule.updateEntitySource.mockReset() + mockModule.updateEntitySource.mockResolvedValue(true) +}) diff --git a/packages/worker/src/jobs/reconcile-artifacts-pushes.ts b/packages/worker/src/jobs/reconcile-artifacts-pushes.ts index 464ce96317..82e5713ab1 100644 --- a/packages/worker/src/jobs/reconcile-artifacts-pushes.ts +++ b/packages/worker/src/jobs/reconcile-artifacts-pushes.ts @@ -2,6 +2,7 @@ import { listEntitySourcesForExternalReconcile, updateEntitySource, } from '#worker/repo/entity-sources.ts' +import { type EntitySourceRow } from '#worker/repo/types.ts' import { resolveArtifactSourceHead } from '#worker/repo/artifacts.ts' import { revokeStaleArtifactsTokens } from '#worker/repo/artifacts-tokens.ts' import { repoSessionRpc } from '#worker/repo/repo-session-do.ts' @@ -16,10 +17,19 @@ export type ReconcileArtifactsPushesResult = { errors: number tokenCleanupErrors: number tokensRevoked: number + batches: number + budgetExhausted: boolean } const defaultBatchSize = 50 const defaultStaleAfterMinutes = 5 +/** + * Wall-clock budget for one cron tick. The reconcile loop keeps pulling + * batches of stale sources until it either drains the backlog or exceeds + * this budget, so throughput scales with available time instead of being + * capped at one batch per five-minute tick. + */ +export const reconcileTimeBudgetMs = 60_000 function minutesAgoIso(now: Date, minutes: number) { return new Date(now.getTime() - minutes * 60_000).toISOString() @@ -34,19 +44,16 @@ export async function reconcileArtifactsPushes(input: { baseUrl: string batchSize?: number staleAfterMinutes?: number + timeBudgetMs?: number now?: Date }): Promise { const now = input.now ?? new Date() - const sources = await listEntitySourcesForExternalReconcile( - input.env.APP_DB, - { - before: minutesAgoIso( - now, - input.staleAfterMinutes ?? defaultStaleAfterMinutes, - ), - limit: input.batchSize ?? defaultBatchSize, - }, + const batchSize = input.batchSize ?? defaultBatchSize + const before = minutesAgoIso( + now, + input.staleAfterMinutes ?? defaultStaleAfterMinutes, ) + const deadlineMs = Date.now() + (input.timeBudgetMs ?? reconcileTimeBudgetMs) const result: ReconcileArtifactsPushesResult = { checked: 0, published: 0, @@ -57,107 +64,153 @@ export async function reconcileArtifactsPushes(input: { errors: 0, tokenCleanupErrors: 0, tokensRevoked: 0, + batches: 0, + budgetExhausted: false, } const shouldCleanupTokens = shouldRunDailyTokenCleanup(now) - for (const source of sources) { - result.checked += 1 - try { - const head = await resolveArtifactSourceHead(input.env, source.repo_id) - if (shouldCleanupTokens) { - try { - const cleanup = await revokeStaleArtifactsTokens( - input.env, - source.repo_id, - { keepAfter: now }, - ) - result.tokensRevoked += cleanup.revoked - } catch (cleanupError) { - result.tokenCleanupErrors += 1 - console.warn('reconcile_artifacts_pushes token cleanup failed', { - sourceId: source.id, - repoId: source.repo_id, - error: - cleanupError instanceof Error - ? cleanupError.message - : String(cleanupError), - }) - } + // Sources whose cursor write failed stay stale and would be re-selected by + // the next batch query; tracking seen ids prevents re-processing them in a + // tight loop within a single tick. + const seenSourceIds = new Set() + while (!result.budgetExhausted) { + const batch = await listEntitySourcesForExternalReconcile( + input.env.APP_DB, + { before, limit: batchSize }, + ) + const sources = batch.filter((source) => !seenSourceIds.has(source.id)) + if (sources.length === 0) break + result.batches += 1 + for (const source of sources) { + seenSourceIds.add(source.id) + await reconcileSource({ + env: input.env, + baseUrl: input.baseUrl, + source, + now, + shouldCleanupTokens, + result, + }) + if (Date.now() >= deadlineMs) { + result.budgetExhausted = true + break } - if (!head.commit) { - result.missingHead += 1 - await updateEntitySource(input.env.APP_DB, { - id: source.id, - userId: source.user_id, - lastExternalCheckAt: now.toISOString(), + } + if (batch.length < batchSize) break + } + console.info('reconcile_artifacts_pushes', result) + return result +} + +async function reconcileSource(input: { + env: Env + baseUrl: string + source: EntitySourceRow + now: Date + shouldCleanupTokens: boolean + result: ReconcileArtifactsPushesResult +}) { + const { env, source, now, result } = input + result.checked += 1 + try { + const head = await resolveArtifactSourceHead(env, source.repo_id) + if (input.shouldCleanupTokens) { + try { + const cleanup = await revokeStaleArtifactsTokens(env, source.repo_id, { + keepAfter: now, }) - continue - } - if (head.commit === source.published_commit) { - result.alreadyPublished += 1 - await updateEntitySource(input.env.APP_DB, { - id: source.id, - userId: source.user_id, - lastExternalCheckAt: now.toISOString(), + result.tokensRevoked += cleanup.revoked + } catch (cleanupError) { + result.tokenCleanupErrors += 1 + console.warn('reconcile_artifacts_pushes token cleanup failed', { + sourceId: source.id, + repoId: source.repo_id, + error: + cleanupError instanceof Error + ? cleanupError.message + : String(cleanupError), }) - continue } - const sessionId = `external-reconcile-${source.id}` - const publishResult = await repoSessionRpc( - input.env, - sessionId, - ).publishFromExternalRef({ - sessionId, - sourceId: source.id, + } + // Unchanged (or unresolvable) HEADs only advance the reconcile cursor; + // the RepoSession Durable Object is spun up solely for changed HEADs. + if (!head.commit) { + result.missingHead += 1 + await updateEntitySource(env.APP_DB, { + id: source.id, userId: source.user_id, - newCommit: head.commit, - expectedHead: head.commit, - allowForce: false, - baseUrl: input.baseUrl, + lastExternalCheckAt: now.toISOString(), }) - switch (publishResult.status) { - case 'already_published': - result.alreadyPublished += 1 - break - case 'published': - result.published += 1 - break - case 'checks_failed': - result.checksFailed += 1 - break - case 'not_fast_forward': - result.notFastForward += 1 - break + return + } + if (head.commit === source.published_commit) { + result.alreadyPublished += 1 + await updateEntitySource(env.APP_DB, { + id: source.id, + userId: source.user_id, + lastExternalCheckAt: now.toISOString(), + }) + return + } + const sessionId = `external-reconcile-${source.id}` + const publishResult = await repoSessionRpc( + env, + sessionId, + ).publishFromExternalRef({ + sessionId, + sourceId: source.id, + userId: source.user_id, + newCommit: head.commit, + expectedHead: head.commit, + allowForce: false, + baseUrl: input.baseUrl, + }) + switch (publishResult.status) { + case 'already_published': + result.alreadyPublished += 1 + break + case 'published': + result.published += 1 + break + case 'checks_failed': + result.checksFailed += 1 + break + case 'not_fast_forward': + result.notFastForward += 1 + break + default: { + const exhaustive: never = publishResult + throw new Error( + `Unknown external publish status: ${String(exhaustive)}`, + ) } - await updateEntitySource(input.env.APP_DB, { + } + await updateEntitySource(env.APP_DB, { + id: source.id, + userId: source.user_id, + lastExternalCheckAt: now.toISOString(), + }) + } catch (error) { + result.errors += 1 + console.warn('reconcile_artifacts_pushes source failed', { + sourceId: source.id, + repoId: source.repo_id, + error: error instanceof Error ? error.message : String(error), + }) + try { + await updateEntitySource(env.APP_DB, { id: source.id, userId: source.user_id, lastExternalCheckAt: now.toISOString(), }) - } catch (error) { - result.errors += 1 - console.warn('reconcile_artifacts_pushes source failed', { + } catch (updateError) { + console.warn('reconcile_artifacts_pushes cursor update failed', { sourceId: source.id, repoId: source.repo_id, - error: error instanceof Error ? error.message : String(error), + error: + updateError instanceof Error + ? updateError.message + : String(updateError), }) - try { - await updateEntitySource(input.env.APP_DB, { - id: source.id, - userId: source.user_id, - lastExternalCheckAt: now.toISOString(), - }) - } catch (updateError) { - console.warn('reconcile_artifacts_pushes cursor update failed', { - sourceId: source.id, - repoId: source.repo_id, - error: - updateError instanceof Error - ? updateError.message - : String(updateError), - }) - } } } - console.info('reconcile_artifacts_pushes', result) - return result } diff --git a/packages/worker/src/mcp/capabilities/registry.node.test.ts b/packages/worker/src/mcp/capabilities/registry.node.test.ts index af11d631d2..1994d23c55 100644 --- a/packages/worker/src/mcp/capabilities/registry.node.test.ts +++ b/packages/worker/src/mcp/capabilities/registry.node.test.ts @@ -3,7 +3,9 @@ import { createMcpCallerContext } from '#mcp/context.ts' import { clearCapabilityRegistryCacheForTests, capabilityMap, + capabilitySpecs, getCapabilityRegistryForContext, + getStaticRegistry, } from '#mcp/capabilities/registry.ts' test('getCapabilityRegistryForContext bypasses connector caches when no connectors are attached', async () => { @@ -33,7 +35,18 @@ test('getCapabilityRegistryForContext bypasses connector caches when no connecto }) expect(get).not.toHaveBeenCalled() - expect(registry.capabilityMap).toBe(capabilityMap) + expect(registry.capabilityMap).toBe(getStaticRegistry().capabilityMap) +}) + +test('getStaticRegistry memoizes and the lazy constant exports mirror it', () => { + expect(getStaticRegistry()).toBe(getStaticRegistry()) + expect(capabilityMap.email_send).toBe( + getStaticRegistry().capabilityMap.email_send, + ) + expect(Object.keys(capabilitySpecs)).toEqual( + Object.keys(getStaticRegistry().capabilitySpecs), + ) + expect('email_send' in capabilityMap).toBe(true) }) test('getCapabilityRegistryForContext filters admin capabilities by current caller roles', async () => { diff --git a/packages/worker/src/mcp/capabilities/registry.ts b/packages/worker/src/mcp/capabilities/registry.ts index 7317f01bc0..2f0e3c7d39 100644 --- a/packages/worker/src/mcp/capabilities/registry.ts +++ b/packages/worker/src/mcp/capabilities/registry.ts @@ -19,23 +19,71 @@ import { type McpServerSnapshot } from '#worker/mcp-client/types.ts' import { getCachedRemoteConnectorSnapshot } from '#worker/remote-connector/snapshot-cache.ts' import { type RemoteConnectorSnapshot } from '#worker/remote-connector/types.ts' -const staticRegistry = buildCapabilityRegistry(builtinDomains) +let staticRegistryMemo: BuiltCapabilityRegistry | null = null -export const capabilityList = staticRegistry.capabilityList +/** + * Building the builtin registry converts ~99 capability Zod schemas to JSON + * Schema, which is far too expensive to run at module scope on every isolate + * cold start. Memoized so the work happens at most once per isolate, on first + * use. + */ +export function getStaticRegistry(): BuiltCapabilityRegistry { + staticRegistryMemo ??= buildCapabilityRegistry(builtinDomains) + return staticRegistryMemo +} + +/** + * Lazy view over one field of the static registry, so the legacy constant + * exports below keep working without forcing the registry build at module + * load. The proxy target only supplies the right base shape (array vs plain + * object); every operation is forwarded to the real, memoized registry. + */ +function lazyStaticRegistryField( + key: Key, + target: object, +): BuiltCapabilityRegistry[Key] { + const resolve = () => getStaticRegistry()[key] as object + return new Proxy(target, { + get(_target, property) { + return Reflect.get(resolve(), property) + }, + has(_target, property) { + return Reflect.has(resolve(), property) + }, + ownKeys() { + return Reflect.ownKeys(resolve()) + }, + getOwnPropertyDescriptor(_target, property) { + return Object.getOwnPropertyDescriptor(resolve(), property) + }, + }) as BuiltCapabilityRegistry[Key] +} + +export const capabilityList = lazyStaticRegistryField('capabilityList', []) -export const capabilityDomains = staticRegistry.capabilityDomains +export const capabilityDomains = lazyStaticRegistryField( + 'capabilityDomains', + [], +) -export const capabilityDomainDescriptionsByName = - staticRegistry.capabilityDomainDescriptionsByName +export const capabilityDomainDescriptionsByName = lazyStaticRegistryField( + 'capabilityDomainDescriptionsByName', + {}, +) -export const capabilityMap = staticRegistry.capabilityMap +export const capabilityMap = lazyStaticRegistryField('capabilityMap', {}) -export const capabilitySpecs = staticRegistry.capabilitySpecs +export const capabilitySpecs = lazyStaticRegistryField('capabilitySpecs', {}) -export const capabilityToolDescriptors = - staticRegistry.capabilityToolDescriptors +export const capabilityToolDescriptors = lazyStaticRegistryField( + 'capabilityToolDescriptors', + {}, +) -export const capabilityHandlers = staticRegistry.capabilityHandlers +export const capabilityHandlers = lazyStaticRegistryField( + 'capabilityHandlers', + {}, +) export const capabilityRegistryCacheTtlMs = 30_000 export const capabilityRegistryCacheLimit = 50 @@ -115,7 +163,7 @@ async function buildCapabilityRegistryForDynamicSources(input: { (domain) => (domain ? [domain.domain] : []), ) if (synthesizedDomains.length === 0) { - return staticRegistry + return getStaticRegistry() } return buildCapabilityRegistry([...builtinDomains, ...synthesizedDomains]) } @@ -174,7 +222,7 @@ export async function getCapabilityRegistryForContext(input: { const userId = input.callerContext.user?.userId ?? null if (!userId) { return filterRegistryForContext({ - registry: staticRegistry, + registry: getStaticRegistry(), callerContext: input.callerContext, }) } @@ -184,7 +232,7 @@ export async function getCapabilityRegistryForContext(input: { }) if (refs.length === 0 && mcpServerRefs.length === 0) { return filterRegistryForContext({ - registry: staticRegistry, + registry: getStaticRegistry(), callerContext: input.callerContext, }) } From 90835fb9f216bfe8c6895fc55f6741034015a49f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:12:44 +0000 Subject: [PATCH 08/14] Document R2 email blobs, retention, reconcile, and export changes Co-authored-by: Kent C. Dodds --- .../contributing/architecture/data-storage.md | 158 ++++++++++++++---- 1 file changed, 127 insertions(+), 31 deletions(-) diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index ccc42d7389..93ab3328dc 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -59,6 +59,12 @@ Deletion must cover these user-owned surfaces: authorized user requests. - **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 + `email_messages.raw_mime_key` while the rows still exist and deleted + best-effort with warnings, mirroring the KV/Vectorize reporting. Rows owned by + `system:email` keep their blobs here (they are not user data); those blobs are + removed when the system-email retention prune deletes the messages through the + shared delete-message helper. - **KV:** published bundle artifact keys, source/manifest snapshot keys, community listing snapshots, and per-user package retriever cache/index keys in `BUNDLE_ARTIFACTS_KV` are deleted before D1 projection rows are removed. @@ -119,6 +125,12 @@ migration-safe chunked interface: StorageRunner `exportStorage({ pageSize, startAfter })` RPC as the dedicated storage export capability. +D1 rows are always read with SQL-level keyset pagination: every query orders by +the table's `rowid`, resumes strictly after an opaque cursor, and applies a SQL +`LIMIT`, so a single query never loads a whole table. `account_export_section` +fetches only the requested page, and the full browser export streams each table +page by page (500 rows per query) before assembling the final JSON document. + Durable Object export behavior: - `StorageRunner` bucket contents are exported with paged entries. These buckets @@ -154,12 +166,16 @@ Relational app data lives in D1. The schema is defined by migrations in `packages/worker/migrations/`: -- `users`: login identity and password hash. There is no persisted mapping from - the stable MCP `userId` (SHA-256 of the normalized email) back to the row; - contextless paths (inbound email) reverse-resolve it by scanning and hashing - (`findUserAccountByStableUserId`). A persisted, indexed `stable_user_id` - column with an app-level backfill is required before onboarding external users - / design partners. +- `users`: login identity and password hash, plus the persisted stable MCP + `userId` (`stable_user_id`, SHA-256 of the normalized email, unique partial + index from migration 0052; written at signup). Inbound email routing does not + reverse-resolve stable ids at all — it uses the indexed username lookup + (`findPublicUserIdentityByUsername`). The remaining contextless paths resolve + stable ids with one indexed point read (`findUserRowByStableUserId` / + `findUserAccountByStableUserId`); legacy rows with a NULL `stable_user_id` + fall back to a scan-and-hash that self-heals by writing the computed id back, + and the `POST /__maintenance/backfill-stable-user-ids` endpoint backfills all + remaining legacy rows in one pass. - `password_resets`: hashed reset tokens with expiry and foreign key to users - `jobs`: persisted job metadata, caller context, schedule state, repo source pointers, and run observability counters/history @@ -195,6 +211,28 @@ snapshots are stored in `BUNDLE_ARTIFACTS_KV`. source ids, published commits, bundle artifact rows, community listing ids, and package ids. +## R2 (`EMAIL_BLOBS`) + +Raw email MIME payloads live in the `EMAIL_BLOBS` R2 bucket instead of D1. +`email_messages` stores an object key in `raw_mime_key` +(`email-raw:v1:{userId}/{messageId}`), and the legacy inline `raw_mime` column +is kept only for rows written before the offload (or when the binding is +unavailable at write time — the write path falls back to inline storage rather +than losing mail, and never throws from that decision). + +- All reads go through `loadRawMime` in `packages/worker/src/email/repo.ts`, + which prefers the inline payload and otherwise fetches the blob by key. + Attachment content extraction re-parses the resolved MIME the same way as + before. +- `deleteEmailMessageById` deletes the message's blob best-effort after the D1 + row delete; blob cleanup never blocks the delete. +- Account deletion enumerates the user's `raw_mime_key` values before D1 rows + are removed and deletes the blobs with warnings on failure. +- Bucket names: `kody-email-blobs` (production), per-preview + `{worker}-email-blobs` buckets created and cleaned up by + `tools/ci/preview-resources.ts`, and the test env reuses the preview-style + name locally (Wrangler/vitest-pool-workers simulate the bucket). + ## Durable Objects (`MCP_OBJECT`) MCP server runtime state is hosted via a Durable Object class (`MCP`) in @@ -216,6 +254,12 @@ Jobs use two Durable Object roles: SQLite state that can be bound to execute calls, jobs, and dedicated storage inspection capabilities +Each `JobManager` alarm processes at most `maxDueJobsPerAlarm` due jobs +(`packages/worker/src/jobs/repo.ts`, oldest `next_run_at` first). When more due +jobs remain after a run, the post-run alarm resync arms a near-immediate +follow-up alarm so large backlogs drain across multiple short invocations +instead of one Durable Object wake. + Storage split: - D1 `jobs` table: job metadata, persisted caller context, schedule, run @@ -286,6 +330,7 @@ Bindings are configured per environment in `packages/worker/wrangler.jsonc` - `APP_DB` (D1) - `OAUTH_KV` (KV) - `BUNDLE_ARTIFACTS_KV` (KV) +- `EMAIL_BLOBS` (R2, raw email MIME blobs) - `MCP_OBJECT` (Durable Objects) - `REMOTE_CONNECTOR_SESSION` (Durable Objects) - `JOB_MANAGER` (Durable Objects) @@ -377,18 +422,30 @@ pushed the session commit to the source Artifacts repo. `packages/worker/src/jobs/reconcile-artifacts-pushes.ts` is a safety net for external pushes that were not followed by an explicit `package_publish_external_push` call. The Worker scheduled handler runs every -five minutes (`wrangler.jsonc` `*/5 * * * *`), selects a small batch of stale -`entity_sources` rows by `last_external_check_at`, resolves the Artifacts -default-branch HEAD, and calls the same external publish path when HEAD differs -from `published_commit`. +five minutes (`wrangler.jsonc` `*/5 * * * *`) and loops over batches of stale +`entity_sources` rows (selected by `last_external_check_at`) until the backlog +is drained or a wall-clock time budget (`reconcileTimeBudgetMs`, ~60 seconds) is +exhausted, so throughput scales with backlog size instead of being capped at one +batch per tick. For each source it resolves the Artifacts default-branch HEAD +cheaply; when HEAD matches `published_commit` (or is unresolvable) it only +advances `last_external_check_at` without any Durable Object work, and it spins +up the RepoSession publish path only when HEAD differs. The reconcile loop is idempotent: if another caller publishes the same commit first, the publish path returns `already_published`. Check failures and non-fast-forward results leave D1/KV untouched and are counted in the one-line -metrics log. Once per day during the 03:00 UTC cron window, reconcile also calls -`revokeStaleArtifactsTokens(...)` for checked repos to clean up expired +metrics log, which also records batches processed and whether the time budget +was exhausted. Once per day during the 03:00 UTC cron window, reconcile also +calls `revokeStaleArtifactsTokens(...)` for checked repos to clean up expired Artifacts tokens. +Reconcile runs as one lane of the scheduled handler in +`packages/worker/src/index.ts`, alongside repo-session cleanup, system-email +retention, general retention, and hourly usage-rollup aggregation. Lane failures +are isolated: each rejected lane is logged with a `scheduled_lane_failed` tag +and reported to Sentry, and the handler never throws, so one broken lane cannot +abort or mask the others. + Production note: - Production deploys warn that the documented Artifacts Worker binding config is @@ -493,7 +550,9 @@ app-owned keys in it. App-owned `BUNDLE_ARTIFACTS_KV` keys are: - `community-snapshot:v1:{listingId}`. - `package-retriever-manifest:v1:{userId}:{packageId}:{revision}`. - `package-retriever-index:v1:{userId}:{scope}` for the legacy combined - retriever index blob. + retriever index blob. No longer written (the combined blob risked the 25MB KV + value limit); refresh and removal delete it best-effort, and it is read only + as a fallback when the KV binding does not support `list()`. - `package-retriever-index-entry:v1:{userId}:{scope}:{packageId}:{retrieverKey}` for per-entry retriever index rows. @@ -501,6 +560,18 @@ Account deletion derives these keys from D1 rows and package ids before deleting D1 projections. New KV prefixes must add corresponding account-deletion coverage or a deliberate retention note. +### R2 key contracts + +App-owned `EMAIL_BLOBS` keys are: + +- `email-raw:v1:{userId}/{messageId}` — raw email MIME for the message row that + stores this key in `email_messages.raw_mime_key`. The `userId` prefix is part + of the per-user isolation contract; account deletion deletes a user's blobs by + the keys stored on their rows. + +New R2 key prefixes must add corresponding account-deletion coverage or a +deliberate retention note, same as KV. + ### Vectorize metadata contracts Vector ids and metadata are conventional and require reindexing when changed. @@ -553,13 +624,17 @@ migration plan. The Worker scheduled handler runs every five minutes, but `packages/worker/src/app/retention.ts` gates the general retention job to the -top of the hour and deletes at most one configured batch per table each time. -This keeps D1 single-writer pressure bounded when a backlog exists; progress is -reported with a one-line `retention-prune` log. The retention module owns the -named constants and manifest, and -`packages/worker/src/app/retention.node.test.ts` fails if a future -growth-pattern D1 table is added without either a policy or a documented -exemption. +top of the hour. Each hourly run loops in round-robin passes over the policy +tables — every pending table gets one configured batch before any table gets a +second one — until every table is drained or the run's time budget +(`retentionRunTimeBudgetMs`, ~20 seconds measured with `Date.now`) is exhausted. +The first pass always completes so a hot table cannot starve the others, and +per-batch sizes stay small to bound D1 single-writer pressure. Progress is +reported with a one-line `retention-prune` log that includes batches-per-table +counts and whether the budget ran out. The retention module owns the named +constants and manifest, and `packages/worker/src/app/retention.node.test.ts` +fails if a future growth-pattern D1 table is added without either a policy or a +documented exemption. Current retention policies: @@ -567,7 +642,11 @@ Current retention policies: and at most 500 runs per `(user_id, package_id)`. Logs are deleted before their runs, and orphan logs are pruned separately. Rows with `status = 'running'`, active invocation/workflow/session references, or a - running child debug run are kept. + running child debug run are kept. The age prune is index-driven; the + per-package cap prune ranks only a bounded set of the largest + `(user_id, package_id)` pairs per batch (`packageRuntimeCapPairsPerBatch`) + instead of window-ranking the whole table, so it stays cheap as the table + grows. - `package_invocations`: keep terminal idempotency rows for 90 days. Rows with `status = 'in_progress'` are never pruned so duplicate requests cannot bypass the in-flight guard. @@ -580,18 +659,35 @@ Current retention policies: - `published_bundle_artifacts`: delete D1 rows and their `BUNDLE_ARTIFACTS_KV` blobs only when the row is older than 30 days, its `published_commit` is no longer current for any matching `entity_sources` row, and there is no active - repo session for the source. Ambiguous publish/edit cases are intentionally - kept. + repo session for the source. When a row is pruned, the matching + `source-snapshot:v1:{sourceId}:{commit}` and + `source-manifest-snapshot:v1:{sourceId}:{commit}` KV keys are deleted under + the same safety conditions, so per-commit snapshots no longer accumulate + forever. Ambiguous publish/edit cases are intentionally kept. - `email_delivery_events`: user-owned delivery events keep 90 days. System email rows under `system:email` remain governed by the system-email retention job, - which prunes messages and delivery events older than 90 days, deletes stale - `system_email_daily_counters`, and caps stored system messages at 5000. + which prunes messages (and their `EMAIL_BLOBS` raw-MIME blobs) and delivery + events older than 90 days in bounded batches within its own time budget, + deletes stale `system_email_daily_counters`, and caps stored system messages + at 5000. +- `email_messages` / `email_attachments` / `email_threads`: user-owned messages + (excluding the `system:email` owner) keep 365 days, deleted oldest first in + batches. For rows with `raw_mime_key`, the `EMAIL_BLOBS` R2 blob is deleted + before the row; if the binding is missing or the blob delete fails, those rows + are skipped and counted as warnings so blobs are never orphaned. 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. +- `usage_rollups`: per user/metric/month rollups keep 24 months by `month` key; + raw Analytics Engine usage events follow platform retention. - `audit_events`: global hashed auth/security audit events keep 180 days. They are not user-owned D1 rows and remain independent of account deletion/export. -`email_messages` and related user inbox tables remain account-deletion scoped; -the scheduled policy currently bounds delivery/audit metadata, not user message -history. `usage_rollups` is bounded by `(user_id, metric, month)`, while raw -Analytics Engine usage events follow platform retention. -`archived_job_artifacts` is explicitly exempt from the retention manifest -because job artifact cleanup is driven by each row's `retain_until` value. +Migration `0055-retention-indexes.sql` adds the global time-column indexes these +prunes order by (`created_at` / `day` / `month` / `started_at` across users) +since the older per-user composite indexes cannot serve them. + +Documented exemptions: `archived_job_artifacts` is exempt because job artifact +cleanup is driven by each row's `retain_until` value, and `mcp_memories` is +exempt because memories are durable user-curated content removed by explicit +user action or account deletion rather than by time-based retention. From 0a0a859cffa9d7dc202ec75301278773623121d2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:13:38 +0000 Subject: [PATCH 09/14] Update primitives map for R2 email blobs and derived usage rollups Co-authored-by: Kent C. Dodds --- .../contributing/architecture/primitives.yaml | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index ad9fba292c..ee4dcb90e9 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -68,13 +68,17 @@ primitives: group: surfaces name: Scheduled handler summary: - Worker cron entry point (Artifacts push reconcile every five minutes, - daily stale-token cleanup around 03:00 UTC, and abandoned repo-session - branch cleanup). + Worker cron entry point with independently isolated lanes (time-budgeted + Artifacts push reconcile, hourly retention pruning, hourly usage rollup + aggregation, system-email retention, daily stale-token cleanup around + 03:00 UTC, and abandoned repo-session branch cleanup); one failing lane + never aborts the others. code: - packages/worker/src/index.ts - packages/worker/src/jobs/reconcile-artifacts-pushes.ts - packages/worker/src/repo/repo-session-cleanup.ts + - packages/worker/src/app/retention.ts + - packages/worker/src/usage/aggregate-rollups.ts docs: - docs/contributing/architecture/data-storage.md @@ -449,16 +453,31 @@ primitives: name: Usage metering summary: Per-user usage events recorded at runtime chokepoints via recordUsage(); - raw events go to the USAGE_EVENTS Analytics Engine dataset and per-user, - per-metric, per-month counters accumulate in D1 usage_rollups. + raw events go to the USAGE_EVENTS Analytics Engine dataset and D1 + usage_rollups is a derived aggregate recomputed hourly from the AE SQL + API (direct D1 upserts only in local dev without the AE binding). Observability only; entitlements use separate source-table counts and daily counters. code: - packages/worker/src/usage/record-usage.ts + - packages/worker/src/usage/aggregate-rollups.ts - packages/worker/migrations/0047-usage-rollups.sql docs: - docs/contributing/architecture/usage-metering.md + - id: email-blobs-r2 + group: storage + name: Email blob storage (R2) + summary: + EMAIL_BLOBS R2 bucket holding raw email MIME under the frozen key + contract email-raw:v1:{userId}/{messageId}; D1 email_messages keeps + metadata plus raw_mime_key, with inline raw_mime as the fallback when + the binding is unavailable. Covered by account deletion and retention. + code: + - packages/worker/src/email/repo.ts + docs: + - docs/contributing/architecture/data-storage.md + invariants: - id: per-user-isolation summary: From a72d101e4f05684001b2947f7e6e39c7aa057089 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:22:25 +0000 Subject: [PATCH 10/14] Format primitives.yaml Co-authored-by: Kent C. Dodds --- docs/contributing/architecture/primitives.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/contributing/architecture/primitives.yaml b/docs/contributing/architecture/primitives.yaml index ee4dcb90e9..d6bc7e78df 100644 --- a/docs/contributing/architecture/primitives.yaml +++ b/docs/contributing/architecture/primitives.yaml @@ -454,8 +454,8 @@ primitives: summary: Per-user usage events recorded at runtime chokepoints via recordUsage(); raw events go to the USAGE_EVENTS Analytics Engine dataset and D1 - usage_rollups is a derived aggregate recomputed hourly from the AE SQL - API (direct D1 upserts only in local dev without the AE binding). + usage_rollups is a derived aggregate recomputed hourly from the AE SQL API + (direct D1 upserts only in local dev without the AE binding). Observability only; entitlements use separate source-table counts and daily counters. code: @@ -469,10 +469,10 @@ primitives: group: storage name: Email blob storage (R2) summary: - EMAIL_BLOBS R2 bucket holding raw email MIME under the frozen key - contract email-raw:v1:{userId}/{messageId}; D1 email_messages keeps - metadata plus raw_mime_key, with inline raw_mime as the fallback when - the binding is unavailable. Covered by account deletion and retention. + EMAIL_BLOBS R2 bucket holding raw email MIME under the frozen key contract + email-raw:v1:{userId}/{messageId}; D1 email_messages keeps metadata plus + raw_mime_key, with inline raw_mime as the fallback when the binding is + unavailable. Covered by account deletion and retention. code: - packages/worker/src/email/repo.ts docs: From 6ec32ca21bd7f7b2c40e131cdae9898e947922db Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 07:42:52 +0000 Subject: [PATCH 11/14] Provision EMAIL_BLOBS fail-soft when the API token lacks R2 access If R2 listing or creation fails, deploys drop the EMAIL_BLOBS binding from the generated Wrangler config and the worker falls back to inline raw MIME storage, instead of failing the resource-ensure run. Co-authored-by: Kent C. Dodds --- tools/ci/preview-resources.ts | 6 +- tools/ci/production-resources.ts | 6 +- tools/ci/resource-utils.node.test.ts | 34 +++++++++++ tools/ci/resource-utils.ts | 88 ++++++++++++++++++++++------ 4 files changed, 111 insertions(+), 23 deletions(-) diff --git a/tools/ci/preview-resources.ts b/tools/ci/preview-resources.ts index 910aae5206..7bac000630 100644 --- a/tools/ci/preview-resources.ts +++ b/tools/ci/preview-resources.ts @@ -276,7 +276,7 @@ async function ensurePreviewResources(options: CliOptions) { d1DatabaseId: d1.id, oauthKvId: oauthKv.id, bundleArtifactsKvId: bundleArtifactsKv.id, - emailBlobsBucketName: emailBlobs.name, + emailBlobsBucketName: emailBlobs.available ? emailBlobs.name : null, workerVars: { CLOUDFLARE_ACCOUNT_ID: process.env.CLOUDFLARE_ACCOUNT_ID, }, @@ -290,7 +290,9 @@ async function ensurePreviewResources(options: CliOptions) { console.log(`oauth_kv_id=${oauthKv.id}`) console.log(`bundle_artifacts_kv_title=${bundleArtifactsKv.title}`) console.log(`bundle_artifacts_kv_id=${bundleArtifactsKv.id}`) - console.log(`email_blobs_bucket_name=${emailBlobs.name}`) + console.log( + `email_blobs_bucket_name=${emailBlobs.available ? emailBlobs.name : ''}`, + ) } async function cleanupPreviewResources(options: CliOptions) { diff --git a/tools/ci/production-resources.ts b/tools/ci/production-resources.ts index 73e2b39861..590092a3db 100644 --- a/tools/ci/production-resources.ts +++ b/tools/ci/production-resources.ts @@ -387,7 +387,7 @@ async function ensureProductionResources(options: CliOptions) { d1DatabaseId: d1.id, oauthKvId: oauthKv.id, bundleArtifactsKvId: bundleArtifactsKv.id, - emailBlobsBucketName: emailBlobs.name, + emailBlobsBucketName: emailBlobs.available ? emailBlobs.name : null, workerVars: { APP_BASE_URL: process.env.APP_BASE_URL, CLOUDFLARE_ACCOUNT_ID: process.env.CLOUDFLARE_ACCOUNT_ID, @@ -402,7 +402,9 @@ async function ensureProductionResources(options: CliOptions) { console.log(`oauth_kv_id=${oauthKv.id}`) console.log(`bundle_artifacts_kv_title=${bundleArtifactsKv.title}`) console.log(`bundle_artifacts_kv_id=${bundleArtifactsKv.id}`) - console.log(`email_blobs_bucket_name=${emailBlobs.name}`) + console.log( + `email_blobs_bucket_name=${emailBlobs.available ? emailBlobs.name : ''}`, + ) } async function main() { diff --git a/tools/ci/resource-utils.node.test.ts b/tools/ci/resource-utils.node.test.ts index 735fcf0131..33222e4fc6 100644 --- a/tools/ci/resource-utils.node.test.ts +++ b/tools/ci/resource-utils.node.test.ts @@ -106,6 +106,40 @@ test('writeGeneratedWranglerConfig preserves migrations and copies environment a } }) +test('writeGeneratedWranglerConfig drops the EMAIL_BLOBS binding when the bucket is unavailable', async () => { + const tempDir = await mkdtemp(path.join(os.tmpdir(), 'kody-resource-utils-')) + + try { + const outPath = path.join(tempDir, 'wrangler-production.generated.json') + await writeGeneratedWranglerConfig({ + baseConfigPath: workerWranglerConfigPath, + outConfigPath: outPath, + envName: 'production', + d1DatabaseName: 'kody', + d1DatabaseId: 'dry-run-kody', + oauthKvId: 'dry-run-kody-oauth', + bundleArtifactsKvId: 'dry-run-kody-bundle-artifacts', + emailBlobsBucketName: null, + }) + + const config = parseJsonc<{ + env?: { + production?: { + r2_buckets?: Array<{ binding: string; bucket_name: string }> + } + } + }>(await readFile(outPath, 'utf8')) + // EMAIL_BLOBS is the only R2 binding, so the r2_buckets key must be + // removed entirely rather than left as an empty array. + expect(config.env?.production?.r2_buckets).toBeUndefined() + expect( + Object.keys(config.env?.production ?? {}).includes('r2_buckets'), + ).toBe(false) + } finally { + await rm(tempDir, { force: true, recursive: true }) + } +}) + test('parseR2BucketListOutput reads bucket names from labelled wrangler output', () => { const output = [ 'Listing buckets...', diff --git a/tools/ci/resource-utils.ts b/tools/ci/resource-utils.ts index 08c88f09b7..4dce0c950c 100644 --- a/tools/ci/resource-utils.ts +++ b/tools/ci/resource-utils.ts @@ -131,43 +131,74 @@ export function parseR2BucketListOutput(output: string): Array { return names } -export function listR2BucketNames(): Array { +/** + * Returns `null` when the listing fails (for example when the API token + * lacks R2 permissions) so callers can degrade gracefully instead of + * aborting the whole provisioning run. + */ +export function listR2BucketNames(): Array | null { const result = runWrangler(['r2', 'bucket', 'list'], { quiet: true }) if (result.status !== 0) { - fail('Failed to list R2 buckets (wrangler r2 bucket list).') + console.error('Failed to list R2 buckets (wrangler r2 bucket list).') + return null } return parseR2BucketListOutput(result.stdout) } +const r2ProvisioningWarning = + "Warning: R2 bucket provisioning failed (the API token may lack R2 permissions). Deploying without the EMAIL_BLOBS binding; raw email MIME will be stored inline in D1. Grant the token 'Workers R2 Storage: Edit' and re-run to enable blob offload." + +export type EnsureR2BucketResult = { + name: string + available: boolean +} + +/** + * The worker degrades gracefully when EMAIL_BLOBS is absent (raw MIME + * stays inline in D1), so R2 provisioning failures are non-fatal: the + * caller gets `available: false` and must drop the binding from the + * generated config instead of failing the deploy. + */ export function ensureR2Bucket({ name, dryRun, }: { name: string dryRun: boolean -}) { +}): EnsureR2BucketResult { if (dryRun) { console.error(`[dry-run] ensure R2 bucket: ${name}`) - return { name } + return { name, available: true } + } + + const bucketNames = listR2BucketNames() + if (bucketNames === null) { + console.error(r2ProvisioningWarning) + return { name, available: false } } - if (listR2BucketNames().includes(name)) { + if (bucketNames.includes(name)) { console.error(`R2 bucket exists: ${name}`) - return { name } + return { name, available: true } } const createResult = runWrangler(['r2', 'bucket', 'create', name], { quiet: true, }) if (createResult.status !== 0) { - fail(`Failed to create R2 bucket: ${name}`) + console.error(`Failed to create R2 bucket: ${name}`) + console.error(r2ProvisioningWarning) + return { name, available: false } } - if (!listR2BucketNames().includes(name)) { - fail(`Created R2 bucket "${name}" but could not find it via list.`) + const namesAfterCreate = listR2BucketNames() + if (namesAfterCreate === null || !namesAfterCreate.includes(name)) { + console.error(`Created R2 bucket "${name}" but could not find it via list.`) + console.error(r2ProvisioningWarning) + return { name, available: false } } console.error(`Created R2 bucket: ${name}`) - return { name } + return { name, available: true } } export function deleteR2Bucket({ @@ -182,7 +213,15 @@ export function deleteR2Bucket({ return } - if (!listR2BucketNames().includes(name)) { + const bucketNames = listR2BucketNames() + if (bucketNames === null) { + console.error( + `Warning: could not list R2 buckets while deleting ${name} (the API token may lack R2 permissions); skipping R2 cleanup.`, + ) + return + } + + if (!bucketNames.includes(name)) { console.error(`R2 bucket already deleted: ${name}`) return } @@ -387,7 +426,7 @@ export async function writeGeneratedWranglerConfig({ d1DatabaseId: string oauthKvId: string bundleArtifactsKvId: string - emailBlobsBucketName: string + emailBlobsBucketName: string | null workerVars?: Record extraMigrations?: Array }) { @@ -507,13 +546,24 @@ export async function writeGeneratedWranglerConfig({ ) } - const emailBlobsEntry = r2Buckets[emailBlobsEntryIndex] as Record< - string, - unknown - > - r2Buckets[emailBlobsEntryIndex] = { - ...emailBlobsEntry, - bucket_name: emailBlobsBucketName, + if (emailBlobsBucketName === null) { + // R2 provisioning was unavailable (for example, the API token lacks + // R2 permissions). Drop the binding so `wrangler deploy` does not + // reference a nonexistent bucket; the worker falls back to inline + // raw MIME storage in D1 when EMAIL_BLOBS is absent. + r2Buckets.splice(emailBlobsEntryIndex, 1) + if (r2Buckets.length === 0) { + delete (targetEnv as Record).r2_buckets + } + } else { + const emailBlobsEntry = r2Buckets[emailBlobsEntryIndex] as Record< + string, + unknown + > + r2Buckets[emailBlobsEntryIndex] = { + ...emailBlobsEntry, + bucket_name: emailBlobsBucketName, + } } const existingVars = (targetEnv as Record).vars From df72712e95d3ee72ea118a4ae1ab7d7950f7854b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 08:35:12 +0000 Subject: [PATCH 12/14] Address reviewer feedback: month bounds, blob-first deletes, search fallbacks Bound the usage aggregation query to the current UTC month on both sides and time out the AE fetch; delete system-email R2 blobs before their rows so failures can retry; run the community semantic fallback when the LIKE prefilter matched nothing; degrade package search to lexical scoring when the vector query fails; derive the repo-checks byte-cap message from its constant. Co-authored-by: Kent C. Dodds --- .../community/community-scoring.node.test.ts | 23 +-- .../community/community-service.node.test.ts | 133 +++++++++++++----- packages/worker/src/community/repo.ts | 4 +- packages/worker/src/community/service.ts | 26 ++-- packages/worker/src/email/system-email.ts | 119 ++++++++++------ .../src/email/system-email.workers.test.ts | 112 +++++++++++++++ .../worker/src/mcp/tools/search.node.test.ts | 72 ++++++++++ packages/worker/src/mcp/tools/search.ts | 18 ++- packages/worker/src/repo/checks.ts | 2 +- .../src/usage/aggregate-rollups.node.test.ts | 40 +++++- .../worker/src/usage/aggregate-rollups.ts | 37 ++++- 11 files changed, 476 insertions(+), 110 deletions(-) diff --git a/packages/worker/src/community/community-scoring.node.test.ts b/packages/worker/src/community/community-scoring.node.test.ts index 0e0e9ae472..637976e35f 100644 --- a/packages/worker/src/community/community-scoring.node.test.ts +++ b/packages/worker/src/community/community-scoring.node.test.ts @@ -1,4 +1,5 @@ import { expect, test, vi } from 'vitest' +import type * as CommunityRepo from './repo.ts' import { type CommunityListingRecord } from './types.ts' const mockModule = vi.hoisted(() => ({ @@ -7,14 +8,20 @@ const mockModule = vi.hoisted(() => ({ countCommunityForksByListingIds: vi.fn(), })) -vi.mock('./repo.ts', () => ({ - listCommunityListingCandidates: (...args: Array) => - mockModule.listCommunityListingCandidates(...args), - getCommunityRatingAggregatesByListingIds: (...args: Array) => - mockModule.getCommunityRatingAggregatesByListingIds(...args), - countCommunityForksByListingIds: (...args: Array) => - mockModule.countCommunityForksByListingIds(...args), -})) +vi.mock('./repo.ts', async (importOriginal) => { + const actual = await importOriginal() + return { + // Pure helper used by the service to decide whether the SQL LIKE + // pre-filter was applied; keep the real implementation. + extractCommunityListingLikeTokens: actual.extractCommunityListingLikeTokens, + listCommunityListingCandidates: (...args: Array) => + mockModule.listCommunityListingCandidates(...args), + getCommunityRatingAggregatesByListingIds: (...args: Array) => + mockModule.getCommunityRatingAggregatesByListingIds(...args), + countCommunityForksByListingIds: (...args: Array) => + mockModule.countCommunityForksByListingIds(...args), + } +}) const { COMMUNITY_SEARCH_CANDIDATE_LIMIT, diff --git a/packages/worker/src/community/community-service.node.test.ts b/packages/worker/src/community/community-service.node.test.ts index aadee1deae..82b1fb63c1 100644 --- a/packages/worker/src/community/community-service.node.test.ts +++ b/packages/worker/src/community/community-service.node.test.ts @@ -1,4 +1,5 @@ import { expect, test, vi } from 'vitest' +import type * as CommunityRepo from './repo.ts' import { type CommunityListingRecord } from './types.ts' const mockModule = vi.hoisted(() => ({ @@ -67,44 +68,50 @@ vi.mock('#worker/repo/artifact-repo-cleanup.ts', () => ({ mockModule.cleanupArtifactReposForPackage(...args), })) -vi.mock('./repo.ts', () => ({ - getCommunityBan: (...args: Array) => - mockModule.getCommunityBan(...args), - getCommunityListingByOwnerAndPackage: (...args: Array) => - mockModule.getCommunityListingByOwnerAndPackage(...args), - getCommunityListingById: (...args: Array) => - mockModule.getCommunityListingById(...args), - listCommunityListingCandidates: (...args: Array) => - mockModule.listCommunityListingCandidates(...args), - getCommunityRatingAggregatesByListingIds: (...args: Array) => - mockModule.getCommunityRatingAggregatesByListingIds(...args), - countCommunityForksByListingIds: (...args: Array) => - mockModule.countCommunityForksByListingIds(...args), - insertCommunityListing: (...args: Array) => - mockModule.insertCommunityListing(...args), - updateCommunityListing: (...args: Array) => - mockModule.updateCommunityListing(...args), - getCommunityForkByListingAndUser: (...args: Array) => - mockModule.getCommunityForkByListingAndUser(...args), - listCommunityForksByListingAndUser: (...args: Array) => - mockModule.listCommunityForksByListingAndUser(...args), - upsertCommunityRating: (...args: Array) => - mockModule.upsertCommunityRating(...args), - insertCommunityReport: (...args: Array) => - mockModule.insertCommunityReport(...args), - getCommunityReportById: (...args: Array) => - mockModule.getCommunityReportById(...args), - insertCommunityFork: (...args: Array) => - mockModule.insertCommunityFork(...args), - deleteCommunityListing: (...args: Array) => - mockModule.deleteCommunityListing(...args), - deleteCommunityRatingsByListingId: (...args: Array) => - mockModule.deleteCommunityRatingsByListingId(...args), - resolveCommunityReportRow: (...args: Array) => - mockModule.resolveCommunityReportRow(...args), - setCommunityListingStatus: (...args: Array) => - mockModule.setCommunityListingStatus(...args), -})) +vi.mock('./repo.ts', async (importOriginal) => { + const actual = await importOriginal() + return { + // Pure helper used by the service to decide whether the SQL LIKE + // pre-filter was applied; keep the real implementation. + extractCommunityListingLikeTokens: actual.extractCommunityListingLikeTokens, + getCommunityBan: (...args: Array) => + mockModule.getCommunityBan(...args), + getCommunityListingByOwnerAndPackage: (...args: Array) => + mockModule.getCommunityListingByOwnerAndPackage(...args), + getCommunityListingById: (...args: Array) => + mockModule.getCommunityListingById(...args), + listCommunityListingCandidates: (...args: Array) => + mockModule.listCommunityListingCandidates(...args), + getCommunityRatingAggregatesByListingIds: (...args: Array) => + mockModule.getCommunityRatingAggregatesByListingIds(...args), + countCommunityForksByListingIds: (...args: Array) => + mockModule.countCommunityForksByListingIds(...args), + insertCommunityListing: (...args: Array) => + mockModule.insertCommunityListing(...args), + updateCommunityListing: (...args: Array) => + mockModule.updateCommunityListing(...args), + getCommunityForkByListingAndUser: (...args: Array) => + mockModule.getCommunityForkByListingAndUser(...args), + listCommunityForksByListingAndUser: (...args: Array) => + mockModule.listCommunityForksByListingAndUser(...args), + upsertCommunityRating: (...args: Array) => + mockModule.upsertCommunityRating(...args), + insertCommunityReport: (...args: Array) => + mockModule.insertCommunityReport(...args), + getCommunityReportById: (...args: Array) => + mockModule.getCommunityReportById(...args), + insertCommunityFork: (...args: Array) => + mockModule.insertCommunityFork(...args), + deleteCommunityListing: (...args: Array) => + mockModule.deleteCommunityListing(...args), + deleteCommunityRatingsByListingId: (...args: Array) => + mockModule.deleteCommunityRatingsByListingId(...args), + resolveCommunityReportRow: (...args: Array) => + mockModule.resolveCommunityReportRow(...args), + setCommunityListingStatus: (...args: Array) => + mockModule.setCommunityListingStatus(...args), + } +}) vi.mock('./snapshot.ts', () => ({ writeCommunitySnapshot: (...args: Array) => @@ -420,6 +427,56 @@ test('searchCommunityListings empty query uses publishedAt tiebreaker', async () ]) }) +test('searchCommunityListings falls back to unfiltered candidates when LIKE prefilter rows all fail matching', async () => { + // The LIKE prefilter matches on raw columns (e.g. readme_content), so it + // can return rows that the in-memory scorer then rejects. The fallback + // must still surface matches among other recent listings. + const prefilterOnlyListing = sampleListing({ + id: 'listing-prefilter-only', + kodyId: 'meal-planner', + name: '@owner/meal-planner', + description: 'Plan weekly meals', + tags: ['meal'], + searchText: 'meal plan grocery', + readmeContent: '# Meal Planner\n\n## Intent\n\nPlan meals.', + }) + const fallbackMatchListing = sampleListing({ id: 'listing-fallback-match' }) + mockModule.listCommunityListingCandidates + .mockResolvedValueOnce([prefilterOnlyListing]) + .mockResolvedValueOnce([prefilterOnlyListing, fallbackMatchListing]) + mockModule.getCommunityRatingAggregatesByListingIds.mockResolvedValue({ + 'listing-fallback-match': { + listingId: 'listing-fallback-match', + ratingCount: 0, + averageStars: null, + averageAdaptationEffort: null, + }, + }) + mockModule.countCommunityForksByListingIds.mockResolvedValue({ + 'listing-fallback-match': 0, + }) + + const results = await searchCommunityListings({ + env: createEnv(), + query: 'discord', + limit: 10, + }) + + expect(results.map((listing) => listing.id)).toEqual([ + 'listing-fallback-match', + ]) + expect(mockModule.listCommunityListingCandidates).toHaveBeenNthCalledWith( + 1, + expect.anything(), + expect.objectContaining({ query: 'discord' }), + ) + expect(mockModule.listCommunityListingCandidates).toHaveBeenNthCalledWith( + 2, + expect.anything(), + expect.not.objectContaining({ query: expect.anything() }), + ) +}) + test('publishCommunityListing accepts Intent heading beyond storage truncation', async () => { const padding = 'x'.repeat(20_000) mockModule.getCommunityBan.mockResolvedValue(null) diff --git a/packages/worker/src/community/repo.ts b/packages/worker/src/community/repo.ts index e20248d51f..b92d33ea37 100644 --- a/packages/worker/src/community/repo.ts +++ b/packages/worker/src/community/repo.ts @@ -123,7 +123,9 @@ const communityListingSearchTextColumns = [ const maxCommunityListingLikeTokens = 8 -function extractCommunityListingLikeTokens(query: string): Array { +export function extractCommunityListingLikeTokens( + query: string, +): Array { // Tokens are restricted to [a-z0-9]+ so they are safe to embed inside // LIKE patterns without wildcard escaping. return Array.from( diff --git a/packages/worker/src/community/service.ts b/packages/worker/src/community/service.ts index 79590f33b9..b7854a1964 100644 --- a/packages/worker/src/community/service.ts +++ b/packages/worker/src/community/service.ts @@ -35,6 +35,7 @@ import { insertCommunityListing, insertCommunityBan, insertCommunityReport, + extractCommunityListingLikeTokens, listCommunityListingCandidates, listCommunityReports as listCommunityReportsFromDb, resolveCommunityReportRow, @@ -508,21 +509,26 @@ export async function searchCommunityListings(input: { .sort(compareCommunityListingsByBayesianAndPublishedAt) .slice(0, input.limit) } - if (listings.length === 0) { - // The SQL LIKE pre-filter found nothing; fall back to a bounded recent - // candidate set so vector-threshold matches keep working. + const matchesQuery = (listing: CommunityListingRecord) => + isCommunityListingSearchMatch({ + query: trimmedQuery, + document: buildListingSearchDocument(listing), + }) + let matched = listings.filter(matchesQuery) + const prefilterApplied = + extractCommunityListingLikeTokens(trimmedQuery).length > 0 + if (matched.length === 0 && prefilterApplied) { + // The SQL LIKE pre-filter produced no scoring matches; fall back to a + // bounded recent candidate set so vector-threshold matches among other + // recent listings keep working. Skipped when the first query was + // already unfiltered (no LIKE tokens), since it would return the same + // candidates. listings = await listCommunityListingCandidates(input.env.APP_DB, { includeDelisted: false, limit: COMMUNITY_SEARCH_CANDIDATE_LIMIT, }) + matched = listings.filter(matchesQuery) } - - const matched = listings.filter((listing) => - isCommunityListingSearchMatch({ - query: trimmedQuery, - document: buildListingSearchDocument(listing), - }), - ) const withAggregates = await attachListingAggregatesBatch( input.env.APP_DB, matched, diff --git a/packages/worker/src/email/system-email.ts b/packages/worker/src/email/system-email.ts index 855eb0c0aa..475de48121 100644 --- a/packages/worker/src/email/system-email.ts +++ b/packages/worker/src/email/system-email.ts @@ -216,7 +216,12 @@ async function deleteSystemEmailMessagesByIds(input: { blobs?: R2Bucket | null messageIds: ReadonlyArray }) { - let deletedRawMimeBlobs = 0 + const result = { + deletedMessages: 0, + deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 0, + } for ( let index = 0; index < input.messageIds.length; @@ -226,53 +231,67 @@ async function deleteSystemEmailMessagesByIds(input: { index, index + systemEmailDeleteIdsMaxParameters, ) - const placeholders = messageIds.map(() => '?').join(', ') - // Capture blob keys before the rows disappear. Blob cleanup is - // best-effort and must never block the D1 delete. - let rawMimeKeys: Array = [] - if (input.blobs) { - const { results } = await input.db - .prepare( - `SELECT raw_mime_key - FROM email_messages - WHERE id IN (${placeholders}) - AND raw_mime_key IS NOT NULL`, - ) - .bind(...messageIds) - .all<{ raw_mime_key: string }>() - rawMimeKeys = (results ?? []).map((row) => row.raw_mime_key) + const chunkPlaceholders = messageIds.map(() => '?').join(', ') + const { results } = await input.db + .prepare( + `SELECT id, raw_mime_key + FROM email_messages + WHERE id IN (${chunkPlaceholders}) + AND raw_mime_key IS NOT NULL`, + ) + .bind(...messageIds) + .all<{ id: string; raw_mime_key: string }>() + const blobRows = results ?? [] + // Delete R2 blobs before their rows, mirroring the user-mail retention + // policy: once a row is gone its raw_mime_key is lost, so a failed blob + // delete would orphan the object forever. Rows whose blob cannot be + // deleted are skipped and retried on a later run. + let deletableIds: ReadonlyArray = messageIds + if (blobRows.length > 0) { + const blobRowIds = new Set(blobRows.map((row) => row.id)) + if (!input.blobs) { + result.skippedMissingBlobBinding += blobRows.length + deletableIds = messageIds.filter((id) => !blobRowIds.has(id)) + console.warn('system-email-raw-mime-binding-missing', { + skipped: blobRows.length, + }) + } else { + try { + await input.blobs.delete(blobRows.map((row) => row.raw_mime_key)) + result.deletedRawMimeBlobs += blobRows.length + } catch (error) { + result.blobDeleteErrors += blobRows.length + deletableIds = messageIds.filter((id) => !blobRowIds.has(id)) + console.warn('system-email-raw-mime-blob-delete-failed', error) + } + } } + if (deletableIds.length === 0) continue + const placeholders = deletableIds.map(() => '?').join(', ') await input.db .prepare( `DELETE FROM email_attachments WHERE message_id IN (${placeholders})`, ) - .bind(...messageIds) + .bind(...deletableIds) .run() await input.db .prepare( `DELETE FROM email_delivery_events WHERE message_id IN (${placeholders})`, ) - .bind(...messageIds) + .bind(...deletableIds) .run() await input.db .prepare( `DELETE FROM email_messages WHERE id IN (${placeholders})`, ) - .bind(...messageIds) + .bind(...deletableIds) .run() - if (input.blobs && rawMimeKeys.length > 0) { - try { - await input.blobs.delete(rawMimeKeys) - deletedRawMimeBlobs += rawMimeKeys.length - } catch (error) { - console.warn('system-email-raw-mime-blob-delete-failed', error) - } - } + result.deletedMessages += deletableIds.length } - return { deletedMessages: input.messageIds.length, deletedRawMimeBlobs } + return result } export type SystemEmailRetentionResult = { @@ -281,6 +300,8 @@ export type SystemEmailRetentionResult = { deletedCounters: number deletedThreads: number deletedRawMimeBlobs: number + skippedMissingBlobBinding: number + blobDeleteErrors: number } export async function pruneSystemEmailRetention(input: { @@ -304,6 +325,20 @@ export async function pruneSystemEmailRetention(input: { deletedCounters: 0, deletedThreads: 0, deletedRawMimeBlobs: 0, + skippedMissingBlobBinding: 0, + blobDeleteErrors: 0, + } + + const addMessageDelete = (batch: { + deletedMessages: number + deletedRawMimeBlobs: number + skippedMissingBlobBinding: number + blobDeleteErrors: number + }) => { + result.deletedMessages += batch.deletedMessages + result.deletedRawMimeBlobs += batch.deletedRawMimeBlobs + result.skippedMissingBlobBinding += batch.skippedMissingBlobBinding + result.blobDeleteErrors += batch.blobDeleteErrors } const oldMessageIds = await listSystemEmailMessageIds({ @@ -311,26 +346,26 @@ export async function pruneSystemEmailRetention(input: { before: cutoffIso, limit: systemEmailLimits.pruneBatchSize, }) - const oldDelete = await deleteSystemEmailMessagesByIds({ - db: input.db, - blobs: input.blobs, - messageIds: oldMessageIds, - }) - result.deletedMessages += oldDelete.deletedMessages - result.deletedRawMimeBlobs += oldDelete.deletedRawMimeBlobs + addMessageDelete( + await deleteSystemEmailMessagesByIds({ + db: input.db, + blobs: input.blobs, + messageIds: oldMessageIds, + }), + ) const overCapMessageIds = await listSystemEmailMessageIds({ db: input.db, offset: systemEmailLimits.maxStoredMessages, limit: systemEmailLimits.pruneBatchSize, }) - const overCapDelete = await deleteSystemEmailMessagesByIds({ - db: input.db, - blobs: input.blobs, - messageIds: overCapMessageIds, - }) - result.deletedMessages += overCapDelete.deletedMessages - result.deletedRawMimeBlobs += overCapDelete.deletedRawMimeBlobs + addMessageDelete( + await deleteSystemEmailMessagesByIds({ + db: input.db, + blobs: input.blobs, + messageIds: overCapMessageIds, + }), + ) let deletedEventsInBatch = 0 do { diff --git a/packages/worker/src/email/system-email.workers.test.ts b/packages/worker/src/email/system-email.workers.test.ts index d77ea9739f..569c9a9c5c 100644 --- a/packages/worker/src/email/system-email.workers.test.ts +++ b/packages/worker/src/email/system-email.workers.test.ts @@ -336,6 +336,118 @@ test('system email retention prunes old operator-owned messages and counters', a ).toMatchObject({ results: [{ id: 'fresh-event' }] }) }) +test('system email retention deletes blobs before rows and keeps rows when the blob delete fails', async () => { + await ensureEmailTestSchema(env.APP_DB) + const now = new Date('2026-07-06T12:00:00.000Z') + const old = new Date( + now.getTime() - (systemEmailLimits.retentionDays + 1) * 24 * 60 * 60 * 1000, + ).toISOString() + const rawMimeKey = `email-raw:v1:${systemEmailOwnerId}/blob-ordering-message` + await env.EMAIL_BLOBS.put(rawMimeKey, 'raw mime payload') + await env.APP_DB.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, from_address, subject, processing_status, raw_mime_key, created_at, updated_at + ) VALUES + ('blob-ordering-message', 'inbound', ?, 'a@example.net', 'Blob', 'stored', ?, ?, ?), + ('plain-old-message', 'inbound', ?, 'b@example.net', 'Plain', 'stored', NULL, ?, ?)`, + ) + .bind( + systemEmailOwnerId, + rawMimeKey, + old, + old, + systemEmailOwnerId, + old, + old, + ) + .run() + const failingBlobs = new Proxy(env.EMAIL_BLOBS, { + get(target, property, receiver) { + if (property === 'delete') { + return async () => { + throw new Error('simulated R2 outage') + } + } + return Reflect.get(target, property, receiver) + }, + }) + + const failed = await pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: failingBlobs, + now, + }) + + // The blob-backed row survives the failed blob delete; only the plain row + // is pruned, and the blob is still readable for the retry. + expect(failed.deletedMessages).toBe(1) + expect(failed.deletedRawMimeBlobs).toBe(0) + expect(failed.blobDeleteErrors).toBe(1) + expect(await env.EMAIL_BLOBS.get(rawMimeKey)).not.toBeNull() + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_messages WHERE id = 'blob-ordering-message'`, + ).first(), + ).toMatchObject({ id: 'blob-ordering-message' }) + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_messages WHERE id = 'plain-old-message'`, + ).first(), + ).toBeNull() + + const retried = await pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + now, + }) + + expect(retried.deletedMessages).toBe(1) + expect(retried.deletedRawMimeBlobs).toBe(1) + expect(retried.blobDeleteErrors).toBe(0) + expect(await env.EMAIL_BLOBS.get(rawMimeKey)).toBeNull() + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_messages WHERE id = 'blob-ordering-message'`, + ).first(), + ).toBeNull() +}) + +test('system email retention skips blob-backed rows when the blobs binding is missing', async () => { + await ensureEmailTestSchema(env.APP_DB) + const now = new Date('2026-07-06T12:00:00.000Z') + const old = new Date( + now.getTime() - (systemEmailLimits.retentionDays + 1) * 24 * 60 * 60 * 1000, + ).toISOString() + const rawMimeKey = `email-raw:v1:${systemEmailOwnerId}/binding-missing-message` + await env.EMAIL_BLOBS.put(rawMimeKey, 'raw mime payload') + await env.APP_DB.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, from_address, subject, processing_status, raw_mime_key, created_at, updated_at + ) VALUES ('binding-missing-message', 'inbound', ?, 'a@example.net', 'Blob', 'stored', ?, ?, ?)`, + ) + .bind(systemEmailOwnerId, rawMimeKey, old, old) + .run() + + const result = await pruneSystemEmailRetention({ db: env.APP_DB, now }) + + expect(result.deletedMessages).toBe(0) + expect(result.skippedMissingBlobBinding).toBe(1) + expect(await env.EMAIL_BLOBS.get(rawMimeKey)).not.toBeNull() + expect( + await env.APP_DB.prepare( + `SELECT id FROM email_messages WHERE id = 'binding-missing-message'`, + ).first(), + ).toMatchObject({ id: 'binding-missing-message' }) + + // Clean up so the skipped row does not leak into other tests. + const cleanup = await pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + now, + }) + expect(cleanup.deletedMessages).toBe(1) +}) + test('system email retention drains delivery-event backlogs larger than one batch', async () => { await ensureEmailTestSchema(env.APP_DB) const now = new Date('2026-07-06T12:00:00.000Z') diff --git a/packages/worker/src/mcp/tools/search.node.test.ts b/packages/worker/src/mcp/tools/search.node.test.ts index d40e1201b7..948157976f 100644 --- a/packages/worker/src/mcp/tools/search.node.test.ts +++ b/packages/worker/src/mcp/tools/search.node.test.ts @@ -898,6 +898,78 @@ test('searchUnified ranks packages via user-filtered package vectors when Vector ).toBe(false) }) +test('searchUnified degrades to lexical package ranking when the vector query throws', async () => { + let packageVectorQueryAttempts = 0 + const env = { + SENTRY_ENVIRONMENT: 'production', + AI: createDeterministicAiBinding(), + CAPABILITY_VECTOR_INDEX: { + async query( + _values: Array, + options: { filter?: Record }, + ) { + const kind = (options.filter as { kind?: { $eq?: string } } | undefined) + ?.kind?.$eq + if (kind === 'package') { + packageVectorQueryAttempts += 1 + throw new Error('vectorize unavailable') + } + return { matches: [] } + }, + }, + } as unknown as Env + const packageRow: PackageSearchRow = { + record: { + id: 'pkg-inbox', + userId: 'user-1', + name: 'inbox-summarizer', + kodyId: 'inbox-summarizer', + description: 'summarize inbox threads', + tags: [], + searchText: null, + sourceId: 'source-pkg-inbox', + hasApp: false, + createdAt: '2026-04-20T00:00:00.000Z', + updatedAt: '2026-04-20T00:00:00.000Z', + }, + projection: { + name: 'inbox-summarizer', + kodyId: 'inbox-summarizer', + description: 'summarize inbox threads', + tags: [], + searchText: null, + hasApp: false, + appEntry: null, + exports: [], + jobs: [], + services: [], + subscriptions: [], + retrievers: [], + }, + readmeSnippet: null, + } + + const result = await searchUnified({ + env, + query: 'summarize inbox threads', + limit: 5, + registry: buildCapabilityRegistry([]), + optionalRows: { + packageRows: [packageRow], + userSecretRows: [], + userValueRows: [], + }, + }) + + expect(result.offline).toBe(false) + expect(packageVectorQueryAttempts).toBe(1) + expect( + result.matches.some( + (match) => match.type === 'package' && match.packageId === 'pkg-inbox', + ), + ).toBe(true) +}) + test('down remote connector statuses surface only disconnected connectors for signed-in users', async () => { const signedInContext = { user: { diff --git a/packages/worker/src/mcp/tools/search.ts b/packages/worker/src/mcp/tools/search.ts index cff1a1091f..32e1ea9a89 100644 --- a/packages/worker/src/mcp/tools/search.ts +++ b/packages/worker/src/mcp/tools/search.ts @@ -978,14 +978,26 @@ async function buildPackageCandidates(input: { }): Promise> { if (input.rows.length === 0) return [] const meaningfulTokens = extractMeaningfulSearchTokens(input.query) - const vectorScoresByRecordId = input.offline - ? null - : await queryPackageVectorScores({ + let vectorScoresByRecordId: Map | null = null + if (!input.offline) { + try { + vectorScoresByRecordId = await queryPackageVectorScores({ env: input.env, query: input.query, rows: input.rows, limit: input.limit, }) + } catch (error) { + // Degrade to lexical/deterministic ranking rather than failing the + // whole search when the vector index is unreachable. + console.warn( + JSON.stringify({ + message: 'package vector query failed, using lexical ranking', + error: error instanceof Error ? error.message : String(error), + }), + ) + } + } const candidates = input.rows .map((entry) => { const exports = Array.isArray(entry.projection.exports) diff --git a/packages/worker/src/repo/checks.ts b/packages/worker/src/repo/checks.ts index 0bb5b3caa8..aac7b87ebd 100644 --- a/packages/worker/src/repo/checks.ts +++ b/packages/worker/src/repo/checks.ts @@ -914,7 +914,7 @@ export async function runRepoChecks(input: { if (totalBytes > repoChecksSourceMaxTotalBytes) { return { ok: false as const, - message: `Repo checks aborted: source root "${sourceRoot || '/'}" exceeds the ${repoChecksSourceMaxTotalBytes}-byte (15 MiB) publish check limit. Remove or shrink large files that should not be published (for example build output, vendored dependencies, or data files) and run the checks again.`, + message: `Repo checks aborted: source root "${sourceRoot || '/'}" exceeds the ${repoChecksSourceMaxTotalBytes}-byte (${Math.round(repoChecksSourceMaxTotalBytes / (1024 * 1024))} MiB) publish check limit. Remove or shrink large files that should not be published (for example build output, vendored dependencies, or data files) and run the checks again.`, } } collected[path] = content diff --git a/packages/worker/src/usage/aggregate-rollups.node.test.ts b/packages/worker/src/usage/aggregate-rollups.node.test.ts index 4c9d734b78..13c4eebb29 100644 --- a/packages/worker/src/usage/aggregate-rollups.node.test.ts +++ b/packages/worker/src/usage/aggregate-rollups.node.test.ts @@ -158,9 +158,14 @@ test('aggregateUsageRollups queries Analytics Engine month-to-date and upserts a expect((init.headers as Record)['authorization']).toBe( 'Bearer token-1', ) + // A hung SQL API must abort instead of stalling the scheduled lane. + expect(init.signal).toBeInstanceOf(AbortSignal) const query = String(init.body) expect(query).toContain('FROM kody_usage_events') - expect(query).toContain(`toDateTime('2026-07-01 00:00:00')`) + // Half-open month bounds: a lower bound alone would let events stamped + // into a later month (clock skew, backdated writes) inflate this month. + expect(query).toContain(`timestamp >= toDateTime('2026-07-01 00:00:00')`) + expect(query).toContain(`timestamp < toDateTime('2026-08-01 00:00:00')`) // Sampling-correct aggregates: counts and sums weight by _sample_interval. expect(query).toContain('sum(_sample_interval) AS event_count') expect(query).toContain( @@ -211,12 +216,12 @@ test('aggregateUsageRollups honors CLOUDFLARE_API_BASE_URL and the preview datas CLOUDFLARE_API_BASE_URL: 'https://cloudflare-mock.local/', SENTRY_ENVIRONMENT: 'preview', }, - new Date('2026-07-15T10:00:00.000Z'), + new Date('2026-12-15T10:00:00.000Z'), ) expect(result).toEqual({ skipped: false, - month: '2026-07', + month: '2026-12', upsertedRows: 0, users: 0, }) @@ -227,7 +232,11 @@ test('aggregateUsageRollups honors CLOUDFLARE_API_BASE_URL and the preview datas expect(url).toBe( 'https://cloudflare-mock.local/client/v4/accounts/account-1/analytics_engine/sql', ) - expect(String(init.body)).toContain('FROM kody_usage_events_preview') + const query = String(init.body) + expect(query).toContain('FROM kody_usage_events_preview') + // The upper bound rolls over the UTC year boundary. + expect(query).toContain(`timestamp >= toDateTime('2026-12-01 00:00:00')`) + expect(query).toContain(`timestamp < toDateTime('2027-01-01 00:00:00')`) expect(batches).toHaveLength(0) }) @@ -262,3 +271,26 @@ test('aggregateUsageRollups batches large result sets and throws on SQL API erro ), ).rejects.toThrow('Analytics Engine SQL query failed (400)') }) + +test('aggregateUsageRollups surfaces a timed-out Analytics Engine fetch as an error', async () => { + // AbortSignal.timeout rejects the fetch with a TimeoutError DOMException; + // it must propagate through the same error path as a failed query. + const fetchMock = vi.fn(async () => { + throw new DOMException('The operation timed out.', 'TimeoutError') + }) + vi.stubGlobal('fetch', fetchMock) + using _restoreFetch = { + [Symbol.dispose]() { + vi.unstubAllGlobals() + }, + } + const { db, batches } = createFakeDb() + + await expect( + aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ), + ).rejects.toThrow('timed out') + expect(batches).toHaveLength(0) +}) diff --git a/packages/worker/src/usage/aggregate-rollups.ts b/packages/worker/src/usage/aggregate-rollups.ts index f8c9dfde88..e6a23962f2 100644 --- a/packages/worker/src/usage/aggregate-rollups.ts +++ b/packages/worker/src/usage/aggregate-rollups.ts @@ -56,6 +56,13 @@ export function resolveUsageEventsDataset(env: { const upsertBatchSize = 50 +/** + * Upper bound for one Analytics Engine SQL API round trip so a hung API can + * never stall the scheduled lane; a timeout aborts the fetch and surfaces + * through the same error path as a failed query. + */ +export const analyticsEngineSqlTimeoutMs = 30_000 + const usageRollupAbsoluteUpsertStatement = ` INSERT INTO usage_rollups ( user_id, metric, month, @@ -72,13 +79,35 @@ ON CONFLICT (user_id, metric, month) DO UPDATE SET updated_at = excluded.updated_at `.trim() +/** + * Half-open [current month start, next month start) UTC bounds for the SQL + * time filter, formatted for `toDateTime`. The explicit upper bound keeps + * events stamped into a later month (clock skew, backdated writes) from + * inflating the current month's rollups. + */ +function utcMonthBounds(now: Date) { + const toDateTimeArgument = (date: Date) => + `${date.toISOString().slice(0, 'YYYY-MM-DD'.length)} 00:00:00` + return { + monthStart: toDateTimeArgument( + new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)), + ), + nextMonthStart: toDateTimeArgument( + new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1)), + ), + } +} + /** * Analytics Engine samples data under load, so every aggregate must weight * by `_sample_interval`: counts are `sum(_sample_interval)` and value sums * are `sum(doubleN * _sample_interval)`. Blob/double positions match the * data point layout in `record-usage.ts`. */ -function buildMonthToDateAggregateQuery(dataset: string, monthStart: string) { +function buildMonthToDateAggregateQuery( + dataset: string, + bounds: { monthStart: string; nextMonthStart: string }, +) { return ` SELECT blob1 AS user_id, @@ -89,7 +118,8 @@ SELECT sum(double2 * _sample_interval) AS total_cpu_ms, sum(double3 * _sample_interval) AS total_bytes FROM ${dataset} -WHERE timestamp >= toDateTime('${monthStart}') +WHERE timestamp >= toDateTime('${bounds.monthStart}') + AND timestamp < toDateTime('${bounds.nextMonthStart}') GROUP BY blob1, blob2 FORMAT JSON `.trim() @@ -118,6 +148,7 @@ async function queryAnalyticsEngineSql(input: { authorization: `Bearer ${input.apiToken}`, }, body: input.query, + signal: AbortSignal.timeout(analyticsEngineSqlTimeoutMs), }) const text = await response.text() if (!response.ok) { @@ -164,7 +195,7 @@ export async function aggregateUsageRollups( env.CLOUDFLARE_API_BASE_URL?.trim() || 'https://api.cloudflare.com', query: buildMonthToDateAggregateQuery( resolveUsageEventsDataset(env), - `${month}-01 00:00:00`, + utcMonthBounds(now), ), }) From 1df1f6c3ff13476558ef9b46b4bfaf5696ef032f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 08:57:43 +0000 Subject: [PATCH 13/14] Address second review round: authoritative rollup recompute, cursor-advancing prunes The hourly rollup aggregation now deletes stale current-month rows absent from the AE result; email and system-email retention advance a keyset cursor past blob-skipped rows so a blocked head cannot wedge the queue; retention hasMore is based on selected rather than deleted rows; artifact reconcile pages strictly forward within a tick even when cursor writes fail; trade-off comments document the lean package search projection and the community browse candidate bound. Co-authored-by: Kent C. Dodds --- .../worker/src/app/retention.node.test.ts | 189 +++++++++++++++++- packages/worker/src/app/retention.ts | 132 +++++++----- packages/worker/src/community/service.ts | 6 +- packages/worker/src/email/system-email.ts | 67 +++++-- .../src/email/system-email.workers.test.ts | 79 ++++++++ .../reconcile-artifacts-pushes.node.test.ts | 114 +++++++++++ .../src/jobs/reconcile-artifacts-pushes.ts | 14 +- packages/worker/src/mcp/tools/search.ts | 5 + packages/worker/src/repo/entity-sources.ts | 35 +++- .../src/usage/aggregate-rollups.node.test.ts | 137 ++++++++++++- .../worker/src/usage/aggregate-rollups.ts | 121 ++++++++--- 11 files changed, 791 insertions(+), 108 deletions(-) diff --git a/packages/worker/src/app/retention.node.test.ts b/packages/worker/src/app/retention.node.test.ts index 396fbd0ff1..ad9d65b7bf 100644 --- a/packages/worker/src/app/retention.node.test.ts +++ b/packages/worker/src/app/retention.node.test.ts @@ -622,8 +622,14 @@ test('package invocation and workflow retention keeps boundary and active idempo .run(id, id, completedAt, status, completedAt, completedAt, completedAt) } - expect(await prunePackageInvocationsForRetention({ db, now })).toBe(2) - expect(await pruneWorkflowRunsForRetention({ db, now })).toBe(2) + expect(await prunePackageInvocationsForRetention({ db, now })).toEqual({ + selected: 2, + deleted: 2, + }) + expect(await pruneWorkflowRunsForRetention({ db, now })).toEqual({ + selected: 2, + deleted: 2, + }) expect(idsForTable(sqlite, 'package_invocations')).toEqual([ 'inv-boundary', @@ -685,9 +691,18 @@ test('memory suppression, email delivery, and audit retention respect boundaries .run(timestamp) } - expect(await pruneMemorySuppressionsForRetention({ db, now })).toBe(1) - expect(await pruneEmailDeliveryEventsForRetention({ db, now })).toBe(1) - expect(await pruneAuditEventsForRetention({ db, now })).toBe(1) + expect(await pruneMemorySuppressionsForRetention({ db, now })).toEqual({ + selected: 1, + deleted: 1, + }) + expect(await pruneEmailDeliveryEventsForRetention({ db, now })).toEqual({ + selected: 1, + deleted: 1, + }) + expect(await pruneAuditEventsForRetention({ db, now })).toEqual({ + selected: 1, + deleted: 1, + }) const memories = sqlite .prepare( @@ -770,6 +785,7 @@ test('email message retention deletes old user rows, R2 blobs, attachments, and skippedMissingBlobBinding: 0, blobDeleteErrors: 0, hasMore: false, + nextCursor: null, }) expect(blobDelete).toHaveBeenCalledWith(['raw-mime/user-1/msg-old-blob']) expect(idsForTable(sqlite, 'email_messages')).toEqual([ @@ -831,6 +847,153 @@ test('email message retention never deletes rows whose blob cannot be deleted fi expect(idsForTable(sqlite, 'email_messages')).toEqual([]) }) +test('email message retention cursor advances past skipped blob rows at the head', async () => { + const { sqlite, db } = createRetentionDb() + // The two oldest expired rows are blob-backed and unpassable without the + // EMAIL_BLOBS binding; the two plain expired rows behind them must still be + // pruned within the same run. + insertEmailMessage(sqlite, { + id: 'msg-blocked-a', + rawMimeKey: 'raw-mime/user-1/msg-blocked-a', + createdAt: daysAgo(emailMessageRetentionDays + 4), + }) + insertEmailMessage(sqlite, { + id: 'msg-blocked-b', + rawMimeKey: 'raw-mime/user-1/msg-blocked-b', + createdAt: daysAgo(emailMessageRetentionDays + 3), + }) + insertEmailMessage(sqlite, { + id: 'msg-plain-a', + createdAt: daysAgo(emailMessageRetentionDays + 2), + }) + insertEmailMessage(sqlite, { + id: 'msg-plain-b', + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + + const first = await pruneUserEmailMessagesForRetention({ + db, + now, + batchSize: 2, + }) + // Full batch selected but nothing deleted: hasMore must stay true and the + // cursor must point past the skipped rows. + expect(first).toMatchObject({ + deletedMessages: 0, + skippedMissingBlobBinding: 2, + hasMore: true, + }) + expect(first.nextCursor).toEqual({ + createdAt: daysAgo(emailMessageRetentionDays + 3), + id: 'msg-blocked-b', + }) + + const second = await pruneUserEmailMessagesForRetention({ + db, + now, + batchSize: 2, + cursor: first.nextCursor, + }) + expect(second).toMatchObject({ + deletedMessages: 2, + skippedMissingBlobBinding: 0, + hasMore: true, + }) + expect(idsForTable(sqlite, 'email_messages')).toEqual([ + 'msg-blocked-a', + 'msg-blocked-b', + ]) + + const third = await pruneUserEmailMessagesForRetention({ + db, + now, + batchSize: 2, + cursor: second.nextCursor, + }) + expect(third).toMatchObject({ deletedMessages: 0, hasMore: false }) +}) + +test('retention run prunes expired plain emails behind a skipped blob-row head', async () => { + const { sqlite, db } = createRetentionDb() + // 251 blob-backed rows fill the first default-size batch and spill into the + // second; 5 plain expired rows sit behind them. + for (let index = 0; index < 251; index += 1) { + insertEmailMessage(sqlite, { + id: `msg-blob-${String(index).padStart(3, '0')}`, + rawMimeKey: `raw-mime/user-1/${index}`, + createdAt: new Date( + now.getTime() - + (emailMessageRetentionDays + 10) * 24 * 60 * 60 * 1000 + + index * 1000, + ).toISOString(), + }) + } + for (let index = 0; index < 5; index += 1) { + insertEmailMessage(sqlite, { + id: `msg-plain-${index}`, + createdAt: daysAgo(emailMessageRetentionDays + 1), + }) + } + const env = { + APP_DB: db, + BUNDLE_ARTIFACTS_KV: { delete: vi.fn(async () => undefined) }, + } as unknown as Pick + + // No EMAIL_BLOBS binding: every blob row is skipped, yet the run-level + // cursor advances past them so the plain rows in the second batch are + // still pruned within the same run. + const result = await pruneRetention({ env, now }) + + expect(result.emailMessages.deletedMessages).toBe(5) + expect(result.emailMessages.skippedMissingBlobBinding).toBe(251) + expect(result.batchesPerTable['email_messages']).toBe(2) + expect(idsForTable(sqlite, 'email_messages')).toHaveLength(251) + expect(idsForTable(sqlite, 'email_messages')).not.toContain('msg-plain-0') +}) + +test('retention prune reports selected separately from deleted when rows vanish mid-batch', async () => { + const { sqlite } = createRetentionDb() + const baseDb = createD1FromSqlite(sqlite) + // Simulate a racing writer deleting one selected row before the batch + // DELETE runs: selected stays at the full batch size so hasMore-style + // decisions keep looping instead of marking the table drained. + const dbWithVanishingRow = { + prepare(query: string) { + const prepared = baseDb.prepare(query) + if (!query.includes('DELETE FROM package_invocations')) return prepared + return { + bind(...params: Array) { + const bound = prepared.bind(...params) + return { + async run() { + sqlite + .prepare(`DELETE FROM package_invocations WHERE id = 'inv-0'`) + .run() + return bound.run() + }, + all: bound.all, + first: bound.first, + } + }, + } + }, + } as unknown as D1Database + for (let index = 0; index < 2; index += 1) { + insertPackageInvocation(sqlite, { + id: `inv-${index}`, + createdAt: daysAgo(120), + }) + } + + expect( + await prunePackageInvocationsForRetention({ + db: dbWithVanishingRow, + now, + batchSize: 2, + }), + ).toEqual({ selected: 2, deleted: 1 }) +}) + test('entitlement counter and usage rollup retention respect boundaries', async () => { const { sqlite, db } = createRetentionDb() const cutoffDay = daysAgo(entitlementDailyCounterRetentionDays).slice(0, 10) @@ -858,8 +1021,14 @@ test('entitlement counter and usage rollup retention respect boundaries', async .run(month, now.toISOString()) } - expect(await pruneEntitlementDailyCountersForRetention({ db, now })).toBe(1) - expect(await pruneUsageRollupsForRetention({ db, now })).toBe(1) + 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`) @@ -1075,11 +1244,11 @@ test('retention pruning deletes only one configured batch per table invocation', expect( await prunePackageInvocationsForRetention({ db, now, batchSize: 2 }), - ).toBe(2) + ).toEqual({ selected: 2, deleted: 2 }) expect(idsForTable(sqlite, 'package_invocations')).toEqual(['inv-2']) expect( await prunePackageInvocationsForRetention({ db, now, batchSize: 2 }), - ).toBe(1) + ).toEqual({ selected: 1, deleted: 1 }) expect(idsForTable(sqlite, 'package_invocations')).toEqual([]) }) @@ -1100,7 +1269,7 @@ test('retention row deletes chunk ids to stay within the D1 binding limit', asyn expect( await prunePackageInvocationsForRetention({ db, now, batchSize: 101 }), - ).toBe(101) + ).toEqual({ selected: 101, deleted: 101 }) expect(idsForTable(sqlite, 'package_invocations')).toEqual([]) }) diff --git a/packages/worker/src/app/retention.ts b/packages/worker/src/app/retention.ts index 51bba96d32..0322b30f4c 100644 --- a/packages/worker/src/app/retention.ts +++ b/packages/worker/src/app/retention.ts @@ -257,6 +257,29 @@ async function selectIds(input: { return (results ?? []).map((row) => row[column] as IdValue) } +/** + * Runs a batched select-then-delete and reports both counts. `selected` (not + * `deleted`) is what callers should feed into hasMore decisions: a full batch + * that deletes fewer rows (racing writers) must not mark a table drained. + */ +async function selectAndDeleteByIds(input: { + db: D1Database + sql: string + bindings: ReadonlyArray + column?: string + table: string + idColumn: string +}): Promise<{ selected: number; deleted: number }> { + const ids = await selectIds(input) + const deleted = await deleteByIds({ + db: input.db, + table: input.table, + idColumn: input.idColumn, + ids, + }) + return { selected: ids.length, deleted } +} + async function deleteByIds(input: { db: D1Database table: string @@ -474,7 +497,7 @@ export async function prunePackageInvocationsForRetention(input: { input.now ?? new Date(), packageInvocationRetentionDays, ) - const ids = await selectIds({ + return selectAndDeleteByIds({ db: input.db, bindings: [cutoff, input.batchSize ?? retentionDefaultBatchSize], sql: `SELECT id @@ -483,12 +506,8 @@ export async function prunePackageInvocationsForRetention(input: { AND status != 'in_progress' ORDER BY created_at ASC, id ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'package_invocations', idColumn: 'id', - ids, }) } @@ -499,7 +518,7 @@ export async function pruneMemorySuppressionsForRetention(input: { }) { const now = input.now ?? new Date() const cutoff = cutoffIso(now, memorySuppressionRetentionDays) - const rowIds = await selectIds({ + return selectAndDeleteByIds({ db: input.db, column: 'rowid', bindings: [ @@ -513,12 +532,8 @@ export async function pruneMemorySuppressionsForRetention(input: { AND last_seen_at < ? ORDER BY last_seen_at ASC, rowid ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'mcp_memory_conversation_suppressions', idColumn: 'rowid', - ids: rowIds, }) } @@ -528,7 +543,7 @@ export async function pruneWorkflowRunsForRetention(input: { batchSize?: number }) { const cutoff = cutoffIso(input.now ?? new Date(), workflowRunRetentionDays) - const ids = await selectIds({ + return selectAndDeleteByIds({ db: input.db, bindings: [cutoff, input.batchSize ?? retentionDefaultBatchSize], sql: `SELECT id @@ -537,12 +552,8 @@ export async function pruneWorkflowRunsForRetention(input: { AND COALESCE(completed_at, updated_at, created_at) < ? ORDER BY COALESCE(completed_at, updated_at, created_at) ASC, id ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'workflow_runs', idColumn: 'id', - ids, }) } @@ -658,7 +669,7 @@ export async function pruneEmailDeliveryEventsForRetention(input: { input.now ?? new Date(), emailDeliveryEventRetentionDays, ) - const ids = await selectIds({ + return selectAndDeleteByIds({ db: input.db, bindings: [ systemEmailOwnerId, @@ -671,15 +682,23 @@ export async function pruneEmailDeliveryEventsForRetention(input: { AND created_at < ? ORDER BY created_at ASC, id ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'email_delivery_events', idColumn: 'id', - ids, }) } +/** + * Keyset cursor over (created_at, id) marking the last email message row + * examined (selected, not necessarily deleted). Threading it through batches + * within one run advances past rows skipped for blob reasons, so a blocked + * head cannot wedge the prune. Cross-run persistence is unnecessary: the next + * hourly run starts from the head and retries the skipped rows. + */ +export type EmailMessageRetentionCursor = { + createdAt: string + id: string +} + export type EmailMessageRetentionResult = { deletedMessages: number deletedAttachments: number @@ -688,6 +707,7 @@ export type EmailMessageRetentionResult = { skippedMissingBlobBinding: number blobDeleteErrors: number hasMore: boolean + nextCursor: EmailMessageRetentionCursor | null } export async function pruneUserEmailMessagesForRetention(input: { @@ -695,20 +715,34 @@ export async function pruneUserEmailMessagesForRetention(input: { blobs?: Pick | undefined now?: Date batchSize?: number + cursor?: EmailMessageRetentionCursor | null }): Promise { const cutoff = cutoffIso(input.now ?? new Date(), emailMessageRetentionDays) const batchSize = input.batchSize ?? retentionDefaultBatchSize + const cursor = input.cursor ?? null + const cursorClause = cursor + ? `AND (created_at > ? OR (created_at = ? AND id > ?))` + : '' + const cursorBindings = cursor + ? [cursor.createdAt, cursor.createdAt, cursor.id] + : [] const { results } = await input.db .prepare( - `SELECT id, user_id, raw_mime_key + `SELECT id, user_id, raw_mime_key, created_at FROM email_messages WHERE user_id != ? AND created_at < ? + ${cursorClause} ORDER BY created_at ASC, id ASC LIMIT ?`, ) - .bind(systemEmailOwnerId, cutoff, batchSize) - .all<{ id: string; user_id: string; raw_mime_key: string | null }>() + .bind(systemEmailOwnerId, cutoff, ...cursorBindings, batchSize) + .all<{ + id: string + user_id: string + raw_mime_key: string | null + created_at: string + }>() const rows = results ?? [] const result: EmailMessageRetentionResult = { deletedMessages: 0, @@ -718,6 +752,7 @@ export async function pruneUserEmailMessagesForRetention(input: { skippedMissingBlobBinding: 0, blobDeleteErrors: 0, hasMore: false, + nextCursor: null, } const rowsWithBlob = rows.filter((row) => row.raw_mime_key !== null) let deletableRows = rows @@ -782,12 +817,16 @@ export async function pruneUserEmailMessagesForRetention(input: { .run() result.deletedThreads += threadDelete.meta.changes ?? 0 } - // Skipped rows sit at the head of the oldest-first queue, so looping again - // in the same run would reselect them; wait for the next run instead. - result.hasMore = - rows.length >= batchSize && - result.skippedMissingBlobBinding === 0 && - result.blobDeleteErrors === 0 + // hasMore is based on the selected count: a full batch means more eligible + // rows may follow even when some rows were skipped rather than deleted. + // nextCursor points past every examined row so the next batch in this run + // moves forward instead of reselecting the skipped head. + const lastRow = rows.at(-1) + result.hasMore = rows.length >= batchSize + result.nextCursor = + result.hasMore && lastRow + ? { createdAt: lastRow.created_at, id: lastRow.id } + : null return result } @@ -800,7 +839,7 @@ export async function pruneEntitlementDailyCountersForRetention(input: { input.now ?? new Date(), entitlementDailyCounterRetentionDays, ).slice(0, 'YYYY-MM-DD'.length) - const rowIds = await selectIds({ + return selectAndDeleteByIds({ db: input.db, column: 'rowid', bindings: [cutoffDay, input.batchSize ?? retentionDefaultBatchSize], @@ -809,12 +848,8 @@ export async function pruneEntitlementDailyCountersForRetention(input: { WHERE day < ? ORDER BY day ASC, rowid ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'entitlement_daily_counters', idColumn: 'rowid', - ids: rowIds, }) } @@ -827,7 +862,7 @@ export async function pruneUsageRollupsForRetention(input: { input.now ?? new Date(), usageRollupRetentionMonths, ) - const rowIds = await selectIds({ + return selectAndDeleteByIds({ db: input.db, column: 'rowid', bindings: [oldestKeptMonth, input.batchSize ?? retentionDefaultBatchSize], @@ -836,12 +871,8 @@ export async function pruneUsageRollupsForRetention(input: { WHERE month < ? ORDER BY month ASC, rowid ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'usage_rollups', idColumn: 'rowid', - ids: rowIds, }) } @@ -851,7 +882,7 @@ export async function pruneAuditEventsForRetention(input: { batchSize?: number }) { const cutoff = cutoffIso(input.now ?? new Date(), auditEventRetentionDays) - const ids = await selectIds({ + return selectAndDeleteByIds({ db: input.db, bindings: [cutoff, input.batchSize ?? retentionDefaultBatchSize], sql: `SELECT id @@ -859,12 +890,8 @@ export async function pruneAuditEventsForRetention(input: { WHERE timestamp < ? ORDER BY timestamp ASC, id ASC LIMIT ?`, - }) - return deleteByIds({ - db: input.db, table: 'audit_events', idColumn: 'id', - ids, }) } @@ -882,6 +909,9 @@ export async function pruneRetention(input: { const startedAtMs = Date.now() const db = input.env.APP_DB const emailBlobs = input.env.EMAIL_BLOBS + // Per-run keyset cursor for the email prune; advances past rows skipped + // for blob reasons so a blocked head cannot wedge later batches. + let emailMessagesCursor: EmailMessageRetentionCursor | null = null const result: RetentionPruneResult = { runtimeRuns: { deletedRuns: 0, deletedLogs: 0, deletedOrphanLogs: 0 }, packageInvocations: 0, @@ -910,15 +940,17 @@ export async function pruneRetention(input: { } const countTask = ( table: string, - prune: () => Promise, - assign: (count: number) => void, + prune: () => Promise<{ selected: number; deleted: number }>, + assign: (deleted: number) => void, ) => ({ table, done: false, run: async () => { - const count = await prune() - assign(count) - return count >= retentionDefaultBatchSize + const batch = await prune() + assign(batch.deleted) + // hasMore keys off the selected count so a full batch that deleted + // fewer rows (racing writers) does not mark the table drained early. + return batch.selected >= retentionDefaultBatchSize }, }) const tasks: Array<{ @@ -989,6 +1021,7 @@ export async function pruneRetention(input: { db, blobs: emailBlobs, now, + cursor: emailMessagesCursor, }) result.emailMessages.deletedMessages += batch.deletedMessages result.emailMessages.deletedAttachments += batch.deletedAttachments @@ -997,6 +1030,7 @@ export async function pruneRetention(input: { result.emailMessages.skippedMissingBlobBinding += batch.skippedMissingBlobBinding result.emailMessages.blobDeleteErrors += batch.blobDeleteErrors + emailMessagesCursor = batch.nextCursor return batch.hasMore }, }, diff --git a/packages/worker/src/community/service.ts b/packages/worker/src/community/service.ts index b7854a1964..97dd45e221 100644 --- a/packages/worker/src/community/service.ts +++ b/packages/worker/src/community/service.ts @@ -72,7 +72,11 @@ const intentHeadingPattern = /^##\s+intent\b/im export const COMMUNITY_SEARCH_VECTOR_MATCH_THRESHOLD = 0.12 // Community search/browse never scores more than this many listings in -// memory; candidates are pre-filtered and recency-ordered in SQL. +// memory; candidates are pre-filtered and recency-ordered in SQL. This is +// a deliberate bound: browse ranks within the newest candidates, so +// listings older than the newest 500 are reachable through search (LIKE +// pre-filter) but not through unfiltered browse pagination. Revisit with +// a materialized score column if the catalog ever approaches this size. export const COMMUNITY_SEARCH_CANDIDATE_LIMIT = 500 export function isCommunityListingSearchMatch(input: { diff --git a/packages/worker/src/email/system-email.ts b/packages/worker/src/email/system-email.ts index 475de48121..7f729eeb6e 100644 --- a/packages/worker/src/email/system-email.ts +++ b/packages/worker/src/email/system-email.ts @@ -185,9 +185,20 @@ export async function countStoredSystemEmailMessages(input: { return Number(row?.count ?? 0) } -async function listSystemEmailMessageIds(input: { +type SystemEmailMessageCursor = { + createdAt: string + id: string +} + +async function listSystemEmailMessages(input: { db: D1Database before?: string + /** + * Keyset cursor (last row examined in the newest-first ordering) so batches + * within one prune run advance past rows skipped for blob reasons instead + * of reselecting the same blocked head. + */ + cursor?: SystemEmailMessageCursor | null offset?: number limit: number }) { @@ -197,18 +208,26 @@ async function listSystemEmailMessageIds(input: { where.push(`created_at < ?`) bindings.push(input.before) } + if (input.cursor) { + where.push(`(created_at < ? OR (created_at = ? AND id < ?))`) + bindings.push( + input.cursor.createdAt, + input.cursor.createdAt, + input.cursor.id, + ) + } const offset = input.offset ?? 0 const result = await input.db .prepare( - `SELECT id + `SELECT id, created_at FROM email_messages WHERE ${where.join(' AND ')} ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?`, ) .bind(...bindings, input.limit, offset) - .all<{ id: string }>() - return (result.results ?? []).map((row) => row.id) + .all<{ id: string; created_at: string }>() + return result.results ?? [] } async function deleteSystemEmailMessagesByIds(input: { @@ -341,20 +360,36 @@ export async function pruneSystemEmailRetention(input: { result.blobDeleteErrors += batch.blobDeleteErrors } - const oldMessageIds = await listSystemEmailMessageIds({ - db: input.db, - before: cutoffIso, - limit: systemEmailLimits.pruneBatchSize, - }) - addMessageDelete( - await deleteSystemEmailMessagesByIds({ + // Age prune loops batches within the time budget. The keyset cursor marks + // the last row examined (not the last deleted), so rows skipped because + // their blob could not be deleted never block later expired rows in the + // same run; the next run retries the skipped rows from the head. + let ageCursor: SystemEmailMessageCursor | null = null + let selectedInBatch = 0 + do { + const oldMessages = await listSystemEmailMessages({ db: input.db, - blobs: input.blobs, - messageIds: oldMessageIds, - }), + before: cutoffIso, + cursor: ageCursor, + limit: systemEmailLimits.pruneBatchSize, + }) + selectedInBatch = oldMessages.length + const lastMessage = oldMessages.at(-1) + if (!lastMessage) break + ageCursor = { createdAt: lastMessage.created_at, id: lastMessage.id } + addMessageDelete( + await deleteSystemEmailMessagesByIds({ + db: input.db, + blobs: input.blobs, + messageIds: oldMessages.map((message) => message.id), + }), + ) + } while ( + selectedInBatch >= systemEmailLimits.pruneBatchSize && + Date.now() < deadlineMs ) - const overCapMessageIds = await listSystemEmailMessageIds({ + const overCapMessages = await listSystemEmailMessages({ db: input.db, offset: systemEmailLimits.maxStoredMessages, limit: systemEmailLimits.pruneBatchSize, @@ -363,7 +398,7 @@ export async function pruneSystemEmailRetention(input: { await deleteSystemEmailMessagesByIds({ db: input.db, blobs: input.blobs, - messageIds: overCapMessageIds, + messageIds: overCapMessages.map((message) => message.id), }), ) diff --git a/packages/worker/src/email/system-email.workers.test.ts b/packages/worker/src/email/system-email.workers.test.ts index 569c9a9c5c..3aae11458e 100644 --- a/packages/worker/src/email/system-email.workers.test.ts +++ b/packages/worker/src/email/system-email.workers.test.ts @@ -448,6 +448,85 @@ test('system email retention skips blob-backed rows when the blobs binding is mi expect(cleanup.deletedMessages).toBe(1) }) +test('system email retention advances past skipped blob rows at the head of a batch', async () => { + await ensureEmailTestSchema(env.APP_DB) + const now = new Date('2026-07-06T12:00:00.000Z') + const oldMs = + now.getTime() - (systemEmailLimits.retentionDays + 1) * 24 * 60 * 60 * 1000 + // A full batch of blob-backed rows sits at the head of the newest-first + // expired ordering; the plain rows behind it are even older. + const blockedCount = systemEmailLimits.pruneBatchSize + const plainCount = 5 + const rows: Array<{ + id: string + rawMimeKey: string | null + createdAt: string + }> = [] + for (let index = 0; index < blockedCount; index += 1) { + rows.push({ + id: `head-blob-${String(index).padStart(3, '0')}`, + rawMimeKey: `email-raw:v1:${systemEmailOwnerId}/head-blob-${index}`, + createdAt: new Date(oldMs + index * 1000).toISOString(), + }) + } + for (let index = 0; index < plainCount; index += 1) { + rows.push({ + id: `head-plain-${index}`, + rawMimeKey: null, + createdAt: new Date(oldMs - 24 * 60 * 60 * 1000).toISOString(), + }) + } + // 5 bindings per row must stay under D1's 100-variable statement limit. + const insertChunkSize = 18 + for (let start = 0; start < rows.length; start += insertChunkSize) { + const chunk = rows.slice(start, start + insertChunkSize) + const values = chunk + .map( + () => `(?, 'inbound', ?, 'a@example.net', 'Head', 'stored', ?, ?, ?)`, + ) + .join(', ') + const bindings = chunk.flatMap((row) => [ + row.id, + systemEmailOwnerId, + row.rawMimeKey, + row.createdAt, + row.createdAt, + ]) + await env.APP_DB.prepare( + `INSERT INTO email_messages ( + id, direction, user_id, from_address, subject, processing_status, raw_mime_key, created_at, updated_at + ) VALUES ${values}`, + ) + .bind(...bindings) + .run() + } + + // No blobs binding: the whole first batch is skipped, but the keyset + // cursor advances so the plain rows in the next batch are still pruned + // within the same run. + const result = await pruneSystemEmailRetention({ db: env.APP_DB, now }) + + expect(result.skippedMissingBlobBinding).toBe(blockedCount) + expect(result.deletedMessages).toBe(plainCount) + const remainingPlain = await env.APP_DB.prepare( + `SELECT COUNT(*) AS count FROM email_messages WHERE id LIKE 'head-plain-%'`, + ).first<{ count: number }>() + expect(Number(remainingPlain?.count ?? -1)).toBe(0) + const remainingBlob = await env.APP_DB.prepare( + `SELECT COUNT(*) AS count FROM email_messages WHERE id LIKE 'head-blob-%'`, + ).first<{ count: number }>() + expect(Number(remainingBlob?.count ?? -1)).toBe(blockedCount) + + // Clean up the skipped rows with a working binding (deleting absent R2 + // keys is a no-op) so they do not leak into other tests. + const cleanup = await pruneSystemEmailRetention({ + db: env.APP_DB, + blobs: env.EMAIL_BLOBS, + now, + }) + expect(cleanup.deletedMessages).toBe(blockedCount) +}) + test('system email retention drains delivery-event backlogs larger than one batch', async () => { await ensureEmailTestSchema(env.APP_DB) const now = new Date('2026-07-06T12:00:00.000Z') diff --git a/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts b/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts index c9a319341d..667c455ff0 100644 --- a/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts +++ b/packages/worker/src/jobs/reconcile-artifacts-pushes.node.test.ts @@ -367,3 +367,117 @@ test('reconcile does not reprocess a source whose cursor write keeps failing', a mockModule.updateEntitySource.mockReset() mockModule.updateEntitySource.mockResolvedValue(true) }) + +test('head sources with failing cursor writes do not starve sources behind them within a tick', async () => { + // Simulated table: two head-of-queue sources whose cursor writes keep + // failing, plus one healthy source ordered behind them. + const rows = [ + source({ + id: 'source-a', + published_commit: 'commit-current', + created_at: '2026-05-04T00:00:00.000Z', + }), + source({ + id: 'source-b', + repo_id: 'repo-2', + published_commit: 'commit-current', + created_at: '2026-05-04T00:01:00.000Z', + }), + source({ + id: 'source-c', + repo_id: 'repo-3', + published_commit: 'commit-current', + created_at: '2026-05-04T00:02:00.000Z', + }), + ] + const checkedAt = new Map() + mockModule.listEntitySourcesForExternalReconcile.mockReset() + // Emulate the real query: staleness filter, keyset `after` cursor, and the + // stale ordering with id tie-break. + mockModule.listEntitySourcesForExternalReconcile.mockImplementation( + async (...args: Array) => { + const query = args[1] as { + before: string + limit: number + after?: { staleOrderedAt: string; id: string } + } + return rows + .map((row) => ({ + ...row, + last_external_check_at: + checkedAt.get(row.id) ?? row.last_external_check_at, + })) + .filter( + (row) => + row.last_external_check_at === null || + row.last_external_check_at < query.before, + ) + .map((row) => ({ + row, + key: row.last_external_check_at ?? row.created_at, + })) + .filter( + ({ row, key }) => + !query.after || + key > query.after.staleOrderedAt || + (key === query.after.staleOrderedAt && row.id > query.after.id), + ) + .sort((a, b) => + a.key === b.key + ? a.row.id.localeCompare(b.row.id) + : a.key.localeCompare(b.key), + ) + .slice(0, query.limit) + .map(({ row }) => row) + }, + ) + mockModule.resolveArtifactSourceHead.mockReset() + mockModule.resolveArtifactSourceHead.mockResolvedValue({ + branch: 'main', + commit: 'commit-current', + }) + mockModule.updateEntitySource.mockReset() + mockModule.updateEntitySource.mockImplementation( + async (...args: Array) => { + const update = args[1] as { id: string; lastExternalCheckAt: string } + if (update.id !== 'source-c') throw new Error('d1 unavailable') + checkedAt.set(update.id, update.lastExternalCheckAt) + return true + }, + ) + + const result = await reconcileArtifactsPushes({ + env: { APP_DB: {} } as Env, + baseUrl: 'https://kody.test', + now: new Date('2026-05-04T02:00:00.000Z'), + batchSize: 2, + }) + + // All three sources are reconciled in a single tick even though the first + // batch's rows never leave the stale window. + expect(result).toEqual( + expect.objectContaining({ + checked: 3, + alreadyPublished: 3, + errors: 2, + batches: 2, + budgetExhausted: false, + }), + ) + expect( + mockModule.listEntitySourcesForExternalReconcile, + ).toHaveBeenCalledTimes(2) + // The second batch query pages past the failing head rows via the keyset. + expect( + mockModule.listEntitySourcesForExternalReconcile, + ).toHaveBeenNthCalledWith(2, expect.anything(), { + before: '2026-05-04T01:55:00.000Z', + limit: 2, + after: { staleOrderedAt: '2026-05-04T00:01:00.000Z', id: 'source-b' }, + }) + // The healthy source's reconcile cursor is persisted. + expect(checkedAt.get('source-c')).toBe('2026-05-04T02:00:00.000Z') + + mockModule.updateEntitySource.mockReset() + mockModule.updateEntitySource.mockResolvedValue(true) +}) diff --git a/packages/worker/src/jobs/reconcile-artifacts-pushes.ts b/packages/worker/src/jobs/reconcile-artifacts-pushes.ts index b45038a3c9..093d1d4716 100644 --- a/packages/worker/src/jobs/reconcile-artifacts-pushes.ts +++ b/packages/worker/src/jobs/reconcile-artifacts-pushes.ts @@ -1,5 +1,6 @@ import { getErrorMessage } from '@kody-internal/shared/error-message.ts' import { + type ExternalReconcileCursor, listEntitySourcesForExternalReconcile, updateEntitySource, } from '#worker/repo/entity-sources.ts' @@ -73,11 +74,22 @@ export async function reconcileArtifactsPushes(input: { // the next batch query; tracking seen ids prevents re-processing them in a // tight loop within a single tick. const seenSourceIds = new Set() + let after: ExternalReconcileCursor | undefined while (!result.budgetExhausted) { const batch = await listEntitySourcesForExternalReconcile( input.env.APP_DB, - { before, limit: batchSize }, + { before, limit: batchSize, after }, ) + const lastRow = batch.at(-1) + if (lastRow) { + // Advance the keyset past every fetched row so a head row whose + // `last_external_check_at` write keeps failing cannot pin the front + // of the stale ordering and starve the sources behind it. + after = { + staleOrderedAt: lastRow.last_external_check_at ?? lastRow.created_at, + id: lastRow.id, + } + } const sources = batch.filter((source) => !seenSourceIds.has(source.id)) if (sources.length === 0) break result.batches += 1 diff --git a/packages/worker/src/mcp/tools/search.ts b/packages/worker/src/mcp/tools/search.ts index 32e1ea9a89..c61fb362f8 100644 --- a/packages/worker/src/mcp/tools/search.ts +++ b/packages/worker/src/mcp/tools/search.ts @@ -210,6 +210,11 @@ export type BuildSavedPackageSearchRowsResult = { * fields (exports, jobs, README) stay empty until a row is hydrated for a * top-ranked match. */ +// Ranking works on this lean D1 projection so broad searches never load +// full package source per package. Export/job/action recall is preserved +// online by the package vectors (built from the full document at publish) +// and by hydrating the top matches; the offline lexical path ranks on +// name/description/tags/searchText only. function buildLeanPackageSearchProjection( record: Awaited>[number], ): PackageSearchProjection { diff --git a/packages/worker/src/repo/entity-sources.ts b/packages/worker/src/repo/entity-sources.ts index 56f6c97418..0b1462751b 100644 --- a/packages/worker/src/repo/entity-sources.ts +++ b/packages/worker/src/repo/entity-sources.ts @@ -179,22 +179,53 @@ export async function upsertEntitySource( await insertEntitySource(db, row) } +/** + * Keyset cursor over the stale ordering used by + * {@link listEntitySourcesForExternalReconcile}: the coalesced + * `last_external_check_at`/`created_at` value of the last row in the previous + * batch, with the row id as a tie-breaker. + */ +export type ExternalReconcileCursor = { + staleOrderedAt: string + id: string +} + export async function listEntitySourcesForExternalReconcile( db: D1Database, input: { before: string limit: number + after?: ExternalReconcileCursor }, ): Promise> { + // The `after` keyset lets callers page strictly forward within one pass + // even when a row's `last_external_check_at` write fails and the row would + // otherwise stay at the head of the stale ordering. + const afterClause = input.after + ? `AND ( + COALESCE(last_external_check_at, created_at) > ? + OR (COALESCE(last_external_check_at, created_at) = ? AND id > ?) + )` + : '' + const bindings: Array = [input.before] + if (input.after) { + bindings.push( + input.after.staleOrderedAt, + input.after.staleOrderedAt, + input.after.id, + ) + } + bindings.push(input.limit) const { results } = await db .prepare( `SELECT * FROM entity_sources WHERE entity_kind = 'package' AND (last_external_check_at IS NULL OR last_external_check_at < ?) - ORDER BY COALESCE(last_external_check_at, created_at) ASC + ${afterClause} + ORDER BY COALESCE(last_external_check_at, created_at) ASC, id ASC LIMIT ?`, ) - .bind(input.before, input.limit) + .bind(...bindings) .all>() return (results ?? []).map(mapEntitySourceRow) } diff --git a/packages/worker/src/usage/aggregate-rollups.node.test.ts b/packages/worker/src/usage/aggregate-rollups.node.test.ts index 13c4eebb29..af38095ea4 100644 --- a/packages/worker/src/usage/aggregate-rollups.node.test.ts +++ b/packages/worker/src/usage/aggregate-rollups.node.test.ts @@ -7,13 +7,54 @@ import { type BoundStatement = { sql: string; params: Array } -function createFakeDb() { +type RollupKeyRow = { user_id: string; metric: string; month: string } + +function createFakeDb(input: { existingRollups?: Array } = {}) { + const rollups = input.existingRollups?.map((row) => ({ ...row })) ?? [] const batches: Array> = [] + const deletes: Array = [] const db = { prepare(sql: string) { return { bind(...params: Array) { - return { sql, params } + return { + sql, + params, + async all() { + if (!sql.includes('SELECT user_id, metric FROM usage_rollups')) { + throw new Error(`Unsupported all query: ${sql}`) + } + return { + results: rollups + .filter((row) => row.month === params[0]) + .map((row) => ({ + user_id: row.user_id, + metric: row.metric, + })), + } + }, + async run() { + if (!sql.startsWith('DELETE FROM usage_rollups')) { + throw new Error(`Unsupported run query: ${sql}`) + } + deletes.push({ sql, params }) + const month = params[0] + let changes = 0 + for (let index = 1; index < params.length; index += 2) { + const rowIndex = rollups.findIndex( + (row) => + row.month === month && + row.user_id === params[index] && + row.metric === params[index + 1], + ) + if (rowIndex >= 0) { + rollups.splice(rowIndex, 1) + changes += 1 + } + } + return { meta: { changes } } + }, + } }, } }, @@ -22,7 +63,7 @@ function createFakeDb() { return [] }, } as unknown as D1Database - return { db, batches } + return { db, batches, deletes, rollups } } function createAggregationEnv(db: D1Database) { @@ -73,7 +114,7 @@ test('resolveUsageEventsDataset picks the preview dataset only for preview', () test('aggregateUsageRollups no-ops when the binding or credentials are missing', async () => { using fetchMock = stubFetchResponse({ body: { data: [] } }) - const { db, batches } = createFakeDb() + const { db, batches, deletes } = createFakeDb() for (const env of [ { APP_DB: db }, @@ -96,6 +137,7 @@ test('aggregateUsageRollups no-ops when the binding or credentials are missing', } expect(fetchMock).not.toHaveBeenCalled() expect(batches).toHaveLength(0) + expect(deletes).toHaveLength(0) }) test('aggregateUsageRollups queries Analytics Engine month-to-date and upserts absolute values', async () => { @@ -143,6 +185,7 @@ test('aggregateUsageRollups queries Analytics Engine month-to-date and upserts a skipped: false, month: '2026-07', upsertedRows: 2, + deletedRows: 0, users: 2, }) @@ -223,6 +266,7 @@ test('aggregateUsageRollups honors CLOUDFLARE_API_BASE_URL and the preview datas skipped: false, month: '2026-12', upsertedRows: 0, + deletedRows: 0, users: 0, }) const [url, init] = fetchMock.mock.calls[0] as unknown as [ @@ -240,6 +284,91 @@ test('aggregateUsageRollups honors CLOUDFLARE_API_BASE_URL and the preview datas expect(batches).toHaveLength(0) }) +test('aggregateUsageRollups deletes current-month rows absent from the Analytics Engine result', async () => { + using _fetchMock = stubFetchResponse({ + body: { + data: [ + { + user_id: 'user-a', + metric: 'execute', + event_count: 5, + error_count: 0, + total_duration_ms: 0, + total_cpu_ms: 0, + total_bytes: 0, + }, + ], + }, + }) + const { db, batches, deletes, rollups } = createFakeDb({ + existingRollups: [ + // Present in the AE result: updated, never deleted. + { user_id: 'user-a', metric: 'execute', month: '2026-07' }, + // Absent from the AE result (for example a straggler from a + // direct D1 upsert whose AE data point was lost): deleted. + { user_id: 'user-a', metric: 'job_run', month: '2026-07' }, + { user_id: 'user-b', metric: 'execute', month: '2026-07' }, + // Other months are never touched, even for stale pairs. + { user_id: 'user-a', metric: 'job_run', month: '2026-06' }, + ], + }) + + const result = await aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ) + + expect(result).toEqual({ + skipped: false, + month: '2026-07', + upsertedRows: 1, + deletedRows: 2, + users: 1, + }) + // The present pair went through the upsert batch, not the delete. + expect(batches[0]?.[0]?.params?.slice(0, 3)).toEqual([ + 'user-a', + 'execute', + '2026-07', + ]) + expect(deletes).toHaveLength(1) + expect(deletes[0]?.sql).toContain('DELETE FROM usage_rollups WHERE month = ?') + expect(deletes[0]?.params).toEqual([ + '2026-07', + 'user-a', + 'job_run', + 'user-b', + 'execute', + ]) + expect(rollups).toEqual([ + { user_id: 'user-a', metric: 'execute', month: '2026-07' }, + { user_id: 'user-a', metric: 'job_run', month: '2026-06' }, + ]) +}) + +test('aggregateUsageRollups chunks stale-row deletes under the bind-parameter cap', async () => { + using _fetchMock = stubFetchResponse({ body: { data: [] } }) + const { db, deletes, rollups } = createFakeDb({ + existingRollups: Array.from({ length: 120 }, (_, index) => ({ + user_id: `stale-user-${index}`, + metric: 'execute', + month: '2026-07', + })), + }) + + const result = await aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ) + + expect(result).toMatchObject({ upsertedRows: 0, deletedRows: 120 }) + // 49 pairs per statement: 1 month param + 2 per pair = 99 binds max. + expect(deletes.map((statement) => statement.params.length)).toEqual([ + 99, 99, 45, + ]) + expect(rollups).toEqual([]) +}) + test('aggregateUsageRollups batches large result sets and throws on SQL API errors', async () => { const manyRows = Array.from({ length: 120 }, (_, index) => ({ user_id: `user-${index}`, diff --git a/packages/worker/src/usage/aggregate-rollups.ts b/packages/worker/src/usage/aggregate-rollups.ts index e6a23962f2..f75d1c39e0 100644 --- a/packages/worker/src/usage/aggregate-rollups.ts +++ b/packages/worker/src/usage/aggregate-rollups.ts @@ -6,9 +6,13 @@ * derived aggregate. `aggregateUsageRollups` recomputes the current UTC * month's rows from Analytics Engine via the SQL API and upserts them with * absolute values — an idempotent recompute, not an increment — so the hourly - * cron can run any number of times without drift. Analytics Engine retention - * (~90 days) always covers a full month, so a month-to-date recompute is - * complete; rows for prior months already in D1 stay untouched. + * cron can run any number of times without drift. The recompute is + * authoritative for the current month: rows whose (user, metric) pair is + * absent from the Analytics Engine result are deleted, so stragglers from + * direct D1 upserts (pre-cutover writes, local-dev rows) cannot linger. + * Analytics Engine retention (~90 days) always covers a full month, so a + * month-to-date recompute is complete; rows for prior months already in D1 + * stay untouched. * * In local dev and tests (no `USAGE_EVENTS` binding or Cloudflare REST * credentials) this is a no-op: `recordUsage` upserts `usage_rollups` @@ -40,7 +44,13 @@ export type UsageAggregationEnv = { export type UsageAggregationResult = | { skipped: true; reason: string } - | { skipped: false; month: string; upsertedRows: number; users: number } + | { + skipped: false + month: string + upsertedRows: number + deletedRows: number + users: number + } /** * The Analytics Engine SQL API dataset (= table name) written by the @@ -56,6 +66,13 @@ export function resolveUsageEventsDataset(env: { const upsertBatchSize = 50 +/** + * Pairs per DELETE statement. Each pair binds two parameters plus one for + * the month, so 49 pairs keeps every statement under D1's ~100 + * bind-parameter cap. + */ +const deleteStatementPairLimit = 49 + /** * Upper bound for one Analytics Engine SQL API round trip so a hung API can * never stall the scheduled lane; a timeout aborts the fetch and surfaces @@ -167,11 +184,62 @@ function toCount(value: number | string) { return Number.isFinite(parsed) ? Math.round(parsed) : 0 } +function rollupPairKey(row: { user_id: string; metric: string }) { + return `${row.user_id}\n${row.metric}` +} + +/** + * Delete current-month rollup rows whose (user_id, metric) pair is absent + * from the Analytics Engine result, making the recompute authoritative for + * the month (stragglers from direct D1 upserts — a pre-cutover write, a row + * whose Analytics Engine data point was lost — would otherwise linger + * forever). A NOT IN over the full pair set cannot be chunked without + * deleting everything, so stale pairs are computed in memory from one SELECT + * and deleted in parameter-capped chunks. Rows for other months are never + * touched. + */ +async function deleteStaleCurrentMonthRollups(input: { + db: D1Database + month: string + presentPairs: ReadonlySet +}) { + const { results } = await input.db + .prepare(`SELECT user_id, metric FROM usage_rollups WHERE month = ?`) + .bind(input.month) + .all<{ user_id: string; metric: string }>() + const stalePairs = (results ?? []).filter( + (row) => !input.presentPairs.has(rollupPairKey(row)), + ) + let deleted = 0 + for ( + let index = 0; + index < stalePairs.length; + index += deleteStatementPairLimit + ) { + const chunk = stalePairs.slice(index, index + deleteStatementPairLimit) + const pairPredicates = chunk + .map(() => '(user_id = ? AND metric = ?)') + .join(' OR ') + const result = await input.db + .prepare( + `DELETE FROM usage_rollups WHERE month = ? AND (${pairPredicates})`, + ) + .bind( + input.month, + ...chunk.flatMap((pair) => [pair.user_id, pair.metric]), + ) + .run() + deleted += result.meta.changes ?? 0 + } + return deleted +} + /** * Recompute the current UTC month's `usage_rollups` rows from Analytics - * Engine and upsert them into D1 with absolute values. No-op (with a debug - * log) when the Analytics Engine binding or the Cloudflare REST credentials - * are unavailable. + * Engine: upsert absolute values for every (user, metric) pair in the result + * and delete current-month rows absent from it. No-op (with a debug log) + * when the Analytics Engine binding or the Cloudflare REST credentials are + * unavailable. */ export async function aggregateUsageRollups( env: UsageAggregationEnv, @@ -200,32 +268,35 @@ export async function aggregateUsageRollups( }) const updatedAt = now.toISOString() - const statements = rows - .filter((row) => row.user_id && row.metric) - .map((row) => - env.APP_DB.prepare(usageRollupAbsoluteUpsertStatement).bind( - row.user_id, - row.metric, - month, - toCount(row.event_count), - toCount(row.error_count), - toCount(row.total_duration_ms), - toCount(row.total_cpu_ms), - toCount(row.total_bytes), - updatedAt, - ), - ) - const users = new Set( - rows.filter((row) => row.user_id && row.metric).map((row) => row.user_id), - ).size + const presentRows = rows.filter((row) => row.user_id && row.metric) + const statements = presentRows.map((row) => + env.APP_DB.prepare(usageRollupAbsoluteUpsertStatement).bind( + row.user_id, + row.metric, + month, + toCount(row.event_count), + toCount(row.error_count), + toCount(row.total_duration_ms), + toCount(row.total_cpu_ms), + toCount(row.total_bytes), + updatedAt, + ), + ) + const users = new Set(presentRows.map((row) => row.user_id)).size for (let index = 0; index < statements.length; index += upsertBatchSize) { await env.APP_DB.batch(statements.slice(index, index + upsertBatchSize)) } + const deletedRows = await deleteStaleCurrentMonthRollups({ + db: env.APP_DB, + month, + presentPairs: new Set(presentRows.map(rollupPairKey)), + }) const result = { skipped: false as const, month, upsertedRows: statements.length, + deletedRows, users, } console.info('usage-rollup-aggregation', JSON.stringify(result)) From b776471a533a31ddb6cc0fcdd52cfb72062753c0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 09:10:09 +0000 Subject: [PATCH 14/14] Guard rollup stale-delete against empty Analytics Engine results Co-authored-by: Kent C. Dodds --- .../src/usage/aggregate-rollups.node.test.ts | 45 ++++++++++++++++++- .../worker/src/usage/aggregate-rollups.ts | 17 ++++--- 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/packages/worker/src/usage/aggregate-rollups.node.test.ts b/packages/worker/src/usage/aggregate-rollups.node.test.ts index af38095ea4..2f100eb6c6 100644 --- a/packages/worker/src/usage/aggregate-rollups.node.test.ts +++ b/packages/worker/src/usage/aggregate-rollups.node.test.ts @@ -347,7 +347,21 @@ test('aggregateUsageRollups deletes current-month rows absent from the Analytics }) test('aggregateUsageRollups chunks stale-row deletes under the bind-parameter cap', async () => { - using _fetchMock = stubFetchResponse({ body: { data: [] } }) + using _fetchMock = stubFetchResponse({ + body: { + data: [ + { + user_id: 'user-live', + metric: 'execute', + event_count: 1, + error_count: 0, + total_duration_ms: 0, + total_cpu_ms: 0, + total_bytes: 0, + }, + ], + }, + }) const { db, deletes, rollups } = createFakeDb({ existingRollups: Array.from({ length: 120 }, (_, index) => ({ user_id: `stale-user-${index}`, @@ -361,7 +375,7 @@ test('aggregateUsageRollups chunks stale-row deletes under the bind-parameter ca new Date('2026-07-15T10:00:00.000Z'), ) - expect(result).toMatchObject({ upsertedRows: 0, deletedRows: 120 }) + expect(result).toMatchObject({ upsertedRows: 1, deletedRows: 120 }) // 49 pairs per statement: 1 month param + 2 per pair = 99 binds max. expect(deletes.map((statement) => statement.params.length)).toEqual([ 99, 99, 45, @@ -369,6 +383,33 @@ test('aggregateUsageRollups chunks stale-row deletes under the bind-parameter ca expect(rollups).toEqual([]) }) +test('aggregateUsageRollups keeps existing rollups when the Analytics Engine result is empty', async () => { + // An empty result is more likely ingestion lag or dataset + // misconfiguration than a real event-free month; the stale-row + // cleanup must not wipe the month's counters. + using _fetchMock = stubFetchResponse({ body: { data: [] } }) + const existingRollups = [ + { user_id: 'user-a', metric: 'execute', month: '2026-07' }, + { user_id: 'user-b', metric: 'job_run', month: '2026-07' }, + ] + const { db, deletes, rollups } = createFakeDb({ existingRollups }) + + const result = await aggregateUsageRollups( + createAggregationEnv(db), + new Date('2026-07-15T10:00:00.000Z'), + ) + + expect(result).toEqual({ + skipped: false, + month: '2026-07', + upsertedRows: 0, + deletedRows: 0, + users: 0, + }) + expect(deletes).toHaveLength(0) + expect(rollups).toEqual(existingRollups) +}) + test('aggregateUsageRollups batches large result sets and throws on SQL API errors', async () => { const manyRows = Array.from({ length: 120 }, (_, index) => ({ user_id: `user-${index}`, diff --git a/packages/worker/src/usage/aggregate-rollups.ts b/packages/worker/src/usage/aggregate-rollups.ts index f75d1c39e0..93e217ad6b 100644 --- a/packages/worker/src/usage/aggregate-rollups.ts +++ b/packages/worker/src/usage/aggregate-rollups.ts @@ -286,11 +286,18 @@ export async function aggregateUsageRollups( for (let index = 0; index < statements.length; index += upsertBatchSize) { await env.APP_DB.batch(statements.slice(index, index + upsertBatchSize)) } - const deletedRows = await deleteStaleCurrentMonthRollups({ - db: env.APP_DB, - month, - presentPairs: new Set(presentRows.map(rollupPairKey)), - }) + // An entirely empty result is more likely Analytics Engine ingestion + // lag (or a dataset misconfiguration) than a genuinely event-free + // month; skip the stale-row cleanup rather than wiping real counters. + // The next hourly run reconciles once events are visible again. + const deletedRows = + presentRows.length === 0 + ? 0 + : await deleteStaleCurrentMonthRollups({ + db: env.APP_DB, + month, + presentPairs: new Set(presentRows.map(rollupPairKey)), + }) const result = { skipped: false as const,