diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 46f2d24564..e4c4b99445 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -261,7 +261,7 @@ jobs: deploy: runs-on: ubuntu-latest name: 🚀 Deploy to production - timeout-minutes: 12 + timeout-minutes: 20 needs: - sha-guard - deploy-jobs-worker @@ -684,20 +684,43 @@ jobs: exit 0 fi base="${DEPLOY_URL%/}" - echo "POST ${base}/__maintenance/reindex-capabilities" - status="$(curl --silent --show-error --location --max-time 120 \ - -X POST \ - -H "Authorization: Bearer ${CAPABILITY_REINDEX_SECRET}" \ - -H "Accept: application/json" \ - --output reindex-response.json \ - --write-out "%{http_code}" \ - "${base}/__maintenance/reindex-capabilities")" - cat reindex-response.json - echo - if [ "$status" -lt 200 ] || [ "$status" -ge 300 ]; then - echo "Capability reindex failed with HTTP ${status}." >&2 - exit 1 - fi + max_sweeps=8 + sweep=1 + payload='{}' + while true; do + echo "POST ${base}/__maintenance/reindex-capabilities (sweep ${sweep})" + status="$(curl --silent --show-error --location --max-time 120 \ + -X POST \ + -H "Authorization: Bearer ${CAPABILITY_REINDEX_SECRET}" \ + -H "Accept: application/json" \ + -H "Content-Type: application/json" \ + --data "${payload}" \ + --output reindex-response.json \ + --write-out "%{http_code}" \ + "${base}/__maintenance/reindex-capabilities")" + cat reindex-response.json + echo + if [ "$status" -lt 200 ] || [ "$status" -ge 300 ]; then + echo "Capability reindex failed with HTTP ${status}." >&2 + exit 1 + fi + complete="$(jq -r '.complete // false' reindex-response.json)" + if [ "$complete" = "true" ]; then + echo "Capability reindex complete after ${sweep} sweep(s)." + break + fi + if [ "$sweep" -ge "$max_sweeps" ]; then + echo "Capability reindex did not finish after ${max_sweeps} sweeps." >&2 + exit 1 + fi + cursor="$(jq -c '.cursor // empty' reindex-response.json)" + if [ -z "$cursor" ] || [ "$cursor" = "null" ]; then + echo "Capability reindex is incomplete but returned no cursor." >&2 + exit 1 + fi + payload="$(jq -nc --argjson cursor "$cursor" '{cursor: $cursor}')" + sweep=$((sweep + 1)) + done deploy-status-worker: runs-on: ubuntu-latest diff --git a/docs/contributing/adding-capabilities.md b/docs/contributing/adding-capabilities.md index 3f9e7b5352..92a8c7fa19 100644 --- a/docs/contributing/adding-capabilities.md +++ b/docs/contributing/adding-capabilities.md @@ -400,7 +400,9 @@ the domain modules. keywords, or schemas, run the guarded `POST /__maintenance/reindex-capabilities` maintenance endpoint so Vectorize has fresh Workers AI embeddings. The endpoint also rebuilds memory, job, and - saved-package vectors with user-scoped metadata. + saved-package vectors with user-scoped metadata. Each POST is time-budgeted; + if the JSON response has `complete: false`, POST again with + `{ "cursor": }` until `complete` is true. Example (assuming `example` exists in `capabilityDomainNames`): diff --git a/docs/contributing/architecture/data-storage.md b/docs/contributing/architecture/data-storage.md index 3fcbc7473d..ad01449197 100644 --- a/docs/contributing/architecture/data-storage.md +++ b/docs/contributing/architecture/data-storage.md @@ -1348,7 +1348,10 @@ deterministic so upserts and deletes target the same vector. Search paths query only per-account namespaces plus the reserved builtin namespace. Vector rows are derived from D1 and can be rebuilt with the bounded `POST /__maintenance/reindex-capabilities` sweep, which keyset-pages memory, -job, and saved-package rows and rebuilds builtins in their reserved namespace. +job, and saved-package rows, rebuilds builtins in their reserved namespace, and +returns before a request-time budget. An incomplete sweep includes +`complete: false` and a `cursor` so the caller can POST again until `complete` +is true. Production deploy CI loops that endpoint after each production ship. ### `entity_sources` and package import contracts diff --git a/docs/contributing/environment-variables.md b/docs/contributing/environment-variables.md index 3a4fb75256..a664eacda2 100644 --- a/docs/contributing/environment-variables.md +++ b/docs/contributing/environment-variables.md @@ -282,10 +282,11 @@ Worker secrets: endpoints. Use the reindex endpoint after changing the embedding model, pooling, or Vectorize index dimensions; it rebuilds built-in capability, memory, job, and saved-package vectors with per-user `userId` metadata on - user-owned rows. Use the re-encrypt endpoint to rewrite remaining pre-AAD - (2-part) secret ciphertexts to v2 without rotating `SECRET_STORE_KEY` (see - [Secret rotation](./secret-rotation.md)). Local dev uses offline search while - `WRANGLER_IS_LOCAL_DEV` is set or the binding is missing. + user-owned rows. Each call is time-budgeted and may return `complete: false` + plus a `cursor` to resume. Use the re-encrypt endpoint to rewrite remaining + pre-AAD (2-part) secret ciphertexts to v2 without rotating `SECRET_STORE_KEY` + (see [Secret rotation](./secret-rotation.md)). Local dev uses offline search + while `WRANGLER_IS_LOCAL_DEV` is set or the binding is missing. - **`JOB_REINDEX_SECRET`** — optional Worker secret; bearer token for `POST /__maintenance/reindex-jobs` when you want a jobs-only Vectorize rebuild (without the full capability/memory/package reindex that diff --git a/docs/contributing/setup-manifest.md b/docs/contributing/setup-manifest.md index eee365ab30..38cef53387 100644 --- a/docs/contributing/setup-manifest.md +++ b/docs/contributing/setup-manifest.md @@ -681,10 +681,12 @@ How to get/set each value: - Generate a long random secret (for example `openssl rand -hex 32`), store it as the repository secret `CAPABILITY_REINDEX_SECRET`, and let the deploy workflow sync it to the Worker. After each production deploy, CI POSTs to - `/__maintenance/reindex-capabilities` with `Authorization: Bearer …` to - refresh built-in capability, memory, job, and saved-package embeddings. Run - the same POST manually after changing the embedding model, pooling, or - Vectorize index dimensions so existing rows are rebuilt with compatible + `/__maintenance/reindex-capabilities` with `Authorization: Bearer …` and + loops on the returned `cursor` until `complete` is true (each POST is + bounded to about 70 seconds; CI follows the cursor for up to 8 sweeps), + refreshing built-in capability, memory, job, and saved-package embeddings. + Run the same POST loop manually after changing the embedding model, pooling, + or Vectorize index dimensions so existing rows are rebuilt with compatible vectors. The same bearer authenticates `POST /__maintenance/reencrypt-secrets` (see [Secret rotation](./secret-rotation.md)). Local and preview environments can diff --git a/packages/worker/src/capability-maintenance.node.test.ts b/packages/worker/src/capability-maintenance.node.test.ts index d88ebbd260..ed51e09deb 100644 --- a/packages/worker/src/capability-maintenance.node.test.ts +++ b/packages/worker/src/capability-maintenance.node.test.ts @@ -47,22 +47,33 @@ function resetMocks() { mockModule.reindexSavedPackageVectors.mockReset() } -function createReindexRequest() { +function completeStep(upserted: number) { + return { upserted, complete: true, afterId: null } +} + +function createReindexRequest( + body?: Record, + input: { method?: string } = {}, +) { return new Request( 'https://kody.example.com/__maintenance/reindex-capabilities', { - method: 'POST', - headers: { Authorization: 'Bearer secret' }, + method: input.method ?? 'POST', + headers: { + Authorization: 'Bearer secret', + 'Content-Type': 'application/json', + }, + body: body === undefined ? undefined : JSON.stringify(body), }, ) } -test('capability reindex maintenance route rebuilds every vector kind', async () => { +test('capability reindex maintenance route rebuilds every vector kind and resumes incomplete sweeps', async () => { resetMocks() - mockModule.reindexCapabilityVectors.mockResolvedValue({ upserted: 3 }) - mockModule.reindexMemoryVectors.mockResolvedValue({ upserted: 2 }) - mockModule.reindexJobVectors.mockResolvedValue({ upserted: 1 }) - mockModule.reindexSavedPackageVectors.mockResolvedValue({ upserted: 4 }) + mockModule.reindexCapabilityVectors.mockResolvedValue(completeStep(3)) + mockModule.reindexMemoryVectors.mockResolvedValue(completeStep(2)) + mockModule.reindexJobVectors.mockResolvedValue(completeStep(1)) + mockModule.reindexSavedPackageVectors.mockResolvedValue(completeStep(4)) const env = { CAPABILITY_REINDEX_SECRET: 'secret', } as Env @@ -75,10 +86,11 @@ test('capability reindex maintenance route rebuilds every vector kind', async () expect(response.status).toBe(200) await expect(response.json()).resolves.toEqual({ ok: true, - capabilities: { upserted: 3 }, - memories: { upserted: 2 }, - jobs: { upserted: 1 }, - packages: { upserted: 4 }, + complete: true, + capabilities: completeStep(3), + memories: completeStep(2), + jobs: completeStep(1), + packages: completeStep(4), }) expect(mockModule.reindexCapabilityVectors).toHaveBeenCalledWith( env, @@ -87,21 +99,94 @@ test('capability reindex maintenance route rebuilds every vector kind', async () name: 'example_capability', }), }), + { + afterId: null, + deadlineMs: expect.any(Number), + }, ) - expect(mockModule.reindexMemoryVectors).toHaveBeenCalledWith(env) - expect(mockModule.reindexJobVectors).toHaveBeenCalledWith(env) + expect(mockModule.reindexMemoryVectors).toHaveBeenCalledWith(env, { + afterId: null, + deadlineMs: expect.any(Number), + }) + expect(mockModule.reindexJobVectors).toHaveBeenCalledWith(env, { + afterId: null, + deadlineMs: expect.any(Number), + }) expect(mockModule.reindexSavedPackageVectors).toHaveBeenCalledWith(env, { baseUrl: 'https://kody.example.com', + afterId: null, + deadlineMs: expect.any(Number), + }) + + resetMocks() + mockModule.reindexCapabilityVectors.mockResolvedValue(completeStep(3)) + mockModule.reindexMemoryVectors.mockResolvedValue({ + upserted: 8, + complete: false, + afterId: 'memory-8', + }) + const incompleteResponse = await handleCapabilityReindexRequest( + createReindexRequest({ timeBudgetMs: 5_000 }), + env, + ) + expect(incompleteResponse.status).toBe(200) + await expect(incompleteResponse.json()).resolves.toEqual({ + ok: true, + complete: false, + cursor: { phase: 'memories', afterId: 'memory-8' }, + capabilities: completeStep(3), + memories: { upserted: 8, complete: false, afterId: 'memory-8' }, + jobs: completeStep(0), + packages: completeStep(0), + }) + expect(mockModule.reindexJobVectors).not.toHaveBeenCalled() + expect(mockModule.reindexSavedPackageVectors).not.toHaveBeenCalled() + + resetMocks() + mockModule.reindexMemoryVectors.mockResolvedValue(completeStep(2)) + mockModule.reindexJobVectors.mockResolvedValue(completeStep(1)) + mockModule.reindexSavedPackageVectors.mockResolvedValue(completeStep(4)) + const resumeResponse = await handleCapabilityReindexRequest( + createReindexRequest({ + cursor: { phase: 'memories', afterId: 'memory-8' }, + }), + env, + ) + expect(resumeResponse.status).toBe(200) + await expect(resumeResponse.json()).resolves.toEqual({ + ok: true, + complete: true, + capabilities: completeStep(0), + memories: completeStep(2), + jobs: completeStep(1), + packages: completeStep(4), + }) + expect(mockModule.reindexCapabilityVectors).not.toHaveBeenCalled() + expect(mockModule.reindexMemoryVectors).toHaveBeenCalledWith(env, { + afterId: 'memory-8', + deadlineMs: expect.any(Number), + }) + + const invalidCursorResponse = await handleCapabilityReindexRequest( + createReindexRequest({ cursor: { phase: 'nope', afterId: null } }), + env, + ) + expect(invalidCursorResponse.status).toBe(400) + await expect(invalidCursorResponse.json()).resolves.toEqual({ + ok: false, + error: 'cursor.phase must be capabilities, memories, jobs, or packages.', }) }) test('capability reindex maintenance route attempts every vector kind before reporting failures', async () => { resetMocks() - mockModule.reindexCapabilityVectors.mockResolvedValue({ upserted: 3 }) + mockModule.reindexCapabilityVectors.mockResolvedValue(completeStep(3)) mockModule.reindexMemoryVectors.mockRejectedValue(new Error('memory failed')) - mockModule.reindexJobVectors.mockResolvedValue({ upserted: 1 }) + mockModule.reindexJobVectors.mockResolvedValue(completeStep(1)) mockModule.reindexSavedPackageVectors.mockResolvedValue({ upserted: 4, + complete: true, + afterId: null, failed: 1, error: '1 saved package vector(s) failed to reindex', }) @@ -117,11 +202,19 @@ test('capability reindex maintenance route attempts every vector kind before rep expect(response.status).toBe(500) await expect(response.json()).resolves.toEqual({ ok: false, - capabilities: { upserted: 3 }, - memories: { upserted: 0, error: 'memory failed' }, - jobs: { upserted: 1 }, + complete: true, + capabilities: completeStep(3), + memories: { + upserted: 0, + complete: false, + afterId: null, + error: 'memory failed', + }, + jobs: completeStep(1), packages: { upserted: 4, + complete: true, + afterId: null, failed: 1, error: '1 saved package vector(s) failed to reindex', }, diff --git a/packages/worker/src/capability-maintenance.ts b/packages/worker/src/capability-maintenance.ts index a57db35eee..549bc4ac4b 100644 --- a/packages/worker/src/capability-maintenance.ts +++ b/packages/worker/src/capability-maintenance.ts @@ -1,5 +1,6 @@ import { handleSecretMaintenanceRequest, + MaintenanceClientError, MaintenanceFailureError, } from './maintenance-handler.ts' import { getStaticRegistry } from './mcp/capabilities/registry.ts' @@ -8,17 +9,17 @@ import { reindexJobVectors } from './jobs/job-reindex.ts' import { reindexMemoryVectors } from './mcp/memory/memory-reindex.ts' import { reindexSavedPackageVectors } from './package-registry/package-reindex.ts' import { getErrorMessage } from '@kody-internal/shared/error-message.ts' +import { + capabilityReindexTimeBudgetMs, + hasReachedReindexDeadline, + isCapabilityReindexPhase, + type CapabilityReindexCursor, + type CapabilityReindexPhase, + type VectorReindexSweepResult, +} from '#worker/vectorize/reindex-sweep.ts' -type ReindexStepResult = { - upserted: number +type ReindexStepResult = VectorReindexSweepResult & { error?: string - failed?: number - warning?: string - failures?: Array<{ - id: string - phase: string - error: string - }> } type ReindexFailureSummary = { @@ -28,36 +29,189 @@ type ReindexFailureSummary = { failures?: ReindexStepResult['failures'] } +type CapabilityReindexKinds = { + capabilities: ReindexStepResult + memories: ReindexStepResult + jobs: ReindexStepResult + packages: ReindexStepResult +} + +const skippedReindexStep: ReindexStepResult = { + upserted: 0, + complete: true, + afterId: null, +} + async function runReindexStep( - run: () => Promise, + run: () => Promise, ): Promise { try { return await run() } catch (error) { return { upserted: 0, + complete: false, + afterId: null, error: getErrorMessage(error), } } } +function parseCapabilityReindexBody(body: unknown): + | { + ok: true + cursor: CapabilityReindexCursor | null + timeBudgetMs: number + } + | { ok: false; error: string } { + if (body == null) { + return { + ok: true, + cursor: null, + timeBudgetMs: capabilityReindexTimeBudgetMs, + } + } + if (typeof body !== 'object' || Array.isArray(body)) { + return { ok: false, error: 'Reindex body must be a JSON object.' } + } + const record = body as Record + let timeBudgetMs = capabilityReindexTimeBudgetMs + if (record.timeBudgetMs !== undefined) { + if ( + typeof record.timeBudgetMs !== 'number' || + !Number.isFinite(record.timeBudgetMs) || + record.timeBudgetMs < 0 + ) { + return { ok: false, error: 'timeBudgetMs must be a non-negative number.' } + } + timeBudgetMs = record.timeBudgetMs + } + if (record.cursor === undefined || record.cursor === null) { + return { ok: true, cursor: null, timeBudgetMs } + } + if (typeof record.cursor !== 'object' || Array.isArray(record.cursor)) { + return { ok: false, error: 'cursor must be an object.' } + } + const cursor = record.cursor as Record + if (!isCapabilityReindexPhase(cursor.phase)) { + return { + ok: false, + error: 'cursor.phase must be capabilities, memories, jobs, or packages.', + } + } + if (cursor.afterId !== null && typeof cursor.afterId !== 'string') { + return { ok: false, error: 'cursor.afterId must be a string or null.' } + } + return { + ok: true, + cursor: { phase: cursor.phase, afterId: cursor.afterId }, + timeBudgetMs, + } +} + +async function runReindexPhase( + env: Env, + input: { + baseUrl: string + phase: CapabilityReindexPhase + afterId: string | null + deadlineMs: number + }, +): Promise { + switch (input.phase) { + case 'capabilities': + return runReindexStep(() => + reindexCapabilityVectors(env, getStaticRegistry().capabilitySpecs, { + afterId: input.afterId, + deadlineMs: input.deadlineMs, + }), + ) + case 'memories': + return runReindexStep(() => + reindexMemoryVectors(env, { + afterId: input.afterId, + deadlineMs: input.deadlineMs, + }), + ) + case 'jobs': + return runReindexStep(() => + reindexJobVectors(env, { + afterId: input.afterId, + deadlineMs: input.deadlineMs, + }), + ) + case 'packages': + return runReindexStep(() => + reindexSavedPackageVectors(env, { + baseUrl: input.baseUrl, + afterId: input.afterId, + deadlineMs: input.deadlineMs, + }), + ) + default: { + const exhaustive: never = input.phase + throw new Error(`Unexpected reindex phase: ${exhaustive}`) + } + } +} + async function reindexAllCapabilitySearchVectors( env: Env, - input: { baseUrl: string }, + input: { + baseUrl: string + cursor: CapabilityReindexCursor | null + deadlineMs: number + }, ) { - const capabilities = await runReindexStep(() => - reindexCapabilityVectors(env, getStaticRegistry().capabilitySpecs), - ) - const memories = await runReindexStep(() => reindexMemoryVectors(env)) - const jobs = await runReindexStep(() => reindexJobVectors(env)) - const packages = await runReindexStep(() => - reindexSavedPackageVectors(env, { baseUrl: input.baseUrl }), - ) + const result: CapabilityReindexKinds = { + capabilities: skippedReindexStep, + memories: skippedReindexStep, + jobs: skippedReindexStep, + packages: skippedReindexStep, + } + const startPhase = input.cursor?.phase ?? 'capabilities' + let afterId = input.cursor?.afterId ?? null + let started = false + + for (const phase of [ + 'capabilities', + 'memories', + 'jobs', + 'packages', + ] as const) { + if (!started) { + if (phase !== startPhase) continue + started = true + } else if (hasReachedReindexDeadline(input.deadlineMs)) { + return { + ...result, + complete: false as const, + cursor: { phase, afterId: null }, + } + } else { + afterId = null + } + + const step = await runReindexPhase(env, { + baseUrl: input.baseUrl, + phase, + afterId, + deadlineMs: input.deadlineMs, + }) + result[phase] = step + if (step.error) continue + if (!step.complete) { + return { + ...result, + complete: false as const, + cursor: { phase, afterId: step.afterId }, + } + } + } + return { - capabilities, - memories, - jobs, - packages, + ...result, + complete: true as const, } } @@ -70,12 +224,26 @@ export async function handleCapabilityReindexRequest( secret: env.CAPABILITY_REINDEX_SECRET, notConfiguredMessage: 'Capability reindex is not configured', run: async () => { + const body = await request.json().catch(() => null) + const parsed = parseCapabilityReindexBody(body) + if (!parsed.ok) { + throw new MaintenanceClientError(parsed.error) + } const result = await reindexAllCapabilitySearchVectors(env, { baseUrl: new URL(request.url).origin, + cursor: parsed.cursor, + deadlineMs: Date.now() + parsed.timeBudgetMs, }) - const failed = Object.entries(result).filter( - (entry): entry is [string, ReindexStepResult & { error: string }] => - typeof entry[1].error === 'string', + const kindResults = ( + ['capabilities', 'memories', 'jobs', 'packages'] as const + ).map((phase) => [phase, result[phase]] as const) + const failed = kindResults.filter( + ( + entry, + ): entry is [ + CapabilityReindexPhase, + ReindexStepResult & { error: string }, + ] => typeof entry[1].error === 'string', ) if (failed.length > 0) { const failedPhases: Array = failed.map( diff --git a/packages/worker/src/jobs/job-reindex.node.test.ts b/packages/worker/src/jobs/job-reindex.node.test.ts index f68be97c39..5f35d9837c 100644 --- a/packages/worker/src/jobs/job-reindex.node.test.ts +++ b/packages/worker/src/jobs/job-reindex.node.test.ts @@ -83,6 +83,8 @@ test('job reindex keeps package job vector ids under the Vectorize limit', async await expect(reindexJobVectors({ APP_DB: {} } as Env)).resolves.toEqual({ upserted: 1, + complete: true, + afterId: null, }) expect(byteLength(`job_${rawJobId}`)).toBeGreaterThan(64) @@ -118,6 +120,8 @@ test('job reindex walks keyset pages and merges the page results', async () => { await expect(reindexJobVectors({ APP_DB: {} } as Env)).resolves.toEqual({ upserted: 201, + complete: true, + afterId: null, }) expect(mockModule.listJobRowsPage).toHaveBeenCalledTimes(2) @@ -154,7 +158,11 @@ test('job reindex retries transient D1 export page errors then surfaces exhausti try { const recoverPromise = reindexJobVectors({ APP_DB: {} } as Env) await vi.advanceTimersByTimeAsync(d1LockRetryBaseDelayMs) - await expect(recoverPromise).resolves.toEqual({ upserted: 1 }) + await expect(recoverPromise).resolves.toEqual({ + upserted: 1, + complete: true, + afterId: null, + }) } finally { vi.useRealTimers() } diff --git a/packages/worker/src/jobs/job-reindex.ts b/packages/worker/src/jobs/job-reindex.ts index 7b536e99c8..49718b33d1 100644 --- a/packages/worker/src/jobs/job-reindex.ts +++ b/packages/worker/src/jobs/job-reindex.ts @@ -3,10 +3,10 @@ import { isCapabilitySearchOffline, } from '#worker/vectorize/embedding.ts' import { - mergeVectorReindexResults, - reindexVectorCandidates, - type VectorReindexResult, -} from '#worker/vectorize/reindex-batches.ts' + reindexPagedVectorRows, + type VectorReindexSweepOptions, + type VectorReindexSweepResult, +} from '#worker/vectorize/reindex-sweep.ts' import { userVectorNamespace } from '#worker/vectorize/vector-namespaces.ts' import { buildJobEmbedText } from '#mcp/jobs-embed.ts' import { jobVectorId } from '#mcp/jobs-vectorize.ts' @@ -20,53 +20,45 @@ const reindexPageSize = 200 export async function reindexJobVectors( env: Env, -): Promise { + options?: VectorReindexSweepOptions, +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') } if (isCapabilitySearchOffline(env)) { - return { upserted: 0 } + return { upserted: 0, complete: true, afterId: null } } - const pageResults: Array = [] - let afterId: string | null = null - while (true) { - const rows = await runD1WithRetry(() => - jobsData(env).listJobsPage({ - 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, - }), - namespace: userVectorNamespace(row.user_id), - metadata: { kind: 'job', userId: row.user_id }, - } - }) - if (candidates.length > 0) { - pageResults.push( - await reindexVectorCandidates({ - env, - index, - kind: 'job', - candidates, + return reindexPagedVectorRows({ + env, + index, + kind: 'job', + pageSize: reindexPageSize, + afterId: options?.afterId, + deadlineMs: options?.deadlineMs, + listPage: ({ afterId, limit }) => + runD1WithRetry(() => + jobsData(env).listJobsPage({ + afterId, + limit, }), - ) - } - if (rows.length < reindexPageSize) break - afterId = rows[rows.length - 1]!.id - } - return mergeVectorReindexResults('job', pageResults) + ), + rowId: (row) => row.id, + toCandidate: (row) => { + if (!row.user_id) return null + const view = toJobView(row.record) + return { + id: jobVectorId(row.id), + text: buildJobEmbedText({ + name: view.name, + scheduleSummary: view.scheduleSummary, + sourceId: view.sourceId, + publishedCommit: view.publishedCommit, + }), + namespace: userVectorNamespace(row.user_id), + metadata: { kind: 'job', userId: row.user_id }, + } + }, + }) } diff --git a/packages/worker/src/maintenance-handler.node.test.ts b/packages/worker/src/maintenance-handler.node.test.ts index f10afb0c92..84ef2eca52 100644 --- a/packages/worker/src/maintenance-handler.node.test.ts +++ b/packages/worker/src/maintenance-handler.node.test.ts @@ -1,6 +1,7 @@ import { expect, test } from 'vitest' import { handleSecretMaintenanceRequest, + MaintenanceClientError, timingSafeEqualString, } from './maintenance-handler.ts' @@ -85,6 +86,21 @@ test('handleSecretMaintenanceRequest enforces auth and reports maintenance resul ok: false, error: 'boom', }) + + const clientErrorResponse = await handleSecretMaintenanceRequest({ + request: createRequest({ authorization: 'Bearer secret' }), + secret: 'secret', + notConfiguredMessage: 'Not configured', + run: async () => { + throw new MaintenanceClientError('bad cursor') + }, + }) + + expect(clientErrorResponse.status).toBe(400) + await expect(clientErrorResponse.json()).resolves.toEqual({ + ok: false, + error: 'bad cursor', + }) }) test('timingSafeEqualString accepts equal secrets and rejects mismatches', async () => { diff --git a/packages/worker/src/maintenance-handler.ts b/packages/worker/src/maintenance-handler.ts index d10e0861b1..6c32bebabb 100644 --- a/packages/worker/src/maintenance-handler.ts +++ b/packages/worker/src/maintenance-handler.ts @@ -12,6 +12,13 @@ export class MaintenanceFailureError extends Error { } } +export class MaintenanceClientError extends Error { + constructor(message: string) { + super(message) + this.name = 'MaintenanceClientError' + } +} + type SecretMaintenanceRequestInput = { request: Request secret: string | null | undefined @@ -84,6 +91,9 @@ export async function handleSecretMaintenanceRequest( const result = await input.run() return Response.json({ ...result, ok: true }) } catch (error) { + if (error instanceof MaintenanceClientError) { + return Response.json({ ok: false, error: error.message }, { status: 400 }) + } if (error instanceof MaintenanceFailureError) { return Response.json( { ...error.result, ok: false, error: error.message }, diff --git a/packages/worker/src/mcp/capabilities/capability-reindex.ts b/packages/worker/src/mcp/capabilities/capability-reindex.ts index a6cc3f348c..3bb685f9cb 100644 --- a/packages/worker/src/mcp/capabilities/capability-reindex.ts +++ b/packages/worker/src/mcp/capabilities/capability-reindex.ts @@ -1,30 +1,39 @@ import { buildCapabilityEmbedText } from './capability-search.ts' import { getCapabilityVectorIndex } from '#worker/vectorize/embedding.ts' import { - reindexVectorCandidates, - type VectorReindexResult, -} from '#worker/vectorize/reindex-batches.ts' + reindexVectorCandidateList, + type VectorReindexSweepOptions, + type VectorReindexSweepResult, +} from '#worker/vectorize/reindex-sweep.ts' import { BUILTIN_VECTOR_NAMESPACE } from '#worker/vectorize/vector-namespaces.ts' import { type CapabilitySpec } from './types.ts' export async function reindexCapabilityVectors( env: Env, specs: Record, -): Promise { + options?: VectorReindexSweepOptions, +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') } - return reindexVectorCandidates({ - env, - index, - kind: 'builtin capability', - candidates: Object.values(specs).map((spec) => ({ + const candidates = Object.values(specs) + .slice() + .sort((left, right) => left.name.localeCompare(right.name)) + .map((spec) => ({ id: spec.name, text: buildCapabilityEmbedText(spec), namespace: BUILTIN_VECTOR_NAMESPACE, metadata: { domain: spec.domain, kind: 'builtin' }, - })), + })) + + return reindexVectorCandidateList({ + env, + index, + kind: 'builtin capability', + candidates, + afterId: options?.afterId, + deadlineMs: options?.deadlineMs, }) } diff --git a/packages/worker/src/mcp/memory/memory-reindex.node.test.ts b/packages/worker/src/mcp/memory/memory-reindex.node.test.ts index e9a1e7eb52..a4e0e105f8 100644 --- a/packages/worker/src/mcp/memory/memory-reindex.node.test.ts +++ b/packages/worker/src/mcp/memory/memory-reindex.node.test.ts @@ -78,6 +78,8 @@ test('memory reindex walks keyset pages and merges the page results', async () = await expect(reindexMemoryVectors({ APP_DB: {} } as Env)).resolves.toEqual({ upserted: 201, + complete: true, + afterId: null, }) expect(mockModule.listMemoriesPage).toHaveBeenCalledTimes(2) @@ -107,6 +109,8 @@ test('memory reindex retries transient D1 export page errors then surfaces exhau mockModule.listMemoriesPage.mockResolvedValueOnce([]) await expect(reindexMemoryVectors({ APP_DB: {} } as Env)).resolves.toEqual({ upserted: 0, + complete: true, + afterId: null, }) resetMocks() @@ -123,7 +127,11 @@ test('memory reindex retries transient D1 export page errors then surfaces exhau try { const recoverPromise = reindexMemoryVectors({ APP_DB: {} } as Env) await vi.advanceTimersByTimeAsync(d1LockRetryBaseDelayMs) - await expect(recoverPromise).resolves.toEqual({ upserted: 1 }) + await expect(recoverPromise).resolves.toEqual({ + upserted: 1, + complete: true, + afterId: null, + }) } finally { vi.useRealTimers() } diff --git a/packages/worker/src/mcp/memory/memory-reindex.ts b/packages/worker/src/mcp/memory/memory-reindex.ts index 3efa470e01..1a42cd4109 100644 --- a/packages/worker/src/mcp/memory/memory-reindex.ts +++ b/packages/worker/src/mcp/memory/memory-reindex.ts @@ -3,10 +3,10 @@ import { isCapabilitySearchOffline, } from '#worker/vectorize/embedding.ts' import { - mergeVectorReindexResults, - reindexVectorCandidates, - type VectorReindexResult, -} from '#worker/vectorize/reindex-batches.ts' + reindexPagedVectorRows, + type VectorReindexSweepOptions, + type VectorReindexSweepResult, +} from '#worker/vectorize/reindex-sweep.ts' import { userVectorNamespace } from '#worker/vectorize/vector-namespaces.ts' import { runD1WithRetry } from '#worker/d1-retry.ts' import { buildMemoryEmbedTextFromRow } from './memory-embed.ts' @@ -19,46 +19,42 @@ const reindexPageSize = 200 export async function reindexMemoryVectors( env: Env, -): Promise { + options?: VectorReindexSweepOptions, +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') } if (isCapabilitySearchOffline(env)) { - return { upserted: 0 } + return { upserted: 0, complete: true, afterId: null } } - const pageResults: Array = [] - let afterId: string | null = null - while (true) { - const rows = await runD1WithRetry(() => - listMemoriesPage({ - db: env.APP_DB, - afterId, - limit: reindexPageSize, - }), - ) - if (rows.length === 0) break - pageResults.push( - await reindexVectorCandidates({ - env, - index, + return reindexPagedVectorRows({ + env, + index, + kind: 'memory', + pageSize: reindexPageSize, + afterId: options?.afterId, + deadlineMs: options?.deadlineMs, + listPage: ({ afterId, limit }) => + runD1WithRetry(() => + listMemoriesPage({ + db: env.APP_DB, + afterId, + limit, + }), + ), + rowId: (row) => row.id, + toCandidate: (row) => ({ + id: memoryVectorId(row.id), + text: buildMemoryEmbedTextFromRow(row), + namespace: userVectorNamespace(row.user_id), + metadata: { kind: 'memory', - candidates: rows.map((row) => ({ - id: memoryVectorId(row.id), - text: buildMemoryEmbedTextFromRow(row), - namespace: userVectorNamespace(row.user_id), - 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) + userId: row.user_id, + status: row.status, + ...(row.category ? { category: row.category } : {}), + }, + }), + }) } diff --git a/packages/worker/src/package-registry/package-reindex-d1-retry.node.test.ts b/packages/worker/src/package-registry/package-reindex-d1-retry.node.test.ts index dd095fe418..8d7cfbc1d9 100644 --- a/packages/worker/src/package-registry/package-reindex-d1-retry.node.test.ts +++ b/packages/worker/src/package-registry/package-reindex-d1-retry.node.test.ts @@ -147,7 +147,11 @@ test('saved package reindex retries transient D1 export errors then reports exha baseUrl: 'https://kody.example.com', }) await vi.advanceTimersByTimeAsync(d1LockRetryBaseDelayMs) - await expect(recoverPromise).resolves.toEqual({ upserted: 1 }) + await expect(recoverPromise).resolves.toEqual({ + upserted: 1, + complete: true, + afterId: null, + }) } finally { vi.useRealTimers() } @@ -177,6 +181,8 @@ test('saved package reindex retries transient D1 export errors then reports exha } await expect(exhaustedPromise).resolves.toEqual({ upserted: 0, + complete: true, + afterId: null, failed: 1, failures: [ { 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 939b002c07..236ef8e9cc 100644 --- a/packages/worker/src/package-registry/package-reindex.node.test.ts +++ b/packages/worker/src/package-registry/package-reindex.node.test.ts @@ -124,7 +124,7 @@ test('saved package reindex embeds full manifests with user-scoped metadata', as reindexSavedPackageVectors(env, { baseUrl: 'https://kody.example.com', }), - ).resolves.toEqual({ upserted: 1 }) + ).resolves.toEqual({ upserted: 1, complete: true, afterId: null }) expect(mockModule.loadPackageManifestBySourceId).toHaveBeenCalledWith({ env, @@ -217,6 +217,8 @@ test('saved package reindex skips failed manifest loads and continues the batch' }), ).resolves.toEqual({ upserted: 1, + complete: true, + afterId: null, failed: 1, failures: [ { @@ -296,7 +298,11 @@ test('saved package reindex retries a transient D1 export error on page listing' baseUrl: 'https://kody.example.com', }) await vi.advanceTimersByTimeAsync(d1LockRetryBaseDelayMs) - await expect(resultPromise).resolves.toEqual({ upserted: 1 }) + await expect(resultPromise).resolves.toEqual({ + upserted: 1, + complete: true, + afterId: null, + }) } finally { vi.useRealTimers() } @@ -366,7 +372,7 @@ test('saved package reindex walks keyset pages and merges the page results', asy reindexSavedPackageVectors({ APP_DB: {} } as Env, { baseUrl: 'https://kody.example.com', }), - ).resolves.toEqual({ upserted: 201 }) + ).resolves.toEqual({ upserted: 201, complete: true, afterId: null }) expect(mockModule.listSavedPackagesPage).toHaveBeenCalledTimes(2) expect(mockModule.listSavedPackagesPage).toHaveBeenNthCalledWith( @@ -385,3 +391,97 @@ test('saved package reindex walks keyset pages and merges the page results', asy expect(upsertedIds).toHaveLength(201) expect(new Set(upsertedIds).size).toBe(201) }) + +test('saved package reindex stops mid-page at the deadline and resumes', 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]), + ) + const firstPage = [ + buildSavedPackage('pkg-a'), + buildSavedPackage('pkg-b'), + buildSavedPackage('pkg-c'), + ] + mockModule.listSavedPackagesPage.mockResolvedValueOnce(firstPage) + + await expect( + reindexSavedPackageVectors({ APP_DB: {} } as Env, { + baseUrl: 'https://kody.example.com', + deadlineMs: 0, + }), + ).resolves.toEqual({ + upserted: 1, + complete: false, + afterId: 'pkg-a', + }) + expect(mockModule.loadPackageManifestBySourceId).toHaveBeenCalledTimes(1) + expect(upsert).toHaveBeenCalledTimes(1) + + mockModule.listSavedPackagesPage.mockResolvedValueOnce([ + buildSavedPackage('pkg-b'), + buildSavedPackage('pkg-c'), + ]) + await expect( + reindexSavedPackageVectors({ APP_DB: {} } as Env, { + baseUrl: 'https://kody.example.com', + afterId: 'pkg-a', + }), + ).resolves.toEqual({ + upserted: 2, + complete: true, + afterId: null, + }) + expect(mockModule.listSavedPackagesPage).toHaveBeenLastCalledWith( + expect.anything(), + { afterId: 'pkg-a', limit: 200 }, + ) +}) + +test('saved package reindex flushes upsert chunks and honors the deadline after a flush', 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') + const packages = Array.from({ length: 20 }, (_, index) => + buildSavedPackage(`pkg-${String(index).padStart(2, '0')}`), + ) + mockModule.listSavedPackagesPage.mockResolvedValueOnce(packages) + let now = 1_000 + const nowSpy = vi.spyOn(Date, 'now').mockImplementation(() => now) + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: unknown, texts: Array) => { + now = 3_000 + return texts.map(() => [0.1]) + }, + ) + + try { + await expect( + reindexSavedPackageVectors({ APP_DB: {} } as Env, { + baseUrl: 'https://kody.example.com', + deadlineMs: 2_000, + }), + ).resolves.toEqual({ + upserted: 16, + complete: false, + afterId: 'pkg-15', + }) + } finally { + nowSpy.mockRestore() + } + + expect(mockModule.loadPackageManifestBySourceId).toHaveBeenCalledTimes(16) + expect(upsert).toHaveBeenCalledTimes(1) + expect(upsert.mock.calls[0]?.[0]).toHaveLength(16) +}) diff --git a/packages/worker/src/package-registry/package-reindex.ts b/packages/worker/src/package-registry/package-reindex.ts index 54bb4f7eed..e249b9607c 100644 --- a/packages/worker/src/package-registry/package-reindex.ts +++ b/packages/worker/src/package-registry/package-reindex.ts @@ -5,10 +5,17 @@ import { import { mergeVectorReindexResults, reindexVectorCandidates, + vectorReindexUpsertBatchSize, type VectorReindexCandidate, type VectorReindexFailure, type VectorReindexResult, } from '#worker/vectorize/reindex-batches.ts' +import { + hasReachedReindexDeadline, + toVectorReindexSweepResult, + type VectorReindexSweepOptions, + type VectorReindexSweepResult, +} from '#worker/vectorize/reindex-sweep.ts' import { userVectorNamespace } from '#worker/vectorize/vector-namespaces.ts' import { runD1WithRetry } from '#worker/d1-retry.ts' import { getErrorMessage } from '@kody-internal/shared/error-message.ts' @@ -24,31 +31,111 @@ import { loadPackageManifestBySourceId } from './source.ts' // bounded regardless of table size (each row also loads its manifest). const reindexPageSize = 200 +type SavedPackageRow = Awaited>[number] + export async function reindexSavedPackageVectors( env: Env, - input: { baseUrl: string }, -): Promise { + input: { baseUrl: string } & VectorReindexSweepOptions, +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') } + const vectorIndex = index if (isCapabilitySearchOffline(env)) { - return { upserted: 0 } + return { upserted: 0, complete: true, afterId: null } } const pageResults: Array = [] - let afterId: string | null = null + let afterId: string | null = input.afterId ?? null + + async function flushChunk(chunk: { + rows: Array + candidates: Array + loadFailures: Array + }) { + if (chunk.loadFailures.length > 0) { + pageResults.push({ + upserted: 0, + failed: chunk.loadFailures.length, + failures: chunk.loadFailures, + failedIds: chunk.loadFailures.map((failure) => failure.id), + }) + } + if (chunk.candidates.length === 0) return + // Snapshot debt generations before upsert so a concurrent publish + // that bumps generation is not wiped by post-page cleanup. + const debtGenerations = new Map() + for (const row of chunk.rows) { + const generation = await getSavedPackageSearchIndexDebtGeneration({ + db: env.APP_DB, + packageId: row.id, + }) + if (generation !== null) { + debtGenerations.set(row.id, generation) + } + } + const pageResult = await reindexVectorCandidates({ + env, + index: vectorIndex, + kind: 'saved package', + candidates: chunk.candidates, + }) + pageResults.push(pageResult) + // Prefer uncapped failedIds (failures is a capped sample for messages). + const failedIds = new Set( + pageResult.failedIds ?? + (pageResult.failures ?? []).map((failure) => failure.id), + ) + // Self-heal deferred upsert debt observed before this chunk upsert. + for (const row of chunk.rows) { + const vectorId = savedPackageVectorId(row.id) + if (failedIds.has(vectorId)) continue + if (chunk.loadFailures.some((failure) => failure.id === vectorId)) { + continue + } + const generation = debtGenerations.get(row.id) + if (generation === undefined) continue + await clearSavedPackageSearchIndexDebt({ + db: env.APP_DB, + packageId: row.id, + generation, + }) + } + } + while (true) { + if (hasReachedReindexDeadline(input.deadlineMs) && pageResults.length > 0) { + return toVectorReindexSweepResult( + mergeVectorReindexResults('saved package', pageResults), + { complete: false, afterId }, + ) + } const rows = await runD1WithRetry(() => listSavedPackagesPage(env.APP_DB, { afterId, limit: reindexPageSize, }), ) - if (rows.length === 0) break + if (rows.length === 0) { + return toVectorReindexSweepResult( + mergeVectorReindexResults('saved package', pageResults), + { complete: true, afterId: null }, + ) + } const loadFailures: Array = [] const candidates: Array = [] + const processedRows: Array = [] + let stoppedEarly = false for (const row of rows) { + if ( + hasReachedReindexDeadline(input.deadlineMs) && + (processedRows.length > 0 || pageResults.length > 0) + ) { + stoppedEarly = true + break + } + processedRows.push(row) try { const { manifest } = await loadPackageManifestBySourceId({ env, @@ -81,56 +168,39 @@ export async function reindexSavedPackageVectors( }), ) } - } - if (loadFailures.length > 0) { - pageResults.push({ - upserted: 0, - failed: loadFailures.length, - failures: loadFailures, - failedIds: loadFailures.map((failure) => failure.id), - }) - } - if (candidates.length > 0) { - // Snapshot debt generations before upsert so a concurrent publish - // that bumps generation is not wiped by post-page cleanup. - const debtGenerations = new Map() - for (const row of rows) { - const generation = await getSavedPackageSearchIndexDebtGeneration({ - db: env.APP_DB, - packageId: row.id, + if (processedRows.length >= vectorReindexUpsertBatchSize) { + await flushChunk({ + rows: processedRows.splice(0), + candidates: candidates.splice(0), + loadFailures: loadFailures.splice(0), }) - if (generation !== null) { - debtGenerations.set(row.id, generation) + afterId = row.id + if (hasReachedReindexDeadline(input.deadlineMs)) { + stoppedEarly = true + break } } - const pageResult = await reindexVectorCandidates({ - env, - index, - kind: 'saved package', + } + if (processedRows.length > 0 || loadFailures.length > 0) { + const lastProcessed = processedRows[processedRows.length - 1] + await flushChunk({ + rows: processedRows, candidates, + loadFailures, }) - pageResults.push(pageResult) - // Prefer uncapped failedIds (failures is a capped sample for messages). - const failedIds = new Set( - pageResult.failedIds ?? - (pageResult.failures ?? []).map((failure) => failure.id), + if (lastProcessed) afterId = lastProcessed.id + } + if (stoppedEarly) { + return toVectorReindexSweepResult( + mergeVectorReindexResults('saved package', pageResults), + { complete: false, afterId }, + ) + } + if (rows.length < reindexPageSize) { + return toVectorReindexSweepResult( + mergeVectorReindexResults('saved package', pageResults), + { complete: true, afterId: null }, ) - // Self-heal deferred upsert debt observed before this page upsert. - for (const row of rows) { - const vectorId = savedPackageVectorId(row.id) - if (failedIds.has(vectorId)) continue - if (loadFailures.some((failure) => failure.id === vectorId)) continue - const generation = debtGenerations.get(row.id) - if (generation === undefined) continue - await clearSavedPackageSearchIndexDebt({ - db: env.APP_DB, - packageId: row.id, - generation, - }) - } } - if (rows.length < reindexPageSize) break - afterId = rows[rows.length - 1]!.id } - return mergeVectorReindexResults('saved package', pageResults) } diff --git a/packages/worker/src/vectorize/reindex-batches.ts b/packages/worker/src/vectorize/reindex-batches.ts index 93e89e0ab6..d2dc7db038 100644 --- a/packages/worker/src/vectorize/reindex-batches.ts +++ b/packages/worker/src/vectorize/reindex-batches.ts @@ -27,7 +27,7 @@ export type VectorReindexResult = { failedIds?: Array } -const upsertBatchSize = 16 +export const vectorReindexUpsertBatchSize = 16 const maxReportedFailures = 20 // Combines per-page reindex results from keyset-paged drivers into a single @@ -143,9 +143,11 @@ export async function reindexVectorCandidates(input: { for ( let offset = 0; offset < input.candidates.length; - offset += upsertBatchSize + offset += vectorReindexUpsertBatchSize ) { - await processBatch(input.candidates.slice(offset, offset + upsertBatchSize)) + await processBatch( + input.candidates.slice(offset, offset + vectorReindexUpsertBatchSize), + ) } if (failed === 0) return { upserted } diff --git a/packages/worker/src/vectorize/reindex-sweep.node.test.ts b/packages/worker/src/vectorize/reindex-sweep.node.test.ts new file mode 100644 index 0000000000..a2bbb7216f --- /dev/null +++ b/packages/worker/src/vectorize/reindex-sweep.node.test.ts @@ -0,0 +1,108 @@ +import { expect, test, vi } from 'vitest' +import { vectorReindexUpsertBatchSize } from './reindex-batches.ts' +import { + reindexPagedVectorRows, + reindexVectorCandidateList, +} from './reindex-sweep.ts' + +const mockModule = vi.hoisted(() => ({ + embedTextsForVectorize: vi.fn(), +})) + +vi.mock('./embedding.ts', () => ({ + embedTextsForVectorize: (...args: Array) => + mockModule.embedTextsForVectorize(...args), +})) + +function candidate(id: string) { + return { + id, + text: id, + namespace: 'ns', + metadata: { kind: 'test' }, + } +} + +test('reindex sweep helpers stop at the deadline and resume from afterId', async () => { + mockModule.embedTextsForVectorize.mockReset() + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: unknown, texts: Array) => texts.map(() => [0.1]), + ) + const upsert = vi.fn() + const index = { upsert } as unknown as VectorizeIndex + const env = {} as Env + const candidates = Array.from( + { length: vectorReindexUpsertBatchSize + 2 }, + (_, index_) => candidate(`cap-${String(index_).padStart(2, '0')}`), + ) + + const first = await reindexVectorCandidateList({ + env, + index, + kind: 'builtin capability', + candidates, + deadlineMs: 0, + }) + expect(first).toEqual({ + upserted: vectorReindexUpsertBatchSize, + complete: false, + afterId: `cap-${String(vectorReindexUpsertBatchSize - 1).padStart(2, '0')}`, + }) + + const second = await reindexVectorCandidateList({ + env, + index, + kind: 'builtin capability', + candidates, + afterId: first.afterId, + }) + expect(second).toEqual({ + upserted: 2, + complete: true, + afterId: null, + }) + + const rows = Array.from( + { length: vectorReindexUpsertBatchSize + 1 }, + (_, index_) => ({ id: `row-${String(index_).padStart(2, '0')}` }), + ) + const listed: Array = [] + const pagedFirst = await reindexPagedVectorRows({ + env, + index, + kind: 'memory', + pageSize: 200, + deadlineMs: 0, + listPage: async ({ afterId }) => { + listed.push(afterId) + return afterId == null ? rows : [] + }, + rowId: (row) => row.id, + toCandidate: (row) => candidate(row.id), + }) + expect(pagedFirst).toEqual({ + upserted: vectorReindexUpsertBatchSize, + complete: false, + afterId: `row-${String(vectorReindexUpsertBatchSize - 1).padStart(2, '0')}`, + }) + expect(listed).toEqual([null]) + + const pagedSecond = await reindexPagedVectorRows({ + env, + index, + kind: 'memory', + pageSize: 200, + afterId: pagedFirst.afterId, + listPage: async ({ afterId }) => { + if (afterId !== pagedFirst.afterId) return [] + return rows.slice(vectorReindexUpsertBatchSize) + }, + rowId: (row) => row.id, + toCandidate: (row) => candidate(row.id), + }) + expect(pagedSecond).toEqual({ + upserted: 1, + complete: true, + afterId: null, + }) +}) diff --git a/packages/worker/src/vectorize/reindex-sweep.ts b/packages/worker/src/vectorize/reindex-sweep.ts new file mode 100644 index 0000000000..f96805eb6f --- /dev/null +++ b/packages/worker/src/vectorize/reindex-sweep.ts @@ -0,0 +1,194 @@ +import { + mergeVectorReindexResults, + reindexVectorCandidates, + vectorReindexUpsertBatchSize, + type VectorReindexCandidate, + type VectorReindexResult, +} from './reindex-batches.ts' + +export const capabilityReindexPhases = [ + 'capabilities', + 'memories', + 'jobs', + 'packages', +] as const + +export type CapabilityReindexPhase = (typeof capabilityReindexPhases)[number] + +export type CapabilityReindexCursor = { + phase: CapabilityReindexPhase + afterId: string | null +} + +export type VectorReindexSweepOptions = { + afterId?: string | null + deadlineMs?: number +} + +export type VectorReindexSweepResult = VectorReindexResult & { + complete: boolean + afterId: string | null +} + +/** + * Wall-clock budget for one `/__maintenance/reindex-capabilities` request. + * Production CI curls with a 120s limit; this leaves headroom to return a + * cursor instead of hanging until the client times out. + */ +export const capabilityReindexTimeBudgetMs = 70_000 + +export function hasReachedReindexDeadline(deadlineMs?: number) { + return deadlineMs !== undefined && Date.now() >= deadlineMs +} + +export function isCapabilityReindexPhase( + value: unknown, +): value is CapabilityReindexPhase { + return ( + typeof value === 'string' && + (capabilityReindexPhases as ReadonlyArray).includes(value) + ) +} + +export function toVectorReindexSweepResult( + result: VectorReindexResult, + input: { complete: boolean; afterId: string | null }, +): VectorReindexSweepResult { + return { + ...result, + complete: input.complete, + afterId: input.afterId, + } +} + +export async function reindexVectorCandidateList(input: { + env: Env + index: VectorizeIndex + kind: string + candidates: ReadonlyArray + afterId?: string | null + deadlineMs?: number +}): Promise { + const startIndex = findResumeIndex(input.candidates, input.afterId) + const pageResults: Array = [] + let afterId = input.afterId ?? null + + for ( + let offset = startIndex; + offset < input.candidates.length; + offset += vectorReindexUpsertBatchSize + ) { + if (hasReachedReindexDeadline(input.deadlineMs) && pageResults.length > 0) { + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: false, afterId }, + ) + } + const batch = input.candidates.slice( + offset, + offset + vectorReindexUpsertBatchSize, + ) + pageResults.push( + await reindexVectorCandidates({ + env: input.env, + index: input.index, + kind: input.kind, + candidates: batch, + }), + ) + afterId = batch[batch.length - 1]!.id + } + + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: true, afterId: null }, + ) +} + +export async function reindexPagedVectorRows(input: { + env: Env + index: VectorizeIndex + kind: string + pageSize: number + afterId?: string | null + deadlineMs?: number + listPage: (input: { + afterId: string | null + limit: number + }) => Promise> + rowId: (row: Row) => string + toCandidate: (row: Row) => VectorReindexCandidate | null +}): Promise { + const pageResults: Array = [] + let afterId = input.afterId ?? null + + while (true) { + if (hasReachedReindexDeadline(input.deadlineMs) && pageResults.length > 0) { + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: false, afterId }, + ) + } + + const rows = await input.listPage({ + afterId, + limit: input.pageSize, + }) + if (rows.length === 0) { + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: true, afterId: null }, + ) + } + + for ( + let offset = 0; + offset < rows.length; + offset += vectorReindexUpsertBatchSize + ) { + if ( + hasReachedReindexDeadline(input.deadlineMs) && + pageResults.length > 0 + ) { + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: false, afterId }, + ) + } + const chunk = rows.slice(offset, offset + vectorReindexUpsertBatchSize) + const candidates = chunk.flatMap((row) => { + const candidate = input.toCandidate(row) + return candidate ? [candidate] : [] + }) + if (candidates.length > 0) { + pageResults.push( + await reindexVectorCandidates({ + env: input.env, + index: input.index, + kind: input.kind, + candidates, + }), + ) + } else { + pageResults.push({ upserted: 0 }) + } + afterId = input.rowId(chunk[chunk.length - 1]!) + } + + if (rows.length < input.pageSize) { + return toVectorReindexSweepResult( + mergeVectorReindexResults(input.kind, pageResults), + { complete: true, afterId: null }, + ) + } + } +} + +function findResumeIndex( + candidates: ReadonlyArray, + afterId: string | null | undefined, +) { + if (afterId == null) return 0 + const index = candidates.findIndex((candidate) => candidate.id === afterId) + return index === -1 ? 0 : index + 1 +}