@vellumai/assistant 0.6.1 → 0.6.3

package/src/__tests__/credential-vault-unit.test.ts

@@ -60,11 +60,11 @@ let slackChannelConfigCalls: Array<{
 }> = [];
 
 mock.module("../oauth/manual-token-connection.js", () => ({
-  syncManualTokenConnection: async (providerKey: string) => {
+  syncManualTokenConnection: async (provider: string) => {
     const { credentialKey } = await import("../security/credential-key.js");
     const { getSecureKeyAsync } = await import("../security/secure-keys.js");
 
-    if (providerKey === "slack_channel") {
+    if (provider === "slack_channel") {
       const hasBotToken = !!(await getSecureKeyAsync(
         credentialKey("slack_channel", "bot_token"),
       ));
@@ -72,9 +72,9 @@ mock.module("../oauth/manual-token-connection.js", () => ({
         credentialKey("slack_channel", "app_token"),
       ));
       if (hasBotToken && hasAppToken) {
-        manualConnectionStore[providerKey] = "active";
+        manualConnectionStore[provider] = "active";
       } else {
-        delete manualConnectionStore[providerKey];
+        delete manualConnectionStore[provider];
       }
     }
   },

package/src/__tests__/credential-vault.test.ts

@@ -50,7 +50,15 @@ mock.module("../tools/registry.js", () => ({
 // ---------------------------------------------------------------------------
 
 let mockRefreshOAuth2Token: ReturnType<
-  typeof mock<() => Promise<{ accessToken: string; expiresIn: number }>>
+  typeof mock<
+    (
+      tokenExchangeUrl: string,
+      clientId: string,
+      refreshToken: string,
+      clientSecret?: string,
+      tokenEndpointAuthMethod?: string,
+    ) => Promise<{ accessToken: string; expiresIn: number }>
+  >
 >;
 
 mock.module("../security/oauth2.js", () => {
@@ -74,7 +82,7 @@ const mockConnections = new Map<
   string,
   {
     id: string;
-    providerKey: string;
+    provider: string;
     oauthAppId: string;
     expiresAt: number | null;
   }
@@ -83,7 +91,7 @@ const mockApps = new Map<
   string,
   {
     id: string;
-    providerKey: string;
+    provider: string;
     clientId: string;
     clientSecretCredentialPath: string;
   }
@@ -92,21 +100,22 @@ const mockProviders = new Map<
   string,
   {
     key: string;
-    tokenUrl: string;
+    tokenExchangeUrl: string;
+    refreshUrl?: string | null;
     tokenEndpointAuthMethod?: string;
   }
 >();
 
 let mockDisconnectOAuthProvider: ReturnType<
   typeof mock<
-    (providerKey: string) => Promise<"disconnected" | "not-found" | "error">
+    (provider: string) => Promise<"disconnected" | "not-found" | "error">
   >
 >;
 
 mock.module("../oauth/oauth-store.js", () => {
-  mockDisconnectOAuthProvider = mock((providerKey: string) =>
+  mockDisconnectOAuthProvider = mock((provider: string) =>
     Promise.resolve(
-      mockConnections.has(providerKey)
+      mockConnections.has(provider)
         ? ("disconnected" as const)
         : ("not-found" as const),
     ),
@@ -746,7 +755,7 @@ describe("credential_store tool", () => {
       // Simulate an active OAuth connection for this service
       mockConnections.set("google", {
         id: "conn-gmail",
-        providerKey: "google",
+        provider: "google",
         oauthAppId: "app-gmail",
         expiresAt: Date.now() + 3600_000,
       });
@@ -1303,7 +1312,7 @@ describe("withValidToken refresh deduplication", () => {
    * Helper: set up a service with an access token, refresh token, and
    * mock DB data so that token refresh can proceed through doRefresh().
    *
-   * OAuth-specific fields (tokenUrl, clientId, expiresAt) are now stored
+   * OAuth-specific fields (tokenExchangeUrl, clientId, expiresAt) are now stored
    * in the SQLite oauth-store. The mock maps simulate the DB layer.
    */
   async function setupService(
@@ -1324,17 +1333,18 @@ describe("withValidToken refresh deduplication", () => {
     );
     mockProviders.set(service, {
       key: service,
-      tokenUrl: "https://oauth.example.com/token",
+      tokenExchangeUrl: "https://oauth.example.com/token",
+      refreshUrl: null,
     });
     mockApps.set(appId, {
       id: appId,
-      providerKey: service,
+      provider: service,
       clientId: "test-client-id",
       clientSecretCredentialPath: `oauth_app/${appId}/client_secret`,
     });
     mockConnections.set(service, {
       id: connId,
-      providerKey: service,
+      provider: service,
       oauthAppId: appId,
       expiresAt: opts?.expired
         ? Date.now() - 60_000 // expired 1 minute ago
@@ -1428,7 +1438,7 @@ describe("withValidToken refresh deduplication", () => {
     let refreshCallCount = 0;
     mockRefreshOAuth2Token.mockImplementation(() => {
       refreshCallCount++;
-      // Both services use the same tokenUrl in this test, so we track by
+      // Both services use the same tokenExchangeUrl in this test, so we track by
       // call order to return the correct deferred promise.
       if (refreshCallCount === 1) return gmailPromise;
       return slackPromise;
@@ -1527,4 +1537,133 @@ describe("withValidToken refresh deduplication", () => {
     // Only one actual refresh attempt
     expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
   });
+
+  // -----------------------------------------------------------------------
+  // refreshUrl resolution — provider.refreshUrl with fallback to tokenExchangeUrl
+  // -----------------------------------------------------------------------
+  describe("refreshUrl resolution", () => {
+    test("uses provider.refreshUrl when set", async () => {
+      await setupService("google");
+      mockProviders.get("google")!.refreshUrl =
+        "https://refresh.example.com/token";
+
+      mockRefreshOAuth2Token.mockImplementation(() =>
+        Promise.resolve({
+          accessToken: "new-token-from-refresh-url",
+          expiresIn: 3600,
+        }),
+      );
+
+      const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+      const callback = async (token: string) => {
+        if (token === "old-access-token") throw err401;
+        return `result-with-${token}`;
+      };
+
+      const result = await withValidToken("google", callback);
+
+      expect(result).toBe("result-with-new-token-from-refresh-url");
+      expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+      // Assert the refresh endpoint passed in is provider.refreshUrl, not
+      // the tokenExchangeUrl fallback.
+      expect(mockRefreshOAuth2Token.mock.calls[0]?.[0]).toBe(
+        "https://refresh.example.com/token",
+      );
+    });
+
+    test("falls back to provider.tokenExchangeUrl when refreshUrl is null", async () => {
+      // setupService sets refreshUrl: null by default — this exercises the
+      // fallback path explicitly.
+      await setupService("google");
+      expect(mockProviders.get("google")!.refreshUrl).toBeNull();
+
+      mockRefreshOAuth2Token.mockImplementation(() =>
+        Promise.resolve({
+          accessToken: "new-token-from-token-exchange-url",
+          expiresIn: 3600,
+        }),
+      );
+
+      const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+      const callback = async (token: string) => {
+        if (token === "old-access-token") throw err401;
+        return `result-with-${token}`;
+      };
+
+      const result = await withValidToken("google", callback);
+
+      expect(result).toBe("result-with-new-token-from-token-exchange-url");
+      expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+      // Assert the refresh endpoint falls back to tokenExchangeUrl.
+      expect(mockRefreshOAuth2Token.mock.calls[0]?.[0]).toBe(
+        "https://oauth.example.com/token",
+      );
+    });
+
+    test("falls back to provider.tokenExchangeUrl when refreshUrl is undefined", async () => {
+      await setupService("google");
+      // Delete the refreshUrl field entirely so the property is `undefined`
+      // rather than `null`. Both representations of "not set" must produce
+      // the fallback behavior.
+      delete mockProviders.get("google")!.refreshUrl;
+      expect(mockProviders.get("google")!.refreshUrl).toBeUndefined();
+
+      mockRefreshOAuth2Token.mockImplementation(() =>
+        Promise.resolve({
+          accessToken: "new-token-from-token-exchange-url",
+          expiresIn: 3600,
+        }),
+      );
+
+      const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+      const callback = async (token: string) => {
+        if (token === "old-access-token") throw err401;
+        return `result-with-${token}`;
+      };
+
+      const result = await withValidToken("google", callback);
+
+      expect(result).toBe("result-with-new-token-from-token-exchange-url");
+      expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+      // Assert the refresh endpoint falls back to tokenExchangeUrl.
+      expect(mockRefreshOAuth2Token.mock.calls[0]?.[0]).toBe(
+        "https://oauth.example.com/token",
+      );
+    });
+
+    test("falls back to provider.tokenExchangeUrl when refreshUrl is empty string", async () => {
+      // Platform's Python `oauth_app.refresh_url or oauth_app.token_exchange_url`
+      // treats an empty string as unset. We use `||` (not `??`) so empty
+      // strings follow the same fallback path and never resolve to an empty
+      // endpoint.
+      await setupService("google");
+      mockProviders.get("google")!.refreshUrl = "";
+
+      mockRefreshOAuth2Token.mockImplementation(() =>
+        Promise.resolve({
+          accessToken: "new-token-from-token-exchange-url",
+          expiresIn: 3600,
+        }),
+      );
+
+      const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+      const callback = async (token: string) => {
+        if (token === "old-access-token") throw err401;
+        return `result-with-${token}`;
+      };
+
+      const result = await withValidToken("google", callback);
+
+      expect(result).toBe("result-with-new-token-from-token-exchange-url");
+      expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+      // Assert the refresh endpoint falls back to tokenExchangeUrl — NOT "".
+      expect(mockRefreshOAuth2Token.mock.calls[0]?.[0]).toBe(
+        "https://oauth.example.com/token",
+      );
+    });
+  });
 });

package/src/__tests__/credentials-cli.test.ts

@@ -148,9 +148,9 @@ let disconnectOAuthProviderResult: "disconnected" | "not-found" | "error" =
 
 mock.module("../oauth/oauth-store.js", () => ({
   disconnectOAuthProvider: async (
-    providerKey: string,
+    provider: string,
   ): Promise<"disconnected" | "not-found" | "error"> => {
-    disconnectOAuthProviderCalls.push(providerKey);
+    disconnectOAuthProviderCalls.push(provider);
     return disconnectOAuthProviderResult;
   },
   getConnectionByProvider: (): undefined => undefined,

package/src/__tests__/date-context.test.ts

@@ -115,7 +115,7 @@ describe("formatTurnTimestamp", () => {
       timeZone: "America/Chicago",
     });
     expect(result).toBe(
-      "2026-04-02 (Thu) 01:52:33 -05:00 (America/Chicago)",
+      "2026-04-02 (Thursday) 01:52:33 -05:00 (America/Chicago)",
     );
   });
 
@@ -124,7 +124,7 @@ describe("formatTurnTimestamp", () => {
       nowMs: THU_APR_02_0652,
       hostTimeZone: "UTC",
     });
-    expect(result).toBe("2026-04-02 (Thu) 06:52:33 +00:00 (UTC)");
+    expect(result).toBe("2026-04-02 (Thursday) 06:52:33 +00:00 (UTC)");
   });
 
   test("handles user timezone override", () => {
@@ -133,7 +133,7 @@ describe("formatTurnTimestamp", () => {
       hostTimeZone: "UTC",
       userTimeZone: "Asia/Tokyo",
     });
-    expect(result).toBe("2026-04-02 (Thu) 15:52:33 +09:00 (Asia/Tokyo)");
+    expect(result).toBe("2026-04-02 (Thursday) 15:52:33 +09:00 (Asia/Tokyo)");
   });
 
   test("handles DST correctly", () => {
@@ -144,7 +144,7 @@ describe("formatTurnTimestamp", () => {
       timeZone: "America/New_York",
     });
     expect(result).toBe(
-      "2026-07-01 (Wed) 08:00:30 -04:00 (America/New_York)",
+      "2026-07-01 (Wednesday) 08:00:30 -04:00 (America/New_York)",
     );
   });
 

package/src/__tests__/embedding-managed-proxy-selection.test.ts

@@ -0,0 +1,256 @@
+/**
+ * Tests for managed proxy Gemini embedding backend selection.
+ *
+ * Verifies that selectEmbeddingBackend correctly routes through the
+ * managed proxy when the feature flag is enabled and managed proxy
+ * prerequisites are satisfied.
+ */
+
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { credentialKey } from "../security/credential-key.js";
+
+// ---------------------------------------------------------------------------
+// Mocks — must be before importing the module under test
+// ---------------------------------------------------------------------------
+
+// Suppress logger output
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Mutable state for managed proxy context
+let mockPlatformBaseUrl = "";
+let mockAssistantApiKey: string | null = null;
+let mockProviderKeys: Record<string, string | null> = {};
+
+mock.module("../config/env.js", () => ({
+  getPlatformBaseUrl: () => mockPlatformBaseUrl,
+  getOllamaBaseUrlEnv: () => "",
+}));
+
+mock.module("../security/secure-keys.js", () => ({
+  getSecureKeyAsync: async (key: string) => {
+    if (key === credentialKey("vellum", "assistant_api_key")) {
+      return mockAssistantApiKey;
+    }
+    // Provider keys are looked up by plain name
+    return mockProviderKeys[key] ?? null;
+  },
+}));
+
+// Feature flag mock
+const mockFeatureFlags: Record<string, boolean> = {};
+
+mock.module("../config/assistant-feature-flags.js", () => ({
+  isAssistantFeatureFlagEnabled: (key: string, _config: unknown) => {
+    return mockFeatureFlags[key] ?? false;
+  },
+}));
+
+import type { AssistantConfig } from "../config/types.js";
+import {
+  clearEmbeddingBackendCache,
+  selectEmbeddingBackend,
+} from "../memory/embedding-backend.js";
+import { GeminiEmbeddingBackend } from "../memory/embedding-gemini.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+const PLATFORM_BASE = "https://platform.example.com";
+const MANAGED_API_KEY = "ast-managed-key-123";
+
+function enableManagedProxy() {
+  mockPlatformBaseUrl = PLATFORM_BASE;
+  mockAssistantApiKey = MANAGED_API_KEY;
+}
+
+function disableManagedProxy() {
+  mockPlatformBaseUrl = "";
+  mockAssistantApiKey = null;
+}
+
+function enableFlag() {
+  mockFeatureFlags["managed-gemini-embeddings-enabled"] = true;
+}
+
+function disableFlag() {
+  mockFeatureFlags["managed-gemini-embeddings-enabled"] = false;
+}
+
+function makeConfig(
+  overrides: {
+    provider?: string;
+    geminiModel?: string;
+    geminiDimensions?: number;
+  } = {},
+): AssistantConfig {
+  return {
+    memory: {
+      embeddings: {
+        provider: overrides.provider ?? "auto",
+        localModel: "Xenova/bge-small-en-v1.5",
+        openaiModel: "text-embedding-3-small",
+        geminiModel: overrides.geminiModel ?? "gemini-embedding-2-preview",
+        geminiDimensions: overrides.geminiDimensions,
+        ollamaModel: "nomic-embed-text",
+      },
+      qdrant: {
+        vectorSize: 384,
+      },
+    },
+    services: {
+      inference: { provider: "anthropic" },
+    },
+  } as unknown as AssistantConfig;
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+beforeEach(() => {
+  disableManagedProxy();
+  disableFlag();
+  mockProviderKeys = {};
+  clearEmbeddingBackendCache();
+});
+
+afterEach(() => {
+  clearEmbeddingBackendCache();
+});
+
+describe("managed proxy Gemini embedding selection", () => {
+  test("selects managed proxy Gemini when flag enabled and proxy context available", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend, reason } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    expect(backend!.provider).toBe("gemini");
+    expect(reason).toBeNull();
+  });
+
+  test("managed proxy backend uses default 3072 dimensions when geminiDimensions not set", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    // Access the private dimensions field to verify default
+    const dimensions = (backend as unknown as { dimensions: number })
+      .dimensions;
+    expect(dimensions).toBe(3072);
+  });
+
+  test("managed proxy backend uses explicit geminiDimensions when set", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ geminiDimensions: 768 });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const dimensions = (backend as unknown as { dimensions: number })
+      .dimensions;
+    expect(dimensions).toBe(768);
+  });
+
+  test("managed proxy backend uses managedBaseUrl (not direct Google API)", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toBe(`${PLATFORM_BASE}/v1/runtime-proxy/gemini`);
+  });
+
+  test("falls back to local when flag is disabled (no managed proxy)", async () => {
+    enableManagedProxy();
+    disableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    // With auto and no provider keys, falls through to local
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("falls back to local when managed proxy context unavailable", async () => {
+    disableManagedProxy();
+    enableFlag();
+    const config = makeConfig();
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("selects managed proxy when provider is explicitly gemini", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ provider: "gemini" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toContain("/v1/runtime-proxy/gemini");
+  });
+
+  test("does not use managed proxy when provider is explicitly local", async () => {
+    enableManagedProxy();
+    enableFlag();
+    const config = makeConfig({ provider: "local" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("local");
+  });
+
+  test("does not use managed proxy when provider is explicitly openai", async () => {
+    enableManagedProxy();
+    enableFlag();
+    mockProviderKeys["openai"] = "user-openai-key";
+    const config = makeConfig({ provider: "openai" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).not.toBeNull();
+    expect(backend!.provider).toBe("openai");
+  });
+
+  test("direct Gemini key still works when flag is off", async () => {
+    disableManagedProxy();
+    disableFlag();
+    mockProviderKeys["gemini"] = "user-gemini-key";
+    const config = makeConfig({ provider: "gemini" });
+
+    const { backend } = await selectEmbeddingBackend(config);
+
+    expect(backend).toBeInstanceOf(GeminiEmbeddingBackend);
+    expect(backend!.provider).toBe("gemini");
+    // Should NOT use managed proxy
+    const managedBaseUrl = (backend as unknown as { managedBaseUrl?: string })
+      .managedBaseUrl;
+    expect(managedBaseUrl).toBeUndefined();
+  });
+});