diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 09c6d73490..02cb2c0134 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -268,8 +268,16 @@ jobs: fi base="${DEPLOY_URL%/}" echo "POST ${base}/__maintenance/reindex-capabilities" - curl --fail --silent --show-error --location --max-time 120 \ + status="$(curl --silent --show-error --location --max-time 120 \ -X POST \ -H "Authorization: Bearer ${CAPABILITY_REINDEX_SECRET}" \ -H "Accept: application/json" \ - "${base}/__maintenance/reindex-capabilities" + --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 diff --git a/packages/worker/src/capability-maintenance.node.test.ts b/packages/worker/src/capability-maintenance.node.test.ts index 87bee14e45..f477081bf6 100644 --- a/packages/worker/src/capability-maintenance.node.test.ts +++ b/packages/worker/src/capability-maintenance.node.test.ts @@ -123,6 +123,20 @@ test('capability reindex maintenance route attempts every vector kind before rep failed: 1, error: '1 saved package vector(s) failed to reindex', }, + failure: { + phase: 'reindex-capability-vectors', + failedPhases: [ + { + phase: 'memories', + cause: 'memory failed', + }, + { + phase: 'packages', + cause: '1 saved package vector(s) failed to reindex', + failed: 1, + }, + ], + }, error: 'Capability search vector reindex failed for memories, packages: memories: memory failed; packages: 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 aed9822337..497ce508c7 100644 --- a/packages/worker/src/capability-maintenance.ts +++ b/packages/worker/src/capability-maintenance.ts @@ -13,10 +13,23 @@ type ReindexStepResult = { upserted: number error?: string failed?: number + warning?: string + failures?: Array<{ + id: string + phase: string + error: string + }> +} + +type ReindexFailureSummary = { + phase: string + cause: string + failed?: number + failures?: ReindexStepResult['failures'] } async function runReindexStep( - run: () => Promise<{ upserted: number }>, + run: () => Promise, ): Promise { try { return await run() @@ -65,13 +78,27 @@ export async function handleCapabilityReindexRequest( typeof entry[1].error === 'string', ) if (failed.length > 0) { + const failedPhases: Array = failed.map( + ([phase, step]) => ({ + phase, + cause: step.error, + ...(typeof step.failed === 'number' ? { failed: step.failed } : {}), + ...(step.failures ? { failures: step.failures } : {}), + }), + ) throw new MaintenanceFailureError( `Capability search vector reindex failed for ${failed .map(([kind]) => kind) .join(', ')}: ${failed .map(([kind, step]) => `${kind}: ${step.error}`) .join('; ')}`, - result, + { + ...result, + failure: { + phase: 'reindex-capability-vectors', + failedPhases, + }, + }, ) } return result diff --git a/packages/worker/src/jobs/job-reindex.ts b/packages/worker/src/jobs/job-reindex.ts index da1be82e15..c02ac3c96a 100644 --- a/packages/worker/src/jobs/job-reindex.ts +++ b/packages/worker/src/jobs/job-reindex.ts @@ -1,17 +1,18 @@ import { - embedTextsForVectorize, getCapabilityVectorIndex, isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' +import { + 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' -const upsertBatchSize = 16 - export async function reindexJobVectors( env: Env, -): Promise<{ upserted: number }> { +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') @@ -38,27 +39,19 @@ export async function reindexJobVectors( ).flat() if (jobs.length === 0) return { upserted: 0 } - let upserted = 0 - for (let offset = 0; offset < jobs.length; offset += upsertBatchSize) { - const batch = jobs.slice(offset, offset + upsertBatchSize) - const texts = batch.map(({ job }) => - buildJobEmbedText({ + 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, }), - ) - const vectors = await embedTextsForVectorize(env, texts) - await index.upsert( - batch.map(({ userId, job }, index_) => ({ - id: jobVectorId(job.id), - values: vectors[index_]!, - metadata: { kind: 'job', userId }, - })), - ) - upserted += batch.length - } - - return { upserted } + metadata: { kind: 'job', userId }, + })), + }) } diff --git a/packages/worker/src/mcp/capabilities/capability-reindex.ts b/packages/worker/src/mcp/capabilities/capability-reindex.ts index 09b6c3a212..110bb6adaf 100644 --- a/packages/worker/src/mcp/capabilities/capability-reindex.ts +++ b/packages/worker/src/mcp/capabilities/capability-reindex.ts @@ -1,38 +1,30 @@ import { buildCapabilityEmbedText, - embedTextsForVectorize, getCapabilityVectorIndex, } from './capability-search.ts' +import { + reindexVectorCandidates, + type VectorReindexResult, +} from './reindex-batches.ts' import { type CapabilitySpec } from './types.ts' -const upsertBatchSize = 16 - export async function reindexCapabilityVectors( env: Env, specs: Record, -): Promise<{ upserted: number }> { +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') } - const list = Object.values(specs) - let upserted = 0 - - for (let offset = 0; offset < list.length; offset += upsertBatchSize) { - const batch = list.slice(offset, offset + upsertBatchSize) - const texts = batch.map((spec) => buildCapabilityEmbedText(spec)) - const rows = await embedTextsForVectorize(env, texts) - - const vectors = batch.map((spec, index_) => ({ + return reindexVectorCandidates({ + env, + index, + kind: 'builtin capability', + candidates: Object.values(specs).map((spec) => ({ id: spec.name, - values: rows[index_]!, + text: buildCapabilityEmbedText(spec), metadata: { domain: spec.domain, kind: 'builtin' }, - })) - - await index.upsert(vectors) - upserted += vectors.length - } - - return { upserted } + })), + }) } diff --git a/packages/worker/src/mcp/capabilities/capability-search.ts b/packages/worker/src/mcp/capabilities/capability-search.ts index 19d648c79c..73cc3be254 100644 --- a/packages/worker/src/mcp/capabilities/capability-search.ts +++ b/packages/worker/src/mcp/capabilities/capability-search.ts @@ -1,4 +1,5 @@ import { type CapabilitySpec } from './types.ts' +import { getErrorMessage } from './error-message.ts' type CapabilityVectorizeEnv = { AI?: Ai @@ -20,9 +21,16 @@ export function getCapabilityVectorIndex(env: Env): VectorizeIndex | undefined { /** Must match Vectorize index dimensions for production indexes. */ export const CAPABILITY_EMBEDDING_DIMENSIONS = 384 export const CAPABILITY_EMBEDDING_MODEL = '@cf/baai/bge-small-en-v1.5' +export const CAPABILITY_EMBEDDING_BATCH_SIZE = 8 +export const CAPABILITY_EMBEDDING_MAX_INPUT_CHARS = 2_000 export const CAPABILITY_SEARCH_RRF_K = 60 +function truncateEmbeddingInput(text: string) { + if (text.length <= CAPABILITY_EMBEDDING_MAX_INPUT_CHARS) return text + return text.slice(0, CAPABILITY_EMBEDDING_MAX_INPUT_CHARS) +} + function fnv1a32(input: string): number { let hash = 2_166_136_261 for (let i = 0; i < input.length; i += 1) { @@ -206,26 +214,40 @@ export async function embedTextsForVectorize( texts: ReadonlyArray, ): Promise>> { if (texts.length === 0) return [] + const truncatedTexts = texts.map((text) => truncateEmbeddingInput(text)) const runtime = env as unknown as CapabilityVectorizeEnv if (!runtime.AI) { if (runtime.SENTRY_ENVIRONMENT !== 'production') { - return texts.map((text) => deterministicEmbedding(text)) + return truncatedTexts.map((text) => deterministicEmbedding(text)) } throw new Error( 'AI binding is required for capability embeddings in production.', ) } + const aiRuntime = runtime as CapabilityVectorizeEnv & { AI: Ai } + + const rows: Array> = [] + for ( + let offset = 0; + offset < truncatedTexts.length; + offset += CAPABILITY_EMBEDDING_BATCH_SIZE + ) { + const batch = truncatedTexts.slice( + offset, + offset + CAPABILITY_EMBEDDING_BATCH_SIZE, + ) + rows.push(...(await embedTextBatchForVectorize(aiRuntime, batch))) + } + return rows +} - const options = - runtime.AI_GATEWAY_ID && runtime.AI_GATEWAY_ID.trim().length > 0 - ? { - gateway: { - id: runtime.AI_GATEWAY_ID.trim(), - }, - } - : undefined - const response = (await runtime.AI.run( +async function runWorkersAiEmbeddingRequest( + runtime: CapabilityVectorizeEnv & { AI: Ai }, + texts: ReadonlyArray, + options?: { gateway: { id: string } }, +) { + return (await runtime.AI.run( CAPABILITY_EMBEDDING_MODEL, { text: [...texts], @@ -233,7 +255,40 @@ export async function embedTextsForVectorize( }, options, )) as WorkersAiEmbeddingResponse +} + +async function embedTextBatchForVectorize( + runtime: CapabilityVectorizeEnv & { AI: Ai }, + texts: ReadonlyArray, +) { + const gatewayId = runtime.AI_GATEWAY_ID?.trim() + if (gatewayId) { + try { + const response = await runWorkersAiEmbeddingRequest(runtime, texts, { + gateway: { id: gatewayId }, + }) + return parseWorkersAiEmbeddingResponse(response, texts.length) + } catch (error) { + console.warn( + JSON.stringify({ + message: + 'Workers AI Gateway embedding request failed; retrying direct Workers AI', + gatewayId, + error: getErrorMessage(error), + }), + ) + try { + const response = await runWorkersAiEmbeddingRequest(runtime, texts) + return parseWorkersAiEmbeddingResponse(response, texts.length) + } catch (directError) { + throw new Error( + `Workers AI embedding request failed after AI Gateway fallback. Gateway error: ${getErrorMessage(error)}. Direct error: ${getErrorMessage(directError)}`, + ) + } + } + } + const response = await runWorkersAiEmbeddingRequest(runtime, texts) return parseWorkersAiEmbeddingResponse(response, texts.length) } diff --git a/packages/worker/src/mcp/capabilities/capability-search.workers.test.ts b/packages/worker/src/mcp/capabilities/capability-search.workers.test.ts index 115bce5ca9..fb720d1ad3 100644 --- a/packages/worker/src/mcp/capabilities/capability-search.workers.test.ts +++ b/packages/worker/src/mcp/capabilities/capability-search.workers.test.ts @@ -1,6 +1,8 @@ -import { expect, test } from 'vitest' +import { expect, test, vi } from 'vitest' import { + CAPABILITY_EMBEDDING_BATCH_SIZE, CAPABILITY_EMBEDDING_DIMENSIONS, + CAPABILITY_EMBEDDING_MAX_INPUT_CHARS, CAPABILITY_EMBEDDING_MODEL, deterministicEmbedding, embedTextsForVectorize, @@ -117,6 +119,85 @@ test('embedding wrapper batches texts through Workers AI and AI Gateway', async ]) }) +test('embedding wrapper falls back to direct Workers AI when AI Gateway fails', async () => { + const consoleWarn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const calls: Array> = [] + const rows = [embeddingRow(1), embeddingRow(2)] + const env = aiEnv( + async (...args) => { + calls.push(args) + if (args[2]) throw new Error('gateway not found') + return { data: rows, shape: [2, CAPABILITY_EMBEDDING_DIMENSIONS] } + }, + { AI_GATEWAY_ID: 'stale-gateway' }, + ) + + try { + await expect( + embedTextsForVectorize(env, ['alpha', 'beta']), + ).resolves.toEqual(rows) + expect(calls).toEqual([ + [ + CAPABILITY_EMBEDDING_MODEL, + { text: ['alpha', 'beta'], pooling: 'cls' }, + { gateway: { id: 'stale-gateway' } }, + ], + [ + CAPABILITY_EMBEDDING_MODEL, + { text: ['alpha', 'beta'], pooling: 'cls' }, + undefined, + ], + ]) + expect(consoleWarn).toHaveBeenCalledWith( + expect.stringContaining('retrying direct Workers AI'), + ) + } finally { + consoleWarn.mockRestore() + } +}) + +test('embedding wrapper chunks large batches and truncates long inputs', async () => { + const calls: Array> = [] + const inputCount = CAPABILITY_EMBEDDING_BATCH_SIZE + 1 + const env = aiEnv(async (...args) => { + calls.push(args) + const payload = args[1] as { text: Array } + return { + data: payload.text.map((_, index) => + embeddingRow(calls.length * 10 + index), + ), + shape: [payload.text.length, CAPABILITY_EMBEDDING_DIMENSIONS], + } + }) + + await expect( + embedTextsForVectorize(env, [ + 'x'.repeat(CAPABILITY_EMBEDDING_MAX_INPUT_CHARS + 500), + ...Array.from({ length: inputCount - 1 }, (_, index) => `text-${index}`), + ]), + ).resolves.toHaveLength(inputCount) + + expect(calls).toHaveLength(2) + expect( + calls.map((args) => (args[1] as { text: Array }).text.length), + ).toEqual([CAPABILITY_EMBEDDING_BATCH_SIZE, 1]) + expect( + (calls[0]![1] as { text: Array }).text.every( + (text) => text.length <= CAPABILITY_EMBEDDING_MAX_INPUT_CHARS, + ), + ).toBe(true) +}) + +test('embedding wrapper reports direct Workers AI failures', async () => { + const env = aiEnv(async () => { + throw new Error('workers ai unavailable') + }) + + await expect(embedTextsForVectorize(env, ['alpha'])).rejects.toThrow( + /workers ai unavailable/, + ) +}) + test('embedding wrapper falls back deterministically outside production when AI is unavailable', async () => { const env = { SENTRY_ENVIRONMENT: 'preview' } as Env diff --git a/packages/worker/src/mcp/capabilities/reindex-batches.node.test.ts b/packages/worker/src/mcp/capabilities/reindex-batches.node.test.ts new file mode 100644 index 0000000000..add77be1df --- /dev/null +++ b/packages/worker/src/mcp/capabilities/reindex-batches.node.test.ts @@ -0,0 +1,98 @@ +import { expect, test, vi } from 'vitest' + +const mockModule = vi.hoisted(() => ({ + embedTextsForVectorize: vi.fn(), +})) + +vi.mock('./capability-search.ts', () => ({ + embedTextsForVectorize: (...args: Array) => + mockModule.embedTextsForVectorize(...args), +})) + +const { reindexVectorCandidates } = await import('./reindex-batches.ts') + +test('reindexVectorCandidates isolates failed embedding items and upserts the rest', async () => { + mockModule.embedTextsForVectorize.mockReset() + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const upsert = vi.fn() + const env = {} as Env + mockModule.embedTextsForVectorize.mockImplementation( + async (_env: Env, texts: Array) => { + if (texts.includes('bad')) throw new Error('embedding input rejected') + return texts.map((text) => [text.length]) + }, + ) + + try { + await expect( + reindexVectorCandidates({ + env, + index: { upsert } as unknown as VectorizeIndex, + kind: 'test', + candidates: [ + { id: 'good-1', text: 'alpha', metadata: { kind: 'test' } }, + { id: 'bad-1', text: 'bad', metadata: { kind: 'test' } }, + { id: 'good-2', text: 'beta', metadata: { kind: 'test' } }, + ], + }), + ).resolves.toEqual({ + upserted: 2, + failed: 1, + failures: [ + { + id: 'bad-1', + phase: 'embed', + error: 'embedding input rejected', + }, + ], + warning: '1 test vector(s) failed to reindex', + }) + + expect(upsert).toHaveBeenCalledWith([ + { + id: 'good-1', + values: [5], + metadata: { kind: 'test' }, + }, + ]) + expect(upsert).toHaveBeenCalledWith([ + { + id: 'good-2', + values: [4], + metadata: { kind: 'test' }, + }, + ]) + expect(consoleError).toHaveBeenCalledWith( + expect.stringContaining('capability vector reindex skipped vector'), + ) + } finally { + consoleError.mockRestore() + } +}) + +test('reindexVectorCandidates marks a phase fatal when every vector fails', async () => { + mockModule.embedTextsForVectorize.mockReset() + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const env = {} as Env + mockModule.embedTextsForVectorize.mockRejectedValue(new Error('ai down')) + + try { + await expect( + reindexVectorCandidates({ + env, + index: { upsert: vi.fn() } as unknown as VectorizeIndex, + kind: 'test', + candidates: [ + { id: 'bad-1', text: 'alpha', metadata: { kind: 'test' } }, + { id: 'bad-2', text: 'beta', metadata: { kind: 'test' } }, + ], + }), + ).resolves.toMatchObject({ + upserted: 0, + failed: 2, + error: '2 test vector(s) failed to reindex', + }) + } finally { + consoleError.mockRestore() + } +}) diff --git a/packages/worker/src/mcp/capabilities/reindex-batches.ts b/packages/worker/src/mcp/capabilities/reindex-batches.ts new file mode 100644 index 0000000000..db7e5a7813 --- /dev/null +++ b/packages/worker/src/mcp/capabilities/reindex-batches.ts @@ -0,0 +1,124 @@ +import { getErrorMessage } from './error-message.ts' +import { embedTextsForVectorize } from './capability-search.ts' + +export type VectorReindexFailurePhase = 'load' | 'embed' | 'upsert' + +export type VectorReindexFailure = { + id: string + phase: VectorReindexFailurePhase + error: string +} + +export type VectorReindexCandidate = { + id: string + text: string + metadata: Record +} + +export type VectorReindexResult = { + upserted: number + failed?: number + warning?: string + error?: string + failures?: Array +} + +const upsertBatchSize = 16 +const maxReportedFailures = 20 + +export async function reindexVectorCandidates(input: { + env: Env + index: VectorizeIndex + kind: string + candidates: ReadonlyArray +}): Promise { + let upserted = 0 + let failed = 0 + const failures: Array = [] + + function recordFailure(input_: { + candidate: VectorReindexCandidate + phase: VectorReindexFailurePhase + error: unknown + }) { + failed += 1 + const failure = { + id: input_.candidate.id, + phase: input_.phase, + error: getErrorMessage(input_.error), + } + if (failures.length < maxReportedFailures) failures.push(failure) + console.error( + JSON.stringify({ + message: 'capability vector reindex skipped vector', + kind: input.kind, + vectorId: input_.candidate.id, + phase: input_.phase, + error: failure.error, + }), + ) + } + + async function splitOrRecord( + batch: ReadonlyArray, + phase: VectorReindexFailurePhase, + error: unknown, + ) { + if (batch.length === 1) { + recordFailure({ candidate: batch[0]!, phase, error }) + return + } + const midpoint = Math.ceil(batch.length / 2) + await processBatch(batch.slice(0, midpoint)) + await processBatch(batch.slice(midpoint)) + } + + async function processBatch(batch: ReadonlyArray) { + let embeddings: Array> + try { + embeddings = await embedTextsForVectorize( + input.env, + batch.map((candidate) => candidate.text), + ) + } catch (error) { + await splitOrRecord(batch, 'embed', error) + return + } + + try { + await input.index.upsert( + batch.map((candidate, index_) => ({ + id: candidate.id, + values: embeddings[index_]!, + metadata: candidate.metadata, + })), + ) + upserted += batch.length + } catch (error) { + await splitOrRecord(batch, 'upsert', error) + } + } + + for ( + let offset = 0; + offset < input.candidates.length; + offset += upsertBatchSize + ) { + await processBatch(input.candidates.slice(offset, offset + upsertBatchSize)) + } + + if (failed === 0) return { upserted } + + const result: VectorReindexResult = { + upserted, + failed, + failures, + } + const message = `${failed} ${input.kind} vector(s) failed to reindex` + if (upserted === 0 && input.candidates.length > 0) { + result.error = message + } else { + result.warning = message + } + return result +} diff --git a/packages/worker/src/mcp/memory/memory-reindex.ts b/packages/worker/src/mcp/memory/memory-reindex.ts index f5c7977f57..657d2217b9 100644 --- a/packages/worker/src/mcp/memory/memory-reindex.ts +++ b/packages/worker/src/mcp/memory/memory-reindex.ts @@ -1,17 +1,18 @@ import { - embedTextsForVectorize, getCapabilityVectorIndex, isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' +import { + reindexVectorCandidates, + type VectorReindexResult, +} from '#mcp/capabilities/reindex-batches.ts' import { buildMemoryEmbedTextFromRow } from './memory-embed.ts' import { listAllMemories } from './repo.ts' import { memoryVectorId } from './memory-vectorize.ts' -const upsertBatchSize = 16 - -export async function reindexMemoryVectors(env: Env): Promise<{ - upserted: number -}> { +export async function reindexMemoryVectors( + env: Env, +): Promise { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') @@ -27,25 +28,19 @@ export async function reindexMemoryVectors(env: Env): Promise<{ return { upserted: 0 } } - let upserted = 0 - for (let offset = 0; offset < rows.length; offset += upsertBatchSize) { - const batch = rows.slice(offset, offset + upsertBatchSize) - const texts = batch.map((row) => buildMemoryEmbedTextFromRow(row)) - const vectors = await embedTextsForVectorize(env, texts) - await index.upsert( - batch.map((row, index_) => ({ - id: memoryVectorId(row.id), - values: vectors[index_]!, - metadata: { - kind: 'memory', - userId: row.user_id, - status: row.status, - ...(row.category ? { category: row.category } : {}), - }, - })), - ) - upserted += batch.length - } - - return { upserted } + return reindexVectorCandidates({ + env, + index, + kind: 'memory', + 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 } : {}), + }, + })), + }) } 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 d60aea0b3c..e7e50adfc1 100644 --- a/packages/worker/src/package-registry/package-reindex.node.test.ts +++ b/packages/worker/src/package-registry/package-reindex.node.test.ts @@ -176,7 +176,14 @@ test('saved package reindex skips failed manifest loads and continues the batch' ).resolves.toEqual({ upserted: 1, failed: 1, - error: '1 saved package vector(s) failed to reindex', + failures: [ + { + id: 'package_pkg-bad', + phase: 'load', + error: 'manifest missing', + }, + ], + warning: '1 saved package vector(s) failed to reindex', }) expect(mockModule.embedTextsForVectorize).toHaveBeenCalledWith(env, [ diff --git a/packages/worker/src/package-registry/package-reindex.ts b/packages/worker/src/package-registry/package-reindex.ts index 137f12b17b..3f7dc12746 100644 --- a/packages/worker/src/package-registry/package-reindex.ts +++ b/packages/worker/src/package-registry/package-reindex.ts @@ -1,19 +1,27 @@ import { - embedTextsForVectorize, getCapabilityVectorIndex, isCapabilitySearchOffline, } from '#mcp/capabilities/capability-search.ts' +import { + reindexVectorCandidates, + type VectorReindexCandidate, + type VectorReindexFailure, +} 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 { loadPackageManifestBySourceId } from './source.ts' -const upsertBatchSize = 16 - export async function reindexSavedPackageVectors( env: Env, input: { baseUrl: string }, -): Promise<{ upserted: number; failed?: number; error?: string }> { +): Promise<{ + upserted: number + failed?: number + warning?: string + error?: string + failures?: Array +}> { const index = getCapabilityVectorIndex(env) if (!index) { throw new Error('CAPABILITY_VECTOR_INDEX binding is not configured') @@ -27,58 +35,59 @@ export async function reindexSavedPackageVectors( return { upserted: 0 } } - let upserted = 0 let failed = 0 - for (let offset = 0; offset < rows.length; offset += upsertBatchSize) { - const batch = rows.slice(offset, offset + upsertBatchSize) - const loaded: Array<{ row: (typeof batch)[number]; embedText: string }> = [] - for (const row of batch) { - try { - const { manifest } = await loadPackageManifestBySourceId({ - env, - baseUrl: input.baseUrl, - userId: row.userId, - sourceId: row.sourceId, - }) - loaded.push({ - row, - embedText: buildSavedPackageEmbedText(manifest), - }) - } catch (error) { - failed += 1 - console.error( - JSON.stringify({ - message: 'saved package vector reindex skipped package', - packageId: row.id, - sourceId: row.sourceId, - error: getErrorMessage(error), - }), - ) - } - } - if (loaded.length === 0) continue - const vectors = await embedTextsForVectorize( - env, - loaded.map(({ embedText }) => embedText), - ) - await index.upsert( - loaded.map(({ row }, index_) => ({ + 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), - values: vectors[index_]!, + text: buildSavedPackageEmbedText(manifest), metadata: { kind: 'package', userId: row.userId, }, - })), - ) - upserted += loaded.length + }) + } 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, + }), + ) + } } - return failed > 0 - ? { - upserted, - failed, - error: `${failed} saved package vector(s) failed to reindex`, - } - : { upserted } + 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 }), + } }