@vellumai/assistant 0.8.10-dev.202606100742.99a7fab → 0.8.10-dev.202606100925.28c1cfb

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@vellumai/assistant",
-  "version": "0.8.10-dev.202606100742.99a7fab",
+  "version": "0.8.10-dev.202606100925.28c1cfb",
   "license": "MIT",
   "type": "module",
   "exports": {

package/src/memory/embedding-backend.test.ts

@@ -1,12 +1,33 @@
-import { afterEach, describe, expect, mock, test } from "bun:test";
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 
 import type { AssistantConfig } from "../config/types.js";
 import {
   clearEmbeddingBackendCache,
+  embedWithBackend,
   resetLocalEmbeddingFailureState,
   selectEmbeddingBackend,
 } from "./embedding-backend.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: {
@@ -73,3 +94,111 @@ describe("embedding backend cache invalidation", () => {
     expect(secondSelection.backend).toBe(firstSelection.backend);
   });
 });
+
+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");
+  });
+});

package/src/memory/embedding-backend.ts

@@ -266,6 +266,56 @@ function geminiCacheExtras(config: AssistantConfig): string[] {
   return extras;
 }
 
+/** Build (or reuse) the direct-API Gemini backend for the given key. */
+function getDirectGeminiBackend(
+  config: AssistantConfig,
+  geminiKey: string,
+): EmbeddingBackend {
+  return getCachedOrCreate(
+    "gemini",
+    config.memory.embeddings.geminiModel,
+    () =>
+      new GeminiEmbeddingBackend(
+        geminiKey,
+        config.memory.embeddings.geminiModel,
+        {
+          taskType: config.memory.embeddings.geminiTaskType,
+          dimensions: config.memory.embeddings.geminiDimensions,
+        },
+      ),
+    geminiCacheExtras(config),
+  );
+}
+
+/**
+ * Build (or reuse) the managed-proxy Gemini backend, or return undefined
+ * when the managed proxy prerequisites (platform URL + assistant API key)
+ * are not satisfied.
+ */
+async function tryGetManagedGeminiBackend(
+  config: AssistantConfig,
+): Promise<EmbeddingBackend | undefined> {
+  const proxyCtx = await resolveManagedProxyContext();
+  const meta = PLATFORM_PROVIDER_META["gemini"];
+  if (!proxyCtx.enabled || !meta?.managed || !meta.proxyPath) return undefined;
+  const managedBaseUrl = `${proxyCtx.platformBaseUrl}${meta.proxyPath}`;
+  return getCachedOrCreate(
+    "gemini",
+    config.memory.embeddings.geminiModel,
+    () =>
+      new GeminiEmbeddingBackend(
+        proxyCtx.assistantApiKey,
+        config.memory.embeddings.geminiModel,
+        {
+          taskType: config.memory.embeddings.geminiTaskType,
+          dimensions: config.memory.embeddings.geminiDimensions ?? 3072,
+          managedBaseUrl,
+        },
+      ),
+    [...geminiCacheExtras(config), "managed"],
+  );
+}
+
 export interface EmbeddingBackendSelection {
   backend: EmbeddingBackend | null;
   reason: string | null;
@@ -305,34 +355,9 @@ export async function selectEmbeddingBackend(
   // at the front of the auto chain so platform assistants use Vellum-managed
   // Gemini embeddings.
   if (requested === "auto" || requested === "gemini") {
-    const proxyCtx = await resolveManagedProxyContext();
-    if (proxyCtx.enabled) {
-      const meta = PLATFORM_PROVIDER_META["gemini"];
-      if (meta?.managed && meta.proxyPath) {
-        const managedBaseUrl = `${proxyCtx.platformBaseUrl}${meta.proxyPath}`;
-        const managedModel = config.memory.embeddings.geminiModel;
-        const managedDimensions =
-          config.memory.embeddings.geminiDimensions ?? 3072;
-        const extras = geminiCacheExtras(config);
-        return {
-          backend: getCachedOrCreate(
-            "gemini",
-            managedModel,
-            () =>
-              new GeminiEmbeddingBackend(
-                proxyCtx.assistantApiKey,
-                managedModel,
-                {
-                  taskType: config.memory.embeddings.geminiTaskType,
-                  dimensions: managedDimensions,
-                  managedBaseUrl,
-                },
-              ),
-            [...extras, "managed"],
-          ),
-          reason: null,
-        };
-      }
+    const managed = await tryGetManagedGeminiBackend(config);
+    if (managed) {
+      return { backend: managed, reason: null };
     }
   }
 
@@ -401,20 +426,7 @@ export async function selectEmbeddingBackend(
           continue;
         }
         return {
-          backend: getCachedOrCreate(
-            "gemini",
-            config.memory.embeddings.geminiModel,
-            () =>
-              new GeminiEmbeddingBackend(
-                geminiKey,
-                config.memory.embeddings.geminiModel,
-                {
-                  taskType: config.memory.embeddings.geminiTaskType,
-                  dimensions: config.memory.embeddings.geminiDimensions,
-                },
-              ),
-            geminiCacheExtras(config),
-          ),
+          backend: getDirectGeminiBackend(config, geminiKey),
           reason: null,
         };
       }
@@ -507,6 +519,21 @@ export async function embedWithBackend(
       ? await selectFallbackBackends(config, selection.backend.provider)
       : [];
 
+  // A managed-proxy Gemini primary can fail at the proxy (e.g. a stale or
+  // revoked platform credential) while a valid direct Gemini key sits in the
+  // credential store. The provider-level chain above never includes Gemini
+  // when Gemini IS the primary, so without this the direct key would never
+  // be tried and the provider would stay down for the whole process.
+  if (
+    selection.backend instanceof GeminiEmbeddingBackend &&
+    selection.backend.managed
+  ) {
+    const geminiKey = await getProviderKeyAsync("gemini");
+    if (geminiKey) {
+      fallbacks.push(getDirectGeminiBackend(config, geminiKey));
+    }
+  }
+
   // ── Compute provider-specific vector cache extras ───────────────
   const vectorExtras =
     primaryProvider === "gemini" ? geminiCacheExtras(config) : undefined;
@@ -657,49 +684,12 @@ async function selectFallbackBackends(
       case "gemini": {
         const geminiKey = await getProviderKeyAsync("gemini");
         if (geminiKey) {
-          backends.push(
-            getCachedOrCreate(
-              "gemini",
-              config.memory.embeddings.geminiModel,
-              () =>
-                new GeminiEmbeddingBackend(
-                  geminiKey,
-                  config.memory.embeddings.geminiModel,
-                  {
-                    taskType: config.memory.embeddings.geminiTaskType,
-                    dimensions: config.memory.embeddings.geminiDimensions,
-                  },
-                ),
-              geminiCacheExtras(config),
-            ),
-          );
+          backends.push(getDirectGeminiBackend(config, geminiKey));
         } else {
           // Try managed proxy Gemini as fallback when no direct key exists.
-          const proxyCtx = await resolveManagedProxyContext();
-          const meta = PLATFORM_PROVIDER_META["gemini"];
-          if (proxyCtx.enabled && meta?.managed && meta.proxyPath) {
-            const managedBaseUrl = `${proxyCtx.platformBaseUrl}${meta.proxyPath}`;
-            const managedModel = config.memory.embeddings.geminiModel;
-            const managedDimensions =
-              config.memory.embeddings.geminiDimensions ?? 3072;
-            const extras = geminiCacheExtras(config);
-            backends.push(
-              getCachedOrCreate(
-                "gemini",
-                managedModel,
-                () =>
-                  new GeminiEmbeddingBackend(
-                    proxyCtx.assistantApiKey,
-                    managedModel,
-                    {
-                      taskType: config.memory.embeddings.geminiTaskType,
-                      dimensions: managedDimensions,
-                      managedBaseUrl,
-                    },
-                  ),
-                [...extras, "managed"],
-              ),
-            );
+          const managed = await tryGetManagedGeminiBackend(config);
+          if (managed) {
+            backends.push(managed);
           } else {
             // Check managed cache variant first, then non-managed, so a warm
             // managed backend survives transient proxy-context blips.

package/src/memory/embedding-gemini.ts

@@ -36,6 +36,11 @@ export class GeminiEmbeddingBackend implements EmbeddingBackend {
     this.managedBaseUrl = options?.managedBaseUrl;
   }
 
+  /** True when requests route through the managed platform proxy. */
+  get managed(): boolean {
+    return Boolean(this.managedBaseUrl);
+  }
+
   async embed(
     inputs: EmbeddingInput[],
     options?: EmbeddingRequestOptions,