@vellumai/assistant 0.10.4-staging.1 → 0.10.4

package/src/persistence/embeddings/__tests__/embedding-identity.test.ts

@@ -0,0 +1,156 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type { AssistantConfig } from "../../../config/types.js";
+
+const CONFIG = {
+  memory: { qdrant: { url: "http://127.0.0.1:6333", vectorSize: 768 } },
+} as unknown as AssistantConfig;
+
+// ── Stub selectEmbeddingBackend + the billing breaker ──────────────
+let backendToReturn: {
+  provider: string;
+  model: string;
+  embed: (inputs: string[]) => Promise<number[][]>;
+} | null = null;
+let breakerOpen = false;
+
+const selectEmbeddingBackendMock = mock(async () => ({
+  backend: backendToReturn,
+  reason: backendToReturn ? null : "no backend",
+}));
+// probeBackendDimension delegates the measurement to resolveBackendDimension;
+// the mock mirrors the real resolver's "embed → vector length, null on throw"
+// contract so these tests drive behavior through the stub backend's embed().
+const resolveBackendDimensionMock = mock(
+  async (backend: { embed: (inputs: string[]) => Promise<number[][]> }) => {
+    try {
+      const [vector] = await backend.embed(["embedding dimension probe"]);
+      return vector?.length ?? null;
+    } catch {
+      return null;
+    }
+  },
+);
+mock.module("../embedding-backend.js", () => ({
+  selectEmbeddingBackend: selectEmbeddingBackendMock,
+  resolveBackendDimension: resolveBackendDimensionMock,
+}));
+
+mock.module("../embedding-billing-breaker.js", () => ({
+  isEmbeddingBillingBreakerOpen: () => breakerOpen,
+}));
+
+// ── Stub the Qdrant REST client ────────────────────────────────────
+let collectionExistsResult = { exists: true };
+let getCollectionResult: unknown = {
+  config: { params: { vectors: { dense: { size: 768 } } } },
+};
+let collectionExistsThrows = false;
+let getCollectionThrows = false;
+
+const collectionExistsMock = mock(async () => {
+  if (collectionExistsThrows) throw new Error("qdrant down");
+  return collectionExistsResult;
+});
+const getCollectionMock = mock(async () => {
+  if (getCollectionThrows) throw new Error("qdrant probe failed");
+  return getCollectionResult;
+});
+mock.module("@qdrant/js-client-rest", () => ({
+  QdrantClient: class {
+    collectionExists = collectionExistsMock;
+    getCollection = getCollectionMock;
+  },
+}));
+
+const { probeBackendDimension, readConceptPageCollectionDim } =
+  await import("../embedding-identity.js");
+
+describe("probeBackendDimension", () => {
+  beforeEach(() => {
+    breakerOpen = false;
+    backendToReturn = {
+      provider: "openai",
+      model: "text-embedding-3-small",
+      embed: async () => [new Array(1536).fill(0)],
+    };
+  });
+  afterEach(() => {
+    selectEmbeddingBackendMock.mockClear();
+  });
+
+  test("returns the measured dimension on success", async () => {
+    const result = await probeBackendDimension(CONFIG);
+    expect(result).toEqual({
+      provider: "openai",
+      model: "text-embedding-3-small",
+      dim: 1536,
+    });
+  });
+
+  test("returns null when the billing breaker is open", async () => {
+    breakerOpen = true;
+    expect(await probeBackendDimension(CONFIG)).toBeNull();
+    // Breaker short-circuits before selecting a backend.
+    expect(selectEmbeddingBackendMock).not.toHaveBeenCalled();
+  });
+
+  test("returns null when no backend is selectable", async () => {
+    backendToReturn = null;
+    expect(await probeBackendDimension(CONFIG)).toBeNull();
+  });
+
+  test("returns null when the backend embed call throws", async () => {
+    backendToReturn = {
+      provider: "openai",
+      model: "text-embedding-3-small",
+      embed: async () => {
+        throw new Error("provider unreachable");
+      },
+    };
+    expect(await probeBackendDimension(CONFIG)).toBeNull();
+  });
+});
+
+describe("readConceptPageCollectionDim", () => {
+  beforeEach(() => {
+    collectionExistsResult = { exists: true };
+    getCollectionResult = {
+      config: { params: { vectors: { dense: { size: 768 } } } },
+    };
+    collectionExistsThrows = false;
+    getCollectionThrows = false;
+  });
+
+  test("returns the committed dense vector size", async () => {
+    expect(await readConceptPageCollectionDim(CONFIG)).toBe(768);
+  });
+
+  test("returns null only when the collection is confirmed absent", async () => {
+    collectionExistsResult = { exists: false };
+    expect(await readConceptPageCollectionDim(CONFIG)).toBeNull();
+  });
+
+  test("throws when an existing collection's dense size is unreadable", async () => {
+    // Collection exists but the dense vector size is missing: "unknown" must
+    // not be reported as "absent" (which the reconcile reads as a fresh
+    // install), so this is an error.
+    getCollectionResult = { config: { params: { vectors: {} } } };
+    expect(readConceptPageCollectionDim(CONFIG)).rejects.toThrow();
+  });
+
+  test("propagates a Qdrant read failure (existence check throws)", async () => {
+    // A transient Qdrant outage must propagate rather than be misread as a
+    // confirmed-absent collection, so the reconcile defers instead of committing
+    // a new dimension while the old collection still exists.
+    collectionExistsThrows = true;
+    expect(readConceptPageCollectionDim(CONFIG)).rejects.toThrow("qdrant down");
+  });
+
+  test("propagates a Qdrant read failure (getCollection throws)", async () => {
+    getCollectionThrows = true;
+    expect(readConceptPageCollectionDim(CONFIG)).rejects.toThrow(
+      "qdrant probe failed",
+    );
+  });
+});

package/src/persistence/embeddings/__tests__/reconcile-decision.test.ts

@@ -0,0 +1,95 @@
+import { describe, expect, test } from "bun:test";
+
+import {
+  decideEmbeddingReconcile,
+  type ReconcileAction,
+} from "../reconcile-decision.js";
+
+describe("decideEmbeddingReconcile", () => {
+  test("fresh install, backend up -> commit-fresh", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: null,
+      probeDim: 3072,
+      configuredProvider: "gemini",
+    });
+    expect(action).toEqual({ kind: "commit-fresh", dim: 3072 });
+  });
+
+  test("fresh install, backend down -> defer-degraded", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: null,
+      probeDim: null,
+      configuredProvider: "auto",
+    });
+    expect(action).toEqual({
+      kind: "defer-degraded",
+      reason: "no reachable embedding backend",
+    });
+  });
+
+  test("committed matches probe -> noop", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 384,
+      probeDim: 384,
+      configuredProvider: "auto",
+    });
+    expect(action).toEqual({ kind: "noop" });
+  });
+
+  test("explicit provider, mismatch upward -> migrate", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 384,
+      probeDim: 3072,
+      configuredProvider: "gemini",
+    });
+    expect(action).toEqual({ kind: "migrate", fromDim: 384, toDim: 3072 });
+  });
+
+  test("explicit provider, mismatch downward -> migrate", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 3072,
+      probeDim: 384,
+      configuredProvider: "local",
+    });
+    expect(action).toEqual({ kind: "migrate", fromDim: 3072, toDim: 384 });
+  });
+
+  test("auto provider, mismatch upward -> noop (no thrash)", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 384,
+      probeDim: 3072,
+      configuredProvider: "auto",
+    });
+    expect(action).toEqual({ kind: "noop" });
+  });
+
+  test("auto provider, mismatch downward -> noop (no thrash)", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 3072,
+      probeDim: 384,
+      configuredProvider: "auto",
+    });
+    expect(action).toEqual({ kind: "noop" });
+  });
+
+  test("existing collection, backend down -> defer-degraded", () => {
+    const action = decideEmbeddingReconcile({
+      committedDim: 3072,
+      probeDim: null,
+      configuredProvider: "gemini",
+    });
+    expect(action).toEqual({
+      kind: "defer-degraded",
+      reason: "no reachable embedding backend",
+    });
+  });
+
+  test("defer-degraded precedes fresh-commit when both committed and probe are null", () => {
+    const action: ReconcileAction = decideEmbeddingReconcile({
+      committedDim: null,
+      probeDim: null,
+      configuredProvider: "gemini",
+    });
+    expect(action.kind).toBe("defer-degraded");
+  });
+});

package/src/persistence/embeddings/embed.ts

@@ -0,0 +1,87 @@
+// ---------------------------------------------------------------------------
+// Shared embedding utility with retry + exponential backoff
+// ---------------------------------------------------------------------------
+
+import type { AssistantConfig } from "../../config/types.js";
+import { getLogger } from "../../util/logger.js";
+import {
+  abortableSleep,
+  computeRetryDelay,
+  isRetryableNetworkError,
+} from "../../util/retry.js";
+import { type EmbeddingInput, embedWithBackend } from "./embedding-backend.js";
+
+const log = getLogger("memory-embed");
+
+export const EMBED_MAX_RETRIES = 3;
+export const EMBED_BASE_DELAY_MS = 500;
+
+/**
+ * Wrap embedWithBackend with retry + exponential backoff for transient failures
+ * (network errors, 429s, 5xx). Aborts immediately if the caller's signal fires.
+ */
+export async function embedWithRetry(
+  config: AssistantConfig,
+  texts: EmbeddingInput[],
+  opts?: { signal?: AbortSignal },
+): ReturnType<typeof embedWithBackend> {
+  let lastError: unknown;
+  for (let attempt = 0; attempt <= EMBED_MAX_RETRIES; attempt++) {
+    try {
+      return await embedWithBackend(config, texts, opts);
+    } catch (err) {
+      lastError = err;
+      if (opts?.signal?.aborted || isAbortError(err)) throw err;
+      const isTransient = isTransientEmbeddingError(err);
+      if (!isTransient || attempt === EMBED_MAX_RETRIES) throw err;
+      const delay = computeRetryDelay(attempt, EMBED_BASE_DELAY_MS);
+      log.warn(
+        { err, attempt: attempt + 1, delayMs: Math.round(delay) },
+        "Transient embedding failure, retrying",
+      );
+      await abortableSleep(delay, opts?.signal);
+      if (opts?.signal?.aborted) throw err;
+    }
+  }
+  throw lastError;
+}
+
+/**
+ * A caller-initiated abort, never a transient failure — must propagate
+ * immediately rather than being retried.
+ */
+export function isAbortError(err: unknown): boolean {
+  if (!(err instanceof Error)) return false;
+  return err.name === "AbortError" || err.name === "APIUserAbortError";
+}
+
+/**
+ * Whether an embedding failure is transient and worth retrying: a retryable
+ * network error (ECONNRESET, timeout, …) or a 429 / 5xx from the provider.
+ * Shared so other embedding paths (e.g. the v3 candidate matcher, which embeds
+ * via `simBatch`) can apply the same bounded-retry policy as
+ * {@link embedWithRetry} instead of degrading on the first blip.
+ */
+export function isTransientEmbeddingError(err: unknown): boolean {
+  return isRetryableNetworkError(err) || isHttpStatusError(err);
+}
+
+function getErrorStatusCode(err: Error): unknown {
+  if ("status" in err) {
+    const status = (err as { status: unknown }).status;
+    if (status != null) return status;
+  }
+  if ("statusCode" in err) return (err as { statusCode: unknown }).statusCode;
+  return undefined;
+}
+
+function isHttpStatusError(err: unknown): boolean {
+  if (!(err instanceof Error)) return false;
+  const status = getErrorStatusCode(err);
+  if (typeof status === "number") {
+    return status === 429 || (status >= 500 && status < 600);
+  }
+  return /\b429\b|(?:failed|error)\s*\((?:429|5\d{2})\)|(?:status|http)\s*(?:code\s*)?:?\s*5\d{2}\b/i.test(
+    err.message,
+  );
+}

package/src/persistence/embeddings/embedding-backend.test.ts

@@ -0,0 +1,355 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type { AssistantConfig } from "../../config/types.js";
+import {
+  clearEmbeddingBackendCache,
+  embedWithBackend,
+  isEmbeddingDimensionAvailable,
+  resetLocalEmbeddingFailureState,
+  selectEmbeddingBackend,
+} from "./embedding-backend.js";
+import {
+  _resetEmbeddingBillingBreaker,
+  recordBillingBlock,
+} from "./embedding-billing-breaker.js";
+
+const getProviderKeyAsyncMock = mock(
+  async (_provider: string): Promise<string | undefined> => undefined,
+);
+mock.module("../../security/secure-keys.js", () => ({
+  getProviderKeyAsync: getProviderKeyAsyncMock,
+}));
+
+const DISABLED_PROXY_CONTEXT = {
+  enabled: false,
+  platformBaseUrl: "",
+  assistantApiKey: "",
+};
+const resolveManagedProxyContextMock = mock(async () => DISABLED_PROXY_CONTEXT);
+mock.module("../../providers/platform-proxy/context.js", () => ({
+  resolveManagedProxyContext: resolveManagedProxyContextMock,
+  hasManagedProxyPrereqs: mock(async () => false),
+  buildManagedBaseUrl: mock(async () => undefined),
+  managedFallbackEnabledFor: mock(async () => false),
+}));
+
+const LOCAL_CONFIG = {
+  memory: {
+    embeddings: {
+      provider: "local",
+      localModel: "BAAI/bge-small-en-v1.5",
+    },
+  },
+} as unknown as AssistantConfig;
+
+describe("embedding backend cache invalidation", () => {
+  afterEach(() => {
+    clearEmbeddingBackendCache();
+  });
+
+  test("clearEmbeddingBackendCache disposes cached backends before clearing", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const dispose = mock();
+    (firstSelection.backend as { dispose?: () => void }).dispose = dispose;
+
+    clearEmbeddingBackendCache();
+
+    expect(dispose).toHaveBeenCalledTimes(1);
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).not.toBe(firstSelection.backend);
+  });
+
+  test("resetLocalEmbeddingFailureState preserves live cached backends", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const dispose = mock();
+    (firstSelection.backend as { dispose?: () => void }).dispose = dispose;
+
+    resetLocalEmbeddingFailureState();
+
+    expect(dispose).not.toHaveBeenCalled();
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).toBe(firstSelection.backend);
+  });
+
+  test("resetLocalEmbeddingFailureState clears poisoned local backend retry state", async () => {
+    const firstSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(firstSelection.backend).not.toBeNull();
+
+    const poisonedInitPromise = Promise.reject(new Error("poisoned"));
+    poisonedInitPromise.catch(() => {});
+
+    const backend = firstSelection.backend as unknown as {
+      delegate: unknown;
+      initPromise: Promise<unknown> | null;
+    };
+    backend.delegate = null;
+    backend.initPromise = poisonedInitPromise;
+
+    resetLocalEmbeddingFailureState();
+
+    expect(backend.initPromise).toBeNull();
+
+    const secondSelection = await selectEmbeddingBackend(LOCAL_CONFIG);
+    expect(secondSelection.backend).toBe(firstSelection.backend);
+  });
+});
+
+describe("isEmbeddingDimensionAvailable", () => {
+  afterEach(() => {
+    clearEmbeddingBackendCache();
+    _resetEmbeddingBillingBreaker();
+    getProviderKeyAsyncMock.mockReset();
+    getProviderKeyAsyncMock.mockImplementation(async () => undefined);
+  });
+
+  function localConfigWithVectorSize(vectorSize: number): AssistantConfig {
+    return {
+      memory: {
+        enabled: true,
+        embeddings: {
+          provider: "local",
+          localModel: "BAAI/bge-small-en-v1.5",
+          required: false,
+        },
+        qdrant: { vectorSize },
+      },
+    } as unknown as AssistantConfig;
+  }
+
+  /** Stub the selected backend's `embed` so the dimension probe is deterministic. */
+  async function stubSelectedBackendProbeDim(
+    config: AssistantConfig,
+    dim: number,
+  ): Promise<ReturnType<typeof mock>> {
+    const { backend } = await selectEmbeddingBackend(config);
+    if (!backend) throw new Error("expected a backend to be selected");
+    const embedMock = mock(async () => [new Array(dim).fill(0)]);
+    (backend as { embed: typeof backend.embed }).embed =
+      embedMock as unknown as typeof backend.embed;
+    return embedMock;
+  }
+
+  test("returns true when the reachable backend matches the committed dimension", async () => {
+    const config = localConfigWithVectorSize(384);
+    await stubSelectedBackendProbeDim(config, 384);
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(true);
+  });
+
+  test("returns false when the reachable backend's dimension differs from the committed one", async () => {
+    // 3072-dim collection committed, but only a 384-dim backend is reachable.
+    const config = localConfigWithVectorSize(3072);
+    await stubSelectedBackendProbeDim(config, 384);
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(false);
+  });
+
+  test("returns false when the backend probe throws (unreachable)", async () => {
+    const config = localConfigWithVectorSize(384);
+    const { backend } = await selectEmbeddingBackend(config);
+    if (!backend) throw new Error("expected a backend to be selected");
+    (backend as { embed: typeof backend.embed }).embed = (async () => {
+      throw new Error("backend unreachable");
+    }) as unknown as typeof backend.embed;
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(false);
+  });
+
+  test("memoizes the probe so repeated calls embed only once", async () => {
+    const config = localConfigWithVectorSize(384);
+    const embedMock = await stubSelectedBackendProbeDim(config, 384);
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(true);
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(true);
+
+    expect(embedMock).toHaveBeenCalledTimes(1);
+  });
+
+  test("returns false when memory is disabled", async () => {
+    const config = {
+      memory: {
+        enabled: false,
+        embeddings: { provider: "local", localModel: "m", required: false },
+        qdrant: { vectorSize: 384 },
+      },
+    } as unknown as AssistantConfig;
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(false);
+  });
+
+  test("under auto, returns true when a configured fallback matches the committed dimension", async () => {
+    // Primary (local) produces 384, the committed collection is 3072, and a
+    // configured OpenAI fallback produces 3072. `embedWithBackend` would fall
+    // through to the fallback, so dense recall must stay available rather than
+    // being disabled because only the primary mismatches.
+    getProviderKeyAsyncMock.mockImplementation(async (p: string) =>
+      p === "openai" ? "sk-test" : undefined,
+    );
+    const config = {
+      // `selectFallbackBackends` probes ollama via `resolveCallSiteConfig`, which
+      // reads `llm.default`; provide a minimal one so the fallback chain builds.
+      llm: { default: {} },
+      memory: {
+        enabled: true,
+        embeddings: {
+          provider: "auto",
+          localModel: "BAAI/bge-small-en-v1.5",
+          openaiModel: "text-embedding-3-small",
+          required: false,
+        },
+        qdrant: { vectorSize: 3072 },
+      },
+    } as unknown as AssistantConfig;
+
+    // Stub the primary (local) → 384.
+    const { backend: primary } = await selectEmbeddingBackend(config);
+    if (!primary) throw new Error("expected a primary backend");
+    (primary as { embed: typeof primary.embed }).embed = mock(async () => [
+      new Array(384).fill(0),
+    ]) as unknown as typeof primary.embed;
+
+    // Warm + stub the OpenAI fallback instance (cached by provider:model, so the
+    // chain probe inside isEmbeddingDimensionAvailable reuses this stub) → 3072.
+    const { backend: openai } = await selectEmbeddingBackend({
+      memory: {
+        embeddings: {
+          provider: "openai",
+          openaiModel: "text-embedding-3-small",
+        },
+        qdrant: { vectorSize: 3072 },
+      },
+    } as unknown as AssistantConfig);
+    if (!openai) throw new Error("expected an openai fallback backend");
+    (openai as { embed: typeof openai.embed }).embed = mock(async () => [
+      new Array(3072).fill(0),
+    ]) as unknown as typeof openai.embed;
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(true);
+  });
+
+  test("returns false (and does not probe) when the billing breaker is open", async () => {
+    const config = localConfigWithVectorSize(384);
+    const embedMock = await stubSelectedBackendProbeDim(config, 384);
+
+    recordBillingBlock(); // open the breaker
+
+    expect(await isEmbeddingDimensionAvailable(config)).toBe(false);
+    // The probe is skipped entirely so a depleted-billing breaker is not
+    // hammered with per-lane dimension probes.
+    expect(embedMock).toHaveBeenCalledTimes(0);
+  });
+});
+
+const GEMINI_CONFIG = {
+  memory: {
+    embeddings: {
+      provider: "gemini",
+      geminiModel: "test-model",
+    },
+    qdrant: {
+      vectorSize: 3,
+    },
+  },
+} as unknown as AssistantConfig;
+
+const ENABLED_PROXY_CONTEXT = {
+  enabled: true,
+  platformBaseUrl: "https://proxy.example.com",
+  assistantApiKey: "stale-platform-key",
+};
+
+function jsonResponse(body: unknown, status: number): Response {
+  return new Response(JSON.stringify(body), {
+    status,
+    headers: { "Content-Type": "application/json" },
+  });
+}
+
+describe("managed-proxy Gemini fallback to direct key", () => {
+  const originalFetch = globalThis.fetch;
+  let fetchMock: ReturnType<typeof mock>;
+
+  beforeEach(() => {
+    clearEmbeddingBackendCache();
+    // Managed-proxy requests (proxy.example.com) are rejected like a platform
+    // proxy with a bad credential; direct Google API requests succeed.
+    fetchMock = mock(async (url: string | URL) => {
+      if (String(url).startsWith("https://proxy.example.com/")) {
+        return jsonResponse({ detail: "Invalid or revoked API key." }, 403);
+      }
+      return jsonResponse({ embedding: { values: [0.1, 0.2, 0.3] } }, 200);
+    });
+    globalThis.fetch = fetchMock as unknown as typeof fetch;
+  });
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+    clearEmbeddingBackendCache();
+    getProviderKeyAsyncMock.mockReset();
+    getProviderKeyAsyncMock.mockImplementation(async () => undefined);
+    resolveManagedProxyContextMock.mockReset();
+    resolveManagedProxyContextMock.mockImplementation(
+      async () => DISABLED_PROXY_CONTEXT,
+    );
+  });
+
+  test("falls back to the direct key when the managed-proxy primary fails", async () => {
+    resolveManagedProxyContextMock.mockImplementation(
+      async () => ENABLED_PROXY_CONTEXT,
+    );
+    getProviderKeyAsyncMock.mockImplementation(async (provider) =>
+      provider === "gemini" ? "direct-key" : undefined,
+    );
+
+    const result = await embedWithBackend(GEMINI_CONFIG, ["hello world"]);
+
+    expect(result.provider).toBe("gemini");
+    expect(result.vectors).toEqual([[0.1, 0.2, 0.3]]);
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    const [managedUrl, managedInit] = fetchMock.mock.calls[0] as [
+      string,
+      RequestInit,
+    ];
+    expect(managedUrl).toBe(
+      "https://proxy.example.com/v1/runtime-proxy/gemini/v1beta/models/test-model:embedContent",
+    );
+    expect(
+      (managedInit.headers as Record<string, string>)["Authorization"],
+    ).toBe("Bearer stale-platform-key");
+    const [directUrl] = fetchMock.mock.calls[1] as [string, RequestInit];
+    expect(directUrl).toContain("generativelanguage.googleapis.com");
+    expect(directUrl).toContain("key=direct-key");
+  });
+
+  test("surfaces the managed-proxy error when no direct key exists", async () => {
+    resolveManagedProxyContextMock.mockImplementation(
+      async () => ENABLED_PROXY_CONTEXT,
+    );
+
+    await expect(
+      embedWithBackend(GEMINI_CONFIG, ["hello world"]),
+    ).rejects.toThrow(/403/);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  test("uses the direct key as primary when the managed proxy is disabled", async () => {
+    getProviderKeyAsyncMock.mockImplementation(async (provider) =>
+      provider === "gemini" ? "direct-key" : undefined,
+    );
+
+    const result = await embedWithBackend(GEMINI_CONFIG, ["another input"]);
+
+    expect(result.vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    const [directUrl] = fetchMock.mock.calls[0] as [string, RequestInit];
+    expect(directUrl).toContain("generativelanguage.googleapis.com");
+    expect(directUrl).toContain("key=direct-key");
+  });
+});