Fix production capability vector reindex resilience

.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

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',
 	})

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<ReindexStepResult>,
 ): Promise<ReindexStepResult> {
 	try {
 		return await run()
@@ -65,13 +78,27 @@ export async function handleCapabilityReindexRequest(
 					typeof entry[1].error === 'string',
 			)
 			if (failed.length > 0) {
+				const failedPhases: Array<ReindexFailureSummary> = 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

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<VectorReindexResult> {
 	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 },
+		})),
+	})
 }

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<string, CapabilitySpec>,
-): Promise<{ upserted: number }> {
+): Promise<VectorReindexResult> {
 	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 }
+		})),
+	})
 }

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,34 +214,81 @@ export async function embedTextsForVectorize(
 	texts: ReadonlyArray<string>,
 ): Promise<Array<Array<number>>> {
 	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<Array<number>> = []
+	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<string>,
+	options?: { gateway: { id: string } },
+) {
+	return (await runtime.AI.run(
 		CAPABILITY_EMBEDDING_MODEL,
 		{
 			text: [...texts],
 			pooling: 'cls',
 		},
 		options,
 	)) as WorkersAiEmbeddingResponse
+}
+
+async function embedTextBatchForVectorize(
+	runtime: CapabilityVectorizeEnv & { AI: Ai },
+	texts: ReadonlyArray<string>,
+) {
+	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)
 }