fix: retry per-attempt timeouts and validate vector dimensions

README.md

@@ -53,15 +53,15 @@ console.log(vectors.length, vectors[0]?.length); // 2 768
 
 ### Config
 
-| Field            | Type                   | Description                                                              |
-| ---------------- | ---------------------- | ------------------------------------------------------------------------ |
-| `baseURL`        | `string`               | Server root with its version path, such as `http://host:11434/v1`.       |
-| `model`          | `string`               | Model name.                                                              |
-| `apiKey`         | `string?`              | Sent as a bearer token.                                                  |
-| `dimensions`     | `number?`              | Requested output dimension. Sent only when set; not all models honor it. |
-| `encodingFormat` | `"float" \| "base64"?` | Wire format. Defaults to `float`. Results are always `number[]`.         |
-| `batchSize`      | `number?`              | Texts per request. Defaults to 32.                                       |
-| `timeoutMs`      | `number?`              | Per-request timeout. Defaults to 30000.                                  |
+| Field            | Type                   | Description                                                                                 |
+| ---------------- | ---------------------- | ------------------------------------------------------------------------------------------- |
+| `baseURL`        | `string`               | Server root with its version path, such as `http://host:11434/v1`.                          |
+| `model`          | `string`               | Model name.                                                                                 |
+| `apiKey`         | `string?`              | Sent as a bearer token.                                                                     |
+| `dimensions`     | `number?`              | Requested output dimension. Sent only when set; a reply of any other dimension is rejected. |
+| `encodingFormat` | `"float" \| "base64"?` | Wire format. Defaults to `float`. Results are always `number[]`.                            |
+| `batchSize`      | `number?`              | Texts per request. Defaults to 32.                                                          |
+| `timeoutMs`      | `number?`              | Per-attempt timeout; a timed-out attempt is retried. Defaults to 30000.                     |
 
 ### Options
 
@@ -78,7 +78,7 @@ A failed request throws `EmbeddingRequestError` with a classified `reason` and t
 
 ### Dimensions
 
-Each model has a fixed output dimension: 768 for `nomic-embed-text`, 1536 for OpenAI `text-embedding-3-small`. If you store vectors, call `probeEmbedDims` at startup. Changing models changes the dimension, so treat it as a schema migration.
+Each model has a fixed output dimension: 768 for `nomic-embed-text`, 1536 for OpenAI `text-embedding-3-small`. If you store vectors, call `probeEmbedDims` at startup. Changing models changes the dimension, so treat it as a schema migration. Every vector from one call has the same non-zero dimension, or the call throws `EmbeddingRequestError` with `protocol_mismatch`.
 
 ## Using with Interchange
 

e2e/retry.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, expect, test } from "bun:test";
+import { createDefaultScheduler } from "@intx/inference";
 import { setupHarness, type Harness } from "@intx/inference-testing";
 
 import { EmbeddingRequestError, embedTexts } from "../src/index";
@@ -95,17 +96,72 @@ test("honors a caller-supplied extractRetryAfterMs", async () => {
   expect(harness.clock.now()).toBeGreaterThanOrEqual(45_000);
 });
 
-test("applies timeoutMs even when the caller passes a signal", async () => {
-  harness = setupHarness();
-  const stream = harness.scenario.createStream();
-  harness.scenario.whenRequestMatches(() => true, stream);
+/** Hangs until aborted for the first `hangs` calls, then replies. */
+function hangingFetch(hangs: number): {
+  fetch: typeof fetch;
+  calls: () => number;
+} {
+  let calls = 0;
+  const hangThenReply = (async (_url: unknown, init?: RequestInit) => {
+    calls++;
+    if (calls > hangs) {
+      return new Response(
+        JSON.stringify({ data: [{ index: 0, embedding: [1] }] }),
+      );
+    }
+    const signal = init?.signal ?? undefined;
+    return await new Promise<Response>((_resolve, reject) => {
+      signal?.addEventListener("abort", () => reject(signal.reason));
+    });
+  }) as typeof fetch;
+  return { fetch: hangThenReply, calls: () => calls };
+}
+
+test("retries an attempt that hits timeoutMs", async () => {
+  const stub = hangingFetch(1);
+  const vectors = await embedTexts(
+    ["a"],
+    { ...CONFIG, timeoutMs: 50 },
+    { deps: { fetch: stub.fetch, scheduler: createDefaultScheduler() } },
+  );
+
+  expect(vectors).toEqual([[1]]);
+  expect(stub.calls()).toBe(2);
+});
 
+test("applies timeoutMs even when the caller passes a signal", async () => {
+  const stub = hangingFetch(Infinity);
   const pending = embedTexts(
     ["a"],
     { ...CONFIG, timeoutMs: 50 },
-    { deps: harness.deps, signal: new AbortController().signal },
+    {
+      deps: { fetch: stub.fetch, scheduler: createDefaultScheduler() },
+      signal: new AbortController().signal,
+    },
+  );
+
+  await expect(pending).rejects.toMatchObject({
+    reason: { category: "timeout" },
+  });
+  expect(stub.calls()).toBe(3);
+});
+
+test("reports a caller abort as aborted, without retrying", async () => {
+  harness = setupHarness({ enableInferenceTimers: true });
+  harness.scenario.whenRequestMatches(
+    () => true,
+    harness.scenario.createStream(),
   );
+  const controller = new AbortController();
+
+  const pending = embedTexts(["a"], CONFIG, {
+    deps: harness.deps,
+    signal: controller.signal,
+  });
+  setTimeout(() => controller.abort(), 20);
   await harness.run();
 
-  await expect(pending).rejects.toBeInstanceOf(EmbeddingRequestError);
+  await expect(pending).rejects.toMatchObject({
+    reason: { category: "aborted" },
+  });
 });

src/embed.test.ts

@@ -115,3 +115,26 @@ test("classifies a 200 that is not JSON instead of leaking a SyntaxError", async
     reason: { category: "protocol_mismatch" },
   });
 });
+
+test("rejects empty, uneven, or off-dimension vectors", async () => {
+  for (const [config, bodies] of [
+    [CONFIG, [{ data: [{ index: 0, embedding: [] }] }]],
+    [
+      { ...CONFIG, batchSize: 1 },
+      [
+        { data: [{ index: 0, embedding: [1, 2] }] },
+        { data: [{ index: 0, embedding: [1] }] },
+      ],
+    ],
+    [{ ...CONFIG, dimensions: 3 }, [{ data: [{ index: 0, embedding: [1] }] }]],
+  ] as const) {
+    harness?.dispose();
+    harness = replies(...bodies);
+    const texts = bodies.map((_, i) => String(i));
+    const pending = embedTexts(texts, config, { deps: harness.deps });
+    await harness.run();
+    await expect(pending).rejects.toMatchObject({
+      reason: { category: "protocol_mismatch" },
+    });
+  }
+});

src/embed.ts

@@ -205,7 +205,20 @@ export async function embedTexts(
       extractRetryAfterMs: extractRetryAfter,
       signal: options.signal,
     });
-    vectors.push(...parseResponse(body, request.url, batch.length));
+    for (const vector of parseResponse(body, request.url, batch.length)) {
+      const width = valid.dimensions ?? vectors[0]?.length ?? vector.length;
+      if (vector.length === 0 || vector.length !== width) {
+        throw new EmbeddingRequestError(
+          classifyProtocolMismatch(
+            vector.length === 0
+              ? "empty embedding"
+              : `expected ${String(width)}-dimensional embeddings, got ${String(vector.length)}`,
+          ),
+          request.url,
+        );
+      }
+      vectors.push(vector);
+    }
   }
   return vectors;
 }

src/request.ts

@@ -85,22 +85,14 @@ export const extractRetryAfterMs: RetryAfterExtractor = (headers) => {
   return Number.isNaN(at) ? undefined : Math.max(0, at - Date.now());
 };
 
-/** Aborts on the caller's signal or the per-attempt timeout, whichever first. */
-function attemptSignal(
-  timeoutMs: number,
-  signal: AbortSignal | undefined,
-): AbortSignal {
-  const timeout = AbortSignal.timeout(timeoutMs);
-  return signal === undefined ? timeout : AbortSignal.any([signal, timeout]);
-}
-
 async function attemptOnce(
   request: BuiltRequest,
   deps: RequestDependencies,
   timeoutMs: number,
   extractRetryAfter: RetryAfterExtractor,
   signal: AbortSignal | undefined,
 ): Promise<Attempt> {
+  const timeout = AbortSignal.timeout(timeoutMs);
   let res: Response;
   let raw: string;
   try {
@@ -109,7 +101,8 @@ async function attemptOnce(
       headers: request.headers,
       body: request.body,
       redirect: "manual",
-      signal: attemptSignal(timeoutMs, signal),
+      signal:
+        signal === undefined ? timeout : AbortSignal.any([signal, timeout]),
     });
 
     if (!res.ok) {
@@ -131,9 +124,20 @@ async function attemptOnce(
     // as a bare SyntaxError.
     raw = await res.text();
   } catch (cause) {
-    const aborted =
-      cause instanceof Error &&
-      (cause.name === "AbortError" || cause.name === "TimeoutError");
+    // A caller abort is final; the per-attempt timeout is retryable.
+    if (signal?.aborted === true) {
+      return { ok: false, error: classifyAbortError() };
+    }
+    if (timeout.aborted) {
+      return {
+        ok: false,
+        error: {
+          category: "timeout",
+          message: `embedding request exceeded ${String(timeoutMs)} ms`,
+        },
+      };
+    }
+    const aborted = cause instanceof Error && cause.name === "AbortError";
     return {
       ok: false,
       error: aborted ? classifyAbortError() : classifyNetworkError(cause),