diff --git a/src/features/workspaces/search/workspace-search-embeddings.test.ts b/src/features/workspaces/search/workspace-search-embeddings.test.ts new file mode 100644 index 00000000..8d7296c2 --- /dev/null +++ b/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); + }); +}); diff --git a/src/features/workspaces/search/workspace-search-embeddings.ts b/src/features/workspaces/search/workspace-search-embeddings.ts index 5790a609..eaf9d479 100644 --- a/src/features/workspaces/search/workspace-search-embeddings.ts +++ b/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; +} + +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 { + 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 { + 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."); diff --git a/src/features/workspaces/search/workspace-search-schema.ts b/src/features/workspaces/search/workspace-search-schema.ts index 2bdd6325..711ac14b 100644 --- a/src/features/workspaces/search/workspace-search-schema.ts +++ b/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)`;