@vellumai/assistant 0.9.0 → 0.9.1-staging.1

package/src/__tests__/token-manager.test.ts

@@ -0,0 +1,519 @@
+import { randomBytes } from "node:crypto";
+import { mkdirSync, readdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  mock,
+  test,
+} from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Mock logger
+// ---------------------------------------------------------------------------
+
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// ---------------------------------------------------------------------------
+// Use encrypted backend with a temp store path
+// ---------------------------------------------------------------------------
+
+import { _resetBackend, setSecureKeyAsync } from "../security/secure-keys.js";
+import { setStorePathForTesting } from "./encrypted-store-test-helpers.js";
+
+const TEST_DIR = join(
+  tmpdir(),
+  `vellum-token-manager-test-${randomBytes(4).toString("hex")}`,
+);
+const STORE_PATH = join(TEST_DIR, "keys.enc");
+
+// ---------------------------------------------------------------------------
+// Mock OAuth2 token refresh so dedup can be observed without network I/O
+// ---------------------------------------------------------------------------
+
+let mockRefreshOAuth2Token: ReturnType<
+  typeof mock<
+    (
+      tokenExchangeUrl: string,
+      clientId: string,
+      refreshToken: string,
+      clientSecret?: string,
+      tokenEndpointAuthMethod?: string,
+    ) => Promise<{ accessToken: string; expiresIn: number }>
+  >
+>;
+
+mock.module("../security/oauth2.js", () => {
+  mockRefreshOAuth2Token = mock(() =>
+    Promise.resolve({
+      accessToken: "refreshed-access-token",
+      expiresIn: 3600,
+    }),
+  );
+  return {
+    refreshOAuth2Token: mockRefreshOAuth2Token,
+  };
+});
+
+// ---------------------------------------------------------------------------
+// Mock oauth-store — token-manager reads refresh config from SQLite
+// ---------------------------------------------------------------------------
+
+/** Mutable per-test map of provider connections for getConnectionByProvider */
+const mockConnections = new Map<
+  string,
+  {
+    id: string;
+    provider: string;
+    oauthAppId: string;
+    expiresAt: number | null;
+  }
+>();
+const mockApps = new Map<
+  string,
+  {
+    id: string;
+    provider: string;
+    clientId: string;
+    clientSecretCredentialPath: string;
+  }
+>();
+const mockProviders = new Map<
+  string,
+  {
+    key: string;
+    tokenExchangeUrl: string;
+    refreshUrl?: string | null;
+    tokenEndpointAuthMethod?: string;
+  }
+>();
+
+mock.module("../oauth/oauth-store.js", () => ({
+  getConnectionByProvider: (service: string) => mockConnections.get(service),
+  getConnection: (id: string) => {
+    for (const conn of mockConnections.values()) {
+      if (conn.id === id) return conn;
+    }
+    return undefined;
+  },
+  getApp: (id: string) => mockApps.get(id),
+  getProvider: (key: string) => mockProviders.get(key),
+  updateConnection: () => {},
+}));
+
+// ---------------------------------------------------------------------------
+// Import the modules under test
+// ---------------------------------------------------------------------------
+
+import {
+  _resetInflightRefreshes,
+  _resetRefreshBreakers,
+  withValidToken,
+} from "../security/token-manager.js";
+import { _setMetadataPath } from "../tools/credentials/metadata-store.js";
+
+describe("withValidToken refresh deduplication", () => {
+  beforeAll(() => {
+    mkdirSync(TEST_DIR, { recursive: true });
+  });
+
+  beforeEach(() => {
+    _resetBackend();
+    for (const entry of readdirSync(TEST_DIR)) {
+      rmSync(join(TEST_DIR, entry), { recursive: true, force: true });
+    }
+    setStorePathForTesting(STORE_PATH);
+    _setMetadataPath(join(TEST_DIR, "metadata.json"));
+    _resetRefreshBreakers();
+    _resetInflightRefreshes();
+    mockRefreshOAuth2Token.mockClear();
+    // Clear mock oauth-store maps
+    mockConnections.clear();
+    mockApps.clear();
+    mockProviders.clear();
+  });
+
+  afterEach(() => {
+    _setMetadataPath(null);
+    setStorePathForTesting(null);
+    _resetBackend();
+    _resetRefreshBreakers();
+    _resetInflightRefreshes();
+    mockConnections.clear();
+    mockApps.clear();
+    mockProviders.clear();
+  });
+
+  afterAll(() => {
+    rmSync(TEST_DIR, { recursive: true, force: true });
+  });
+
+  /**
+   * 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 (tokenExchangeUrl, clientId, expiresAt) are stored
+   * in the SQLite oauth-store. The mock maps simulate the DB layer.
+   */
+  async function setupService(
+    service: string,
+    opts?: { expired?: boolean; accessToken?: string },
+  ) {
+    const accessToken = opts?.accessToken ?? "old-access-token";
+
+    // Seed mock oauth-store maps so token-manager can resolve refresh config
+    const appId = `app-${service}`;
+    const connId = `conn-${service}`;
+
+    // Store access token under the oauth_connection key path that
+    // withValidToken reads (not the legacy credentialKey path).
+    await setSecureKeyAsync(
+      `oauth_connection/${connId}/access_token`,
+      accessToken,
+    );
+    mockProviders.set(service, {
+      key: service,
+      tokenExchangeUrl: "https://oauth.example.com/token",
+      refreshUrl: null,
+    });
+    mockApps.set(appId, {
+      id: appId,
+      provider: service,
+      clientId: "test-client-id",
+      clientSecretCredentialPath: `oauth_app/${appId}/client_secret`,
+    });
+    mockConnections.set(service, {
+      id: connId,
+      provider: service,
+      oauthAppId: appId,
+      expiresAt: opts?.expired
+        ? Date.now() - 60_000 // expired 1 minute ago
+        : Date.now() + 3600_000, // expires in 1 hour
+    });
+    // Store refresh token and client_secret in secure keys (token-manager reads them)
+    await setSecureKeyAsync(
+      `oauth_connection/${connId}/refresh_token`,
+      "valid-refresh-token",
+    );
+    await setSecureKeyAsync(
+      `oauth_app/${appId}/client_secret`,
+      "test-client-secret",
+    );
+  }
+
+  test("3 concurrent 401 refreshes for the same service call doRefresh exactly once", async () => {
+    await setupService("google");
+
+    let resolveRefresh!: (value: {
+      accessToken: string;
+      expiresIn: number;
+    }) => void;
+    const refreshPromise = new Promise<{
+      accessToken: string;
+      expiresIn: number;
+    }>((resolve) => {
+      resolveRefresh = resolve;
+    });
+
+    mockRefreshOAuth2Token.mockImplementation(() => refreshPromise);
+
+    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}`;
+    };
+
+    // Launch 3 concurrent withValidToken calls — all will get a non-expired
+    // token first, call the callback, get a 401, and then try to refresh.
+    const p1 = withValidToken("google", callback);
+    const p2 = withValidToken("google", callback);
+    const p3 = withValidToken("google", callback);
+
+    // Let the event loop tick so all 3 calls enter the 401 retry path
+    await new Promise((r) => setTimeout(r, 10));
+
+    // Resolve the single refresh attempt
+    resolveRefresh({ accessToken: "new-token-123", expiresIn: 3600 });
+
+    const results = await Promise.all([p1, p2, p3]);
+
+    // All 3 should succeed with the refreshed token
+    expect(results).toEqual([
+      "result-with-new-token-123",
+      "result-with-new-token-123",
+      "result-with-new-token-123",
+    ]);
+
+    // refreshOAuth2Token should have been called exactly once
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(1);
+  });
+
+  test("concurrent refreshes for different services proceed independently", async () => {
+    await setupService("google");
+    await setupService("slack");
+
+    let resolveGmail!: (value: {
+      accessToken: string;
+      expiresIn: number;
+    }) => void;
+    let resolveSlack!: (value: {
+      accessToken: string;
+      expiresIn: number;
+    }) => void;
+
+    const gmailPromise = new Promise<{
+      accessToken: string;
+      expiresIn: number;
+    }>((resolve) => {
+      resolveGmail = resolve;
+    });
+    const slackPromise = new Promise<{
+      accessToken: string;
+      expiresIn: number;
+    }>((resolve) => {
+      resolveSlack = resolve;
+    });
+
+    let refreshCallCount = 0;
+    mockRefreshOAuth2Token.mockImplementation(() => {
+      refreshCallCount++;
+      // 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;
+    });
+
+    const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+    const gmailCallback = async (token: string) => {
+      if (token === "old-access-token") throw err401;
+      return `gmail-${token}`;
+    };
+    const slackCallback = async (token: string) => {
+      if (token === "old-access-token") throw err401;
+      return `slack-${token}`;
+    };
+
+    const p1 = withValidToken("google", gmailCallback);
+    const p2 = withValidToken("slack", slackCallback);
+
+    await new Promise((r) => setTimeout(r, 10));
+
+    // Resolve both independently
+    resolveGmail({ accessToken: "gmail-new-token", expiresIn: 3600 });
+    resolveSlack({ accessToken: "slack-new-token", expiresIn: 3600 });
+
+    const [r1, r2] = await Promise.all([p1, p2]);
+
+    expect(r1).toBe("gmail-gmail-new-token");
+    expect(r2).toBe("slack-slack-new-token");
+
+    // Both services should have triggered their own refresh
+    expect(mockRefreshOAuth2Token).toHaveBeenCalledTimes(2);
+  });
+
+  test("deduplication cleans up after refresh completes, allowing subsequent refreshes", async () => {
+    await setupService("google");
+
+    let refreshCount = 0;
+    mockRefreshOAuth2Token.mockImplementation(() => {
+      refreshCount++;
+      return Promise.resolve({
+        accessToken: `token-${refreshCount}`,
+        expiresIn: 3600,
+      });
+    });
+
+    const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+    // First call triggers a refresh (old token → 401 → refresh → token-1)
+    const r1 = await withValidToken("google", async (token: string) => {
+      if (token !== "token-1") throw err401;
+      return token;
+    });
+    expect(r1).toBe("token-1");
+    expect(refreshCount).toBe(1);
+
+    // Second call also triggers a 401 to verify dedup state was cleaned up
+    // and a new refresh is allowed (not deduplicated with the first).
+    const r2 = await withValidToken("google", async (token: string) => {
+      if (token !== "token-2") throw err401;
+      return token;
+    });
+    expect(r2).toBe("token-2");
+    // Second refresh should have happened (not deduplicated with the first,
+    // since the first already completed)
+    expect(refreshCount).toBe(2);
+  });
+
+  test("deduplication propagates refresh errors to all waiting callers", async () => {
+    await setupService("google");
+
+    mockRefreshOAuth2Token.mockImplementation(() =>
+      Promise.reject(
+        Object.assign(
+          new Error("OAuth2 token refresh failed (HTTP 401: invalid_grant)"),
+        ),
+      ),
+    );
+
+    const err401 = Object.assign(new Error("Unauthorized"), { status: 401 });
+
+    const callback = async (token: string) => {
+      if (token === "old-access-token") throw err401;
+      return "should-not-reach";
+    };
+
+    // Launch 2 concurrent calls — both should fail with the same error
+    const p1 = withValidToken("google", callback);
+    const p2 = withValidToken("google", callback);
+
+    const results = await Promise.allSettled([p1, p2]);
+
+    expect(results[0].status).toBe("rejected");
+    expect(results[1].status).toBe("rejected");
+
+    // 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__/tool-executor.test.ts

@@ -587,36 +587,6 @@ describe("isSideEffectTool", () => {
     expect(isSideEffectTool("nonexistent_tool")).toBe(false);
     expect(isSideEffectTool("")).toBe(false);
   });
-
-  describe("action-aware classification for mixed-action tools", () => {
-    test("credential_store store is a side-effect", () => {
-      expect(isSideEffectTool("credential_store", { action: "store" })).toBe(
-        true,
-      );
-    });
-
-    test("credential_store delete is a side-effect", () => {
-      expect(isSideEffectTool("credential_store", { action: "delete" })).toBe(
-        true,
-      );
-    });
-
-    test("credential_store prompt is a side-effect", () => {
-      expect(isSideEffectTool("credential_store", { action: "prompt" })).toBe(
-        true,
-      );
-    });
-
-    test("credential_store list is NOT a side-effect", () => {
-      expect(isSideEffectTool("credential_store", { action: "list" })).toBe(
-        false,
-      );
-    });
-
-    test("credential_store without input is NOT a side-effect", () => {
-      expect(isSideEffectTool("credential_store")).toBe(false);
-    });
-  });
 });
 
 // ---------------------------------------------------------------------------
@@ -751,10 +721,6 @@ describe("ToolExecutor forcePromptSideEffects enforcement", () => {
       { name: "document_create", input: { title: "doc", content: "body" } },
       { name: "document_update", input: { id: "doc-1", content: "updated" } },
       { name: "document_delete", input: { surface_id: "doc-1" } },
-      {
-        name: "credential_store",
-        input: { action: "store", name: "api-key", value: "secret" },
-      },
     ];
 
     for (const { name, input } of sideEffectTools) {
@@ -817,51 +783,6 @@ describe("ToolExecutor forcePromptSideEffects enforcement", () => {
     expect(promptCalled).toBe(true);
   });
 
-  // ── Credential store action-aware (PR fix9) ──────────
-
-  test("credential_store store forces prompt under forcePromptSideEffects", async () => {
-    checkResultOverride = { decision: "allow", reason: "Matched trust rule" };
-
-    const executor = new ToolExecutor(makeTrackingPrompter());
-    const result = await executor.execute(
-      "credential_store",
-      { action: "store", name: "api-key", value: "sk-secret-123" },
-      makeContext({ forcePromptSideEffects: true }),
-    );
-
-    expect(result.isError).toBe(false);
-    expect(promptCalled).toBe(true);
-  });
-
-  test("credential_store delete forces prompt under forcePromptSideEffects", async () => {
-    checkResultOverride = { decision: "allow", reason: "Matched trust rule" };
-
-    const executor = new ToolExecutor(makeTrackingPrompter());
-    const result = await executor.execute(
-      "credential_store",
-      { action: "delete", name: "api-key" },
-      makeContext({ forcePromptSideEffects: true }),
-    );
-
-    expect(result.isError).toBe(false);
-    expect(promptCalled).toBe(true);
-  });
-
-  test("credential_store list does NOT force prompt under forcePromptSideEffects", async () => {
-    checkResultOverride = { decision: "allow", reason: "Matched trust rule" };
-
-    const executor = new ToolExecutor(makeTrackingPrompter());
-    const result = await executor.execute(
-      "credential_store",
-      { action: "list" },
-      makeContext({ forcePromptSideEffects: true }),
-    );
-
-    expect(result.isError).toBe(false);
-    // list is read-only — must NOT trigger forced prompting
-    expect(promptCalled).toBe(false);
-  });
-
   // ── Workspace mode + forcePromptSideEffects interaction ──────────
 
   test("workspace mode allow → prompt promotion still works for side-effect tools under forcePromptSideEffects", async () => {

package/src/__tests__/trusted-contact-multichannel.test.ts

@@ -266,7 +266,7 @@ for (const config of CHANNEL_CONFIGS) {
 
       const contactResult = findContactChannel({
         channelType: config.channel,
-        externalUserId: config.senderExternalUserId,
+        address: config.senderExternalUserId,
       });
 
       expect(contactResult).not.toBeNull();
@@ -288,7 +288,7 @@ for (const config of CHANNEL_CONFIGS) {
       // Should be found on this channel
       const sameChanResult = findContactChannel({
         channelType: config.channel,
-        externalUserId: config.senderExternalUserId,
+        address: config.senderExternalUserId,
       });
       expect(sameChanResult).not.toBeNull();
 
@@ -296,7 +296,7 @@ for (const config of CHANNEL_CONFIGS) {
       const otherChannel = config.channel === "telegram" ? "slack" : "telegram";
       const crossChanResult = findContactChannel({
         channelType: otherChannel,
-        externalUserId: config.senderExternalUserId,
+        address: config.senderExternalUserId,
       });
       expect(crossChanResult).toBeNull();
     });

package/src/__tests__/trusted-contact-verification.test.ts

@@ -102,7 +102,7 @@ describe("trusted contact verification → member activation", () => {
     // Verify: active member record exists
     const contactResult = findContactChannel({
       channelType: "telegram",
-      externalUserId: "requester-user-123",
+      address: "requester-user-123",
     });
 
     expect(contactResult).not.toBeNull();
@@ -236,7 +236,7 @@ describe("trusted contact verification → member activation", () => {
     // Simulate the ACL check that inbound-message-handler performs
     const contactResult = findContactChannel({
       channelType: "telegram",
-      externalUserId: "requester-user-456",
+      address: "requester-user-456",
     });
 
     expect(contactResult).not.toBeNull();
@@ -274,7 +274,7 @@ describe("trusted contact verification → member activation", () => {
     // Member should be found via contacts
     const contactResult = findContactChannel({
       channelType: "telegram",
-      externalUserId: "user-cross-test",
+      address: "user-cross-test",
     });
     expect(contactResult).not.toBeNull();
     expect(contactResult!.channel.status).toBe("active");
@@ -282,7 +282,7 @@ describe("trusted contact verification → member activation", () => {
     // Member should NOT be found for a different channel type
     const otherChannel = findContactChannel({
       channelType: "slack",
-      externalUserId: "user-cross-test",
+      address: "user-cross-test",
     });
     expect(otherChannel).toBeNull();
   });
@@ -306,7 +306,7 @@ describe("trusted contact verification → member activation", () => {
     // Verify the member is indeed revoked (ACL would reject)
     const revokedResult = findContactChannel({
       channelType: "telegram",
-      externalUserId: "user-revoked",
+      address: "user-revoked",
     });
     expect(revokedResult).not.toBeNull();
     expect(revokedResult!.channel.status).toBe("revoked");
@@ -345,7 +345,7 @@ describe("trusted contact verification → member activation", () => {
     // Verify: member is now active again
     const reactivatedResult = findContactChannel({
       channelType: "telegram",
-      externalUserId: "user-revoked",
+      address: "user-revoked",
     });
     expect(reactivatedResult).not.toBeNull();
     expect(reactivatedResult!.channel.status).toBe("active");