diff --git a/src/services/embedding.ts b/src/services/embedding.ts index efadf9ee..e9065a46 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -13,6 +13,18 @@ const requireFromHere = createRuntimeRequire(import.meta); const TIMEOUT_MS = 30000; const GLOBAL_EMBEDDING_KEY = Symbol.for("opencode-mem.embedding.instance"); const MAX_CACHE_SIZE = 100; +/** + * Per-attempt ceiling for one remote embedding HTTP call. Kept at 8s so + * 3 attempts + linear backoff stay under the outer embedWithTimeout budget + * (TIMEOUT_MS = 30s). A single hung request costs one attempt instead of the + * whole budget (#400: one transient hang previously aborted an entire import + * with zero memories written). + */ +const REMOTE_ATTEMPT_TIMEOUT_MS = 8000; +/** Attempts per remote embedding call: 1 try + 2 retries on transient failures. */ +const REMOTE_MAX_ATTEMPTS = 3; +/** Base backoff between remote embedding attempts; scales linearly (backoffMs * attempt). */ +const REMOTE_RETRY_BACKOFF_MS = 500; // ONNX weight variants in descending precision, mirroring transformers.js // DEFAULT_DTYPE_SUFFIX_MAPPING (utils/dtypes.js). fp32 has no filename suffix. @@ -243,6 +255,97 @@ function withTimeout(promise: Promise, ms: number): Promise { ]); } +/** Remote embedding call failed with an HTTP status. */ +export class EmbeddingHttpError extends Error { + constructor( + public readonly status: number, + statusText: string + ) { + super(`API embedding failed: ${status} ${statusText}`); + this.name = "EmbeddingHttpError"; + } +} + +/** + * Classify an embedding fetch failure as retryable. Only transient conditions + * are retried: HTTP 429/5xx, per-attempt AbortSignal timeouts, and network + * TypeErrors from fetch. JSON parse errors, other 4xx, and unexpected errors + * fail fast so programming/shape bugs are not amplified across attempts. + */ +function isRetryableEmbeddingError(error: unknown): boolean { + if (error instanceof EmbeddingHttpError) { + return error.status === 429 || error.status >= 500; + } + if (error instanceof DOMException) { + return error.name === "TimeoutError" || error.name === "AbortError"; + } + // fetch network failures surface as TypeError; AbortSignal.timeout may too + // depending on runtime (Bun/Node) — accept TimeoutError/AbortError names. + if (error instanceof Error) { + if (error.name === "TimeoutError" || error.name === "AbortError") return true; + if (error instanceof TypeError) return true; + } + return false; +} + +/** + * Perform a remote embedding request with bounded retries. Each attempt uses + * AbortSignal.timeout so a hung connection is cancelled (not orphaned via + * Promise.race); transient failures (network, abort/timeout, 429, 5xx) are + * retried with linear backoff, everything else fails fast. + * + * `fetchImpl` is injectable for tests. + */ +export async function fetchRemoteEmbeddingWithRetry( + fetchImpl: typeof fetch, + input: string, + options: { + url?: string; + apiKey?: string; + model?: string; + attemptTimeoutMs?: number; + maxAttempts?: number; + backoffMs?: number; + } = {} +): Promise { + const url = options.url ?? CONFIG.embeddingApiUrl; + const apiKey = options.apiKey ?? CONFIG.embeddingApiKey; + const model = options.model ?? CONFIG.embeddingModel; + const attemptTimeoutMs = options.attemptTimeoutMs ?? REMOTE_ATTEMPT_TIMEOUT_MS; + const maxAttempts = options.maxAttempts ?? REMOTE_MAX_ATTEMPTS; + const backoffMs = options.backoffMs ?? REMOTE_RETRY_BACKOFF_MS; + + let lastError: unknown; + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + try { + const response = await fetchImpl(`${url}/embeddings`, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify({ input, model }), + signal: AbortSignal.timeout(attemptTimeoutMs), + }); + if (!response.ok) { + throw new EmbeddingHttpError(response.status, response.statusText); + } + try { + return await response.json(); + } catch (parseError) { + throw new Error("Invalid embedding API JSON", { cause: parseError }); + } + } catch (error) { + lastError = error; + if (attempt === maxAttempts || !isRetryableEmbeddingError(error)) { + throw error instanceof Error ? error : new Error(String(error)); + } + await new Promise((resolve) => setTimeout(resolve, backoffMs * attempt)); + } + } + throw lastError instanceof Error ? lastError : new Error(String(lastError)); +} + export class EmbeddingService { private pipe: any = null; private initPromise: Promise | null = null; @@ -270,29 +373,11 @@ export class EmbeddingService { private async initializeModel(progressCallback?: (progress: any) => void): Promise { try { if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) { - // Send a probe request to verify the API endpoint is actually reachable - // Uses a minimal embedding of "ping" to test the full request pipeline - const probeResponse = await withTimeout( - fetch(`${CONFIG.embeddingApiUrl}/embeddings`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${CONFIG.embeddingApiKey}`, - }, - body: JSON.stringify({ - input: "ping", - model: CONFIG.embeddingModel, - }), - }), - TIMEOUT_MS - ); - - if (!probeResponse.ok) { - throw new Error( - `Embedding API health check failed: ${probeResponse.status} ${probeResponse.statusText}` - ); - } - + // Send a probe request to verify the API endpoint is actually reachable. + // Uses a minimal embedding of "ping" to test the full request pipeline. + // Retried on transient failure (#400): a single hung probe previously + // burned the whole 30s budget and set the sticky initError. + await fetchRemoteEmbeddingWithRetry(fetch, "ping"); this.isWarmedUp = true; this.initError = null; return; @@ -351,23 +436,7 @@ export class EmbeddingService { let result: Float32Array; if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) { - const response = await fetch(`${CONFIG.embeddingApiUrl}/embeddings`, { - method: "POST", - headers: { - "Content-Type": "application/json", - Authorization: `Bearer ${CONFIG.embeddingApiKey}`, - }, - body: JSON.stringify({ - input, - model: CONFIG.embeddingModel, - }), - }); - - if (!response.ok) { - throw new Error(`API embedding failed: ${response.statusText}`); - } - - const data: any = await response.json(); + const data: any = await fetchRemoteEmbeddingWithRetry(fetch, input); result = new Float32Array(data.data[0].embedding); } else { const output = await this.pipe(input, { diff --git a/tests/embedding-remote-retry.test.ts b/tests/embedding-remote-retry.test.ts new file mode 100644 index 00000000..d9d72f75 --- /dev/null +++ b/tests/embedding-remote-retry.test.ts @@ -0,0 +1,173 @@ +import { describe, it, expect } from "bun:test"; +import { fetchRemoteEmbeddingWithRetry, EmbeddingHttpError } from "../src/services/embedding.js"; + +function jsonResponse(status: number): Response { + return { + ok: status >= 200 && status < 300, + status, + statusText: status === 200 ? "OK" : "Error", + json: async () => ({ data: [{ embedding: [0.1, 0.2] }] }), + } as unknown as Response; +} + +const RETRY_OPTS = { + url: "https://fake.local/v1", + apiKey: "test-key", + model: "test-model", + attemptTimeoutMs: 50, + maxAttempts: 3, + backoffMs: 5, +}; + +describe("fetchRemoteEmbeddingWithRetry", () => { + it("recovers when the first attempt hangs past the per-attempt AbortSignal timeout", async () => { + let calls = 0; + const fetchImpl = async (_url: unknown, init?: { signal?: AbortSignal }) => { + calls++; + if (calls === 1) { + expect(init?.signal).toBeDefined(); + await new Promise((_resolve, reject) => { + const signal = init!.signal!; + if (signal.aborted) { + reject(signal.reason ?? new DOMException("The operation was aborted.", "TimeoutError")); + return; + } + signal.addEventListener( + "abort", + () => { + reject( + signal.reason ?? new DOMException("The operation was aborted.", "TimeoutError") + ); + }, + { once: true } + ); + }); + } + return jsonResponse(200); + }; + const data = await fetchRemoteEmbeddingWithRetry( + fetchImpl as unknown as typeof fetch, + "hello", + RETRY_OPTS + ); + expect(data.data[0].embedding).toEqual([0.1, 0.2]); + expect(calls).toBe(2); + }); + + it("retries on network TypeError and succeeds on a later attempt", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + if (calls === 1) throw new TypeError("fetch failed"); + return jsonResponse(200); + }; + const data = await fetchRemoteEmbeddingWithRetry( + fetchImpl as unknown as typeof fetch, + "hello", + RETRY_OPTS + ); + expect(data.data[0].embedding).toEqual([0.1, 0.2]); + expect(calls).toBe(2); + }); + + it("fails fast on invalid JSON without consuming another attempt", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + return { + ok: true, + status: 200, + statusText: "OK", + json: async () => { + throw new SyntaxError("Unexpected token"); + }, + } as unknown as Response; + }; + await expect( + fetchRemoteEmbeddingWithRetry(fetchImpl as unknown as typeof fetch, "hello", RETRY_OPTS) + ).rejects.toThrow("Invalid embedding API JSON"); + expect(calls).toBe(1); + }); + + it("passes AbortSignal.timeout to each fetch attempt", async () => { + let seenSignal: AbortSignal | undefined; + const fetchImpl = async (_url: unknown, init?: { signal?: AbortSignal }) => { + seenSignal = init?.signal; + return jsonResponse(200); + }; + await fetchRemoteEmbeddingWithRetry(fetchImpl as unknown as typeof fetch, "hello", RETRY_OPTS); + expect(seenSignal).toBeDefined(); + expect(typeof seenSignal!.aborted).toBe("boolean"); + }); + + it("retries on 429 and succeeds on a later attempt", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + return calls === 1 ? jsonResponse(429) : jsonResponse(200); + }; + const data = await fetchRemoteEmbeddingWithRetry( + fetchImpl as unknown as typeof fetch, + "hello", + RETRY_OPTS + ); + expect(data.data[0].embedding).toEqual([0.1, 0.2]); + expect(calls).toBe(2); + }); + + it("retries on 500 and succeeds on a later attempt", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + return calls === 1 ? jsonResponse(500) : jsonResponse(200); + }; + await fetchRemoteEmbeddingWithRetry(fetchImpl as unknown as typeof fetch, "hello", RETRY_OPTS); + expect(calls).toBe(2); + }); + + it("fails fast on non-retryable 4xx without consuming another attempt", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + return jsonResponse(401); + }; + await expect( + fetchRemoteEmbeddingWithRetry(fetchImpl as unknown as typeof fetch, "hello", RETRY_OPTS) + ).rejects.toBeInstanceOf(EmbeddingHttpError); + expect(calls).toBe(1); + }); + + it("throws the last error after exhausting all attempts on persistent 5xx", async () => { + let calls = 0; + const fetchImpl = async () => { + calls++; + return jsonResponse(503); + }; + await expect( + fetchRemoteEmbeddingWithRetry(fetchImpl as unknown as typeof fetch, "hello", RETRY_OPTS) + ).rejects.toBeInstanceOf(EmbeddingHttpError); + expect(calls).toBe(3); + }); + + it("preserves the request shape: base url, bearer auth, prefixed input, model", async () => { + let seen: { url: string; auth: string | undefined; body: any } | null = null; + const fetchImpl = async (url: unknown, init?: { headers?: Record }) => { + seen = { + url: String(url), + auth: init?.headers?.["Authorization"], + body: JSON.parse(String(init?.body)), + }; + return jsonResponse(200); + }; + await fetchRemoteEmbeddingWithRetry( + fetchImpl as unknown as typeof fetch, + "search_query: hello", + RETRY_OPTS + ); + expect(seen).not.toBeNull(); + expect(seen!.url).toBe("https://fake.local/v1/embeddings"); + expect(seen!.auth).toBe("Bearer test-key"); + expect(seen!.body.input).toBe("search_query: hello"); + expect(seen!.body.model).toBe("test-model"); + }); +});