fix(workspace-search): repair wedged search indexing

src/features/workspaces/search/workspace-search-embeddings.test.ts

@@ -0,0 +1,66 @@
+import { describe, expect, it, vi } from "vitest";
+
+import {
+	embedWorkspaceSearchTexts,
+	isRetryableEmbeddingError,
+} from "#/features/workspaces/search/workspace-search-embeddings";
+
+function aiError(message: string): Error {
+	const error = new Error(message);
+	error.name = "AiError";
+	return error;
+}
+
+function embeddingResponse(count: number) {
+	return { data: Array.from({ length: count }, () => [0.1, 0.2, 0.3]) };
+}
+
+describe("workspace search embeddings", () => {
+	it("classifies transient Workers AI errors as retryable", () => {
+		expect(isRetryableEmbeddingError(aiError("3043: Internal server error"))).toBe(true);
+		expect(isRetryableEmbeddingError(aiError("5006: Service temporarily unavailable"))).toBe(true);
+		expect(isRetryableEmbeddingError(aiError("503: upstream unavailable"))).toBe(true);
+	});
+
+	it("does not retry permanent client-side errors", () => {
+		expect(isRetryableEmbeddingError(aiError("2001: invalid input"))).toBe(false);
+		expect(isRetryableEmbeddingError(aiError("400: bad request"))).toBe(false);
+		expect(isRetryableEmbeddingError(new Error("some other failure"))).toBe(false);
+		expect(isRetryableEmbeddingError("not an error")).toBe(false);
+	});
+
+	it("retries a transient failure and returns the eventual embeddings", async () => {
+		const run = vi
+			.fn()
+			.mockRejectedValueOnce(aiError("3043: Internal server error"))
+			.mockResolvedValueOnce(embeddingResponse(2));
+		const ai = { run } as unknown as Ai;
+
+		const embeddings = await embedWorkspaceSearchTexts(ai, ["one", "two"], {
+			sleep: () => Promise.resolve(),
+		});
+
+		expect(run).toHaveBeenCalledTimes(2);
+		expect(embeddings).toHaveLength(2);
+	});
+
+	it("gives up after the attempt budget and rethrows the transient error", async () => {
+		const run = vi.fn().mockRejectedValue(aiError("3043: Internal server error"));
+		const ai = { run } as unknown as Ai;
+
+		await expect(
+			embedWorkspaceSearchTexts(ai, ["one"], { maxAttempts: 3, sleep: () => Promise.resolve() }),
+		).rejects.toThrow("3043");
+		expect(run).toHaveBeenCalledTimes(3);
+	});
+
+	it("fails fast on permanent errors without retrying", async () => {
+		const run = vi.fn().mockRejectedValue(aiError("400: bad request"));
+		const ai = { run } as unknown as Ai;
+
+		await expect(
+			embedWorkspaceSearchTexts(ai, ["one"], { sleep: () => Promise.resolve() }),
+		).rejects.toThrow("400");
+		expect(run).toHaveBeenCalledTimes(1);
+	});
+});

src/features/workspaces/search/workspace-search-embeddings.ts

@@ -2,19 +2,95 @@ import { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspa
 
 const workspaceSearchEmbeddingModel = "@cf/baai/bge-m3";
 const embeddingBatchSize = 16;
+const maximumEmbeddingAttempts = 4;
+const embeddingRetryBaseDelayMs = 250;
+const embeddingRetryMaxDelayMs = 2_000;
 
-export async function embedWorkspaceSearchTexts(ai: Ai, texts: string[]) {
+export interface WorkspaceSearchEmbeddingRetryOptions {
+	maxAttempts?: number;
+	sleep?: (ms: number) => Promise<void>;
+}
+
+export async function embedWorkspaceSearchTexts(
+	ai: Ai,
+	texts: string[],
+	options?: WorkspaceSearchEmbeddingRetryOptions,
+) {
 	const embeddings: number[][] = [];
 	for (const batch of batchWorkspaceSearchValues(texts, embeddingBatchSize)) {
-		const output: unknown = await ai.run(workspaceSearchEmbeddingModel, {
-			text: batch,
-			truncate_inputs: true,
-		});
+		const output = await runEmbeddingWithRetry(ai, batch, options);
 		embeddings.push(...readEmbeddingData(output));
 	}
 	return embeddings;
 }
 
+async function runEmbeddingWithRetry(
+	ai: Ai,
+	batch: string[],
+	options?: WorkspaceSearchEmbeddingRetryOptions,
+): Promise<unknown> {
+	const maxAttempts = options?.maxAttempts ?? maximumEmbeddingAttempts;
+	const sleep = options?.sleep ?? defaultEmbeddingRetrySleep;
+
+	for (let attempt = 1; ; attempt += 1) {
+		try {
+			return await ai.run(workspaceSearchEmbeddingModel, {
+				text: batch,
+				truncate_inputs: true,
+			});
+		} catch (error) {
+			// A transient Workers AI hiccup should not consume an index attempt or
+			// page anyone: retry the retryable ones in-process before bubbling up.
+			if (attempt >= maxAttempts || !isRetryableEmbeddingError(error)) {
+				throw error;
+			}
+			await sleep(embeddingRetryDelayMs(attempt));
+		}
+	}
+}
+
+/**
+ * Workers AI surfaces transient upstream failures as an `AiError` whose message
+ * leads with a numeric code, e.g. `3043: Internal server error`. Only the
+ * server-side classes — 3xxx internal errors and 5xx-class upstream failures —
+ * are worth retrying; client/validation codes are permanent and must bubble up.
+ */
+export function isRetryableEmbeddingError(error: unknown): boolean {
+	if (!(error instanceof Error) || !isAiErrorName(error.name)) {
+		return false;
+	}
+	const code = readAiErrorCode(error.message);
+	return code !== null && isRetryableAiErrorCode(code);
+}
+
+function isAiErrorName(name: string): boolean {
+	return name === "AiError" || name === "InferenceUpstreamError" || name === "AiInternalError";
+}
+
+function readAiErrorCode(message: string): number | null {
+	const match = /^\s*(\d{3,4})\b/.exec(message);
+	if (!match?.[1]) {
+		return null;
+	}
+	return Number.parseInt(match[1], 10);
+}
+
+function isRetryableAiErrorCode(code: number): boolean {
+	return (
+		(code >= 3000 && code <= 3999) || (code >= 5000 && code <= 5999) || (code >= 500 && code <= 599)
+	);
+}
+
+function embeddingRetryDelayMs(attempt: number): number {
+	return Math.min(embeddingRetryBaseDelayMs * 2 ** (attempt - 1), embeddingRetryMaxDelayMs);
+}
+
+function defaultEmbeddingRetrySleep(ms: number): Promise<void> {
+	return new Promise((resolve) => {
+		setTimeout(resolve, ms);
+	});
+}
+
 function readEmbeddingData(output: unknown): number[][] {
 	if (!isRecord(output) || !Array.isArray(output.data)) {
 		throw new Error("Workspace search embedding response is missing vector data.");

src/features/workspaces/search/workspace-search-schema.ts

@@ -11,6 +11,16 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) {
 			attempts INTEGER NOT NULL DEFAULT 0
 		)
 	`;
+	// The delete queue survives version resets (its rows reference vectors still
+	// live in Vectorize), so it is never rebuilt by the reset below. Objects whose
+	// table predates the `attempts` column must gain it here, or every flush throws
+	// "no such column: attempts" and indexing stalls forever.
+	const hasAttemptsColumn = sql<{ name: string }>`
+		PRAGMA table_info(kernel_search_vector_deletes)
+	`.some((column) => column.name === "attempts");
+	if (!hasAttemptsColumn) {
+		sql`ALTER TABLE kernel_search_vector_deletes ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0`;
+	}
 	sql`CREATE INDEX IF NOT EXISTS kernel_search_vector_deletes_pending_idx
 		ON kernel_search_vector_deletes (requested_at)`;