@vellumai/assistant 0.4.45 → 0.4.48

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

@@ -48,6 +48,7 @@ mock.module("../tools/registry.js", () => ({
 // Imports under test
 // ---------------------------------------------------------------------------
 
+import { credentialKey } from "../security/credential-key.js";
 import { getSecureKey, setSecureKey } from "../security/secure-keys.js";
 import { CredentialBroker } from "../tools/credentials/broker.js";
 import {
@@ -108,7 +109,7 @@ describe("CredentialBroker transient credentials", () => {
     const result = broker.consume(auth.token.tokenId);
     expect(result.success).toBe(true);
     expect(result.value).toBe("one-time-secret");
-    expect(result.storageKey).toBe("credential:svc:key");
+    expect(result.storageKey).toBe(credentialKey("svc", "key"));
 
     // Second authorize + consume should NOT have the transient value
     const auth2 = broker.authorize({
@@ -148,7 +149,7 @@ describe("CredentialBroker transient credentials", () => {
     upsertCredentialMetadata("github", "token", {
       allowedTools: ["browser_fill_credential"],
     });
-    setSecureKey("credential:github:token", "stored-value");
+    setSecureKey(credentialKey("github", "token"), "stored-value");
     broker.injectTransient("github", "token", "transient-value");
 
     // First fill uses transient
@@ -411,7 +412,7 @@ describe("credential_store tool — prompt action", () => {
     expect(result.content).not.toContain("prompt-secret-val");
 
     // Verify stored
-    expect(getSecureKey("credential:test-prompt:api_key")).toBe(
+    expect(getSecureKey(credentialKey("test-prompt", "api_key"))).toBe(
       "prompt-secret-val",
     );
   });
@@ -595,16 +596,18 @@ describe("credential_store tool — oauth2_connect error paths", () => {
     expect(result.content).toContain("client_id is required");
   });
 
-  test("uses stored client_id from secure storage", async () => {
-    // Store both client_id and client_secret for the service — the
-    // requiresClientSecret guardrail will short-circuit if client_secret
-    // is missing, so we need both to validate that stored client_id
-    // is resolved correctly.
+  test("uses stored client_id from metadata", async () => {
+    // Store client_id in metadata (the canonical source) and client_secret
+    // in the secure store — the requiresClientSecret guardrail will
+    // short-circuit if client_secret is missing, so we need both to
+    // validate that stored client_id is resolved correctly.
+    upsertCredentialMetadata("integration:gmail", "access_token", {
+      oauth2ClientId: "stored-client-id-123",
+    });
     setSecureKey(
-      "credential:integration:gmail:client_id",
-      "stored-client-id-123",
+      credentialKey("integration:gmail", "client_secret"),
+      "test-secret",
     );
-    setSecureKey("credential:integration:gmail:client_secret", "test-secret");
 
     const result = await credentialStoreTool.execute(
       {
@@ -624,12 +627,11 @@ describe("credential_store tool — oauth2_connect error paths", () => {
   });
 
   test("rejects when client_secret is missing for service that requires it", async () => {
-    // Store only client_id — client_secret is intentionally absent to
-    // validate the requiresClientSecret guardrail.
-    setSecureKey(
-      "credential:integration:gmail:client_id",
-      "stored-client-id-456",
-    );
+    // Store only client_id in metadata — client_secret is intentionally
+    // absent to validate the requiresClientSecret guardrail.
+    upsertCredentialMetadata("integration:gmail", "access_token", {
+      oauth2ClientId: "stored-client-id-456",
+    });
 
     const result = await credentialStoreTool.execute(
       {
@@ -786,7 +788,7 @@ describe("credential_store tool — store validation edge cases", () => {
     );
 
     // Verify stored
-    expect(getSecureKey("credential:del-test:key")).toBe("secret");
+    expect(getSecureKey(credentialKey("del-test", "key"))).toBe("secret");
     const { getCredentialMetadata } =
       await import("../tools/credentials/metadata-store.js");
     expect(getCredentialMetadata("del-test", "key")).toBeDefined();
@@ -803,7 +805,7 @@ describe("credential_store tool — store validation edge cases", () => {
     expect(result.isError).toBe(false);
 
     // Both should be gone
-    expect(getSecureKey("credential:del-test:key")).toBeUndefined();
+    expect(getSecureKey(credentialKey("del-test", "key"))).toBeUndefined();
     expect(getCredentialMetadata("del-test", "key")).toBeUndefined();
   });
 });
@@ -892,7 +894,7 @@ describe("CredentialBroker — serverUseById edge cases", () => {
         },
       ],
     });
-    setSecureKey("credential:multi:api_key", "multi-secret");
+    setSecureKey(credentialKey("multi", "api_key"), "multi-secret");
 
     const result = broker.serverUseById({
       credentialId: meta.credentialId,

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

@@ -45,20 +45,47 @@ mock.module("../tools/registry.js", () => ({
   registerTool: () => {},
 }));
 
+// ---------------------------------------------------------------------------
+// Mock OAuth2 token refresh for token-manager deduplication tests
+// ---------------------------------------------------------------------------
+
+let mockRefreshOAuth2Token: ReturnType<
+  typeof mock<() => Promise<{ accessToken: string; expiresIn: number }>>
+>;
+
+mock.module("../security/oauth2.js", () => {
+  mockRefreshOAuth2Token = mock(() =>
+    Promise.resolve({
+      accessToken: "refreshed-access-token",
+      expiresIn: 3600,
+    }),
+  );
+  return {
+    refreshOAuth2Token: mockRefreshOAuth2Token,
+  };
+});
+
 // ---------------------------------------------------------------------------
 // Import the module under test
 // ---------------------------------------------------------------------------
 
 // getCredentialValue is no longer exported (sealed in PR 17) — use getSecureKey directly
 
+import { credentialKey } from "../security/credential-key.js";
 import {
   deleteSecureKey,
   getSecureKey,
   setSecureKey,
 } from "../security/secure-keys.js";
+import {
+  _resetInflightRefreshes,
+  _resetRefreshBreakers,
+  withValidToken,
+} from "../security/token-manager.js";
 import {
   _setMetadataPath,
   getCredentialMetadata,
+  upsertCredentialMetadata,
 } from "../tools/credentials/metadata-store.js";
 import { credentialStoreTool } from "../tools/credentials/vault.js";
 import type { ToolContext } from "../tools/types.js";
@@ -120,7 +147,7 @@ async function executeVault(
         };
       }
 
-      const key = `credential:${service}:${field}`;
+      const key = credentialKey(service, field);
       const ok = setSecureKey(key, value);
       if (!ok) {
         return { content: "Error: failed to store credential", isError: true };
@@ -151,7 +178,7 @@ async function executeVault(
         };
       }
 
-      const key = `credential:${service}:${field}`;
+      const key = credentialKey(service, field);
       const result = deleteSecureKey(key);
       if (result !== "deleted") {
         return {
@@ -553,7 +580,7 @@ describe("credential_store tool", () => {
 
       // Delete the secret directly without going through the tool (simulates
       // a divergence where metadata write failed after secret deletion)
-      deleteSecureKey("credential:svc-a:key");
+      deleteSecureKey(credentialKey("svc-a", "key"));
 
       const result = await credentialStoreTool.execute(
         { action: "list" },
@@ -596,7 +623,7 @@ describe("credential_store tool", () => {
   // -----------------------------------------------------------------------
   describe("delete action", () => {
     test("deletes a stored credential", async () => {
-      setSecureKey("credential:gmail:password", "secret");
+      setSecureKey(credentialKey("gmail", "password"), "secret");
 
       const result = await executeVault({
         action: "delete",
@@ -607,7 +634,7 @@ describe("credential_store tool", () => {
       expect(result.content).toBe("Deleted credential for gmail/password.");
 
       // Verify it's actually gone
-      expect(getSecureKey("credential:gmail:password")).toBeUndefined();
+      expect(getSecureKey(credentialKey("gmail", "password"))).toBeUndefined();
     });
 
     test("returns error for non-existent credential", async () => {
@@ -644,12 +671,14 @@ describe("credential_store tool", () => {
   // -----------------------------------------------------------------------
   describe("credential value access", () => {
     test("credential values are stored via secure keys", () => {
-      setSecureKey("credential:github:token", "ghp_abc123");
-      expect(getSecureKey("credential:github:token")).toBe("ghp_abc123");
+      setSecureKey(credentialKey("github", "token"), "ghp_abc123");
+      expect(getSecureKey(credentialKey("github", "token"))).toBe("ghp_abc123");
     });
 
     test("returns undefined for non-existent credential", () => {
-      expect(getSecureKey("credential:nonexistent:field")).toBeUndefined();
+      expect(
+        getSecureKey(credentialKey("nonexistent", "field")),
+      ).toBeUndefined();
     });
   });
 
@@ -1094,8 +1123,12 @@ describe("credential_store tool", () => {
         value: "github-pass",
       });
 
-      expect(getSecureKey("credential:gmail:password")).toBe("gmail-pass");
-      expect(getSecureKey("credential:github:password")).toBe("github-pass");
+      expect(getSecureKey(credentialKey("gmail", "password"))).toBe(
+        "gmail-pass",
+      );
+      expect(getSecureKey(credentialKey("github", "password"))).toBe(
+        "github-pass",
+      );
     });
 
     test("same service with different fields do not collide", async () => {
@@ -1112,10 +1145,249 @@ describe("credential_store tool", () => {
         value: "backup@example.com",
       });
 
-      expect(getSecureKey("credential:gmail:password")).toBe("pass123");
-      expect(getSecureKey("credential:gmail:recovery_email")).toBe(
+      expect(getSecureKey(credentialKey("gmail", "password"))).toBe("pass123");
+      expect(getSecureKey(credentialKey("gmail", "recovery_email"))).toBe(
         "backup@example.com",
       );
     });
   });
 });
+
+// ---------------------------------------------------------------------------
+// Token refresh deduplication tests
+// ---------------------------------------------------------------------------
+
+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 });
+    }
+    _setStorePath(STORE_PATH);
+    _setMetadataPath(join(TEST_DIR, "metadata.json"));
+    _resetRefreshBreakers();
+    _resetInflightRefreshes();
+    mockRefreshOAuth2Token.mockClear();
+  });
+
+  afterEach(() => {
+    _setMetadataPath(null);
+    _setStorePath(null);
+    _resetBackend();
+    _resetRefreshBreakers();
+    _resetInflightRefreshes();
+  });
+
+  afterAll(() => {
+    rmSync(TEST_DIR, { recursive: true, force: true });
+  });
+
+  /**
+   * Helper: set up a service with an access token, refresh token, and OAuth2
+   * metadata so that token refresh can proceed through doRefresh().
+   */
+  function setupService(
+    service: string,
+    opts?: { expired?: boolean; accessToken?: string },
+  ) {
+    const accessToken = opts?.accessToken ?? "old-access-token";
+    setSecureKey(credentialKey(service, "access_token"), accessToken);
+    setSecureKey(
+      credentialKey(service, "refresh_token"),
+      "valid-refresh-token",
+    );
+    upsertCredentialMetadata(service, "access_token", {
+      oauth2TokenUrl: "https://oauth.example.com/token",
+      oauth2ClientId: "test-client-id",
+      ...(opts?.expired
+        ? { expiresAt: Date.now() - 60_000 } // expired 1 minute ago
+        : { expiresAt: Date.now() + 3600_000 }), // expires in 1 hour
+    });
+  }
+
+  test("3 concurrent 401 refreshes for the same service call doRefresh exactly once", async () => {
+    setupService("integration:gmail");
+
+    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("integration:gmail", callback);
+    const p2 = withValidToken("integration:gmail", callback);
+    const p3 = withValidToken("integration:gmail", 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 () => {
+    setupService("integration:gmail");
+    setupService("integration: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 tokenUrl 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("integration:gmail", gmailCallback);
+    const p2 = withValidToken("integration: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 () => {
+    setupService("integration:gmail");
+
+    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
+    const r1 = await withValidToken(
+      "integration:gmail",
+      async (token: string) => {
+        if (token === "old-access-token") throw err401;
+        return token;
+      },
+    );
+    expect(r1).toBe("token-1");
+    expect(refreshCount).toBe(1);
+
+    // Set up so the next call will also get a 401 (token-1 stored from first refresh)
+    const r2 = await withValidToken(
+      "integration:gmail",
+      async (token: string) => {
+        if (token === "token-1") 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 () => {
+    setupService("integration:gmail");
+
+    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("integration:gmail", callback);
+    const p2 = withValidToken("integration:gmail", 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);
+  });
+});

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

@@ -2,6 +2,7 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 
 import { Command } from "commander";
 
+import { credentialKey } from "../security/credential-key.js";
 import type { CredentialMetadata } from "../tools/credentials/metadata-store.js";
 
 // ---------------------------------------------------------------------------
@@ -238,7 +239,7 @@ function seedCredential(
     ...extra,
   };
   metadataStore.push(record);
-  secureKeyStore.set(`credential:${service}:${field}`, secret);
+  secureKeyStore.set(credentialKey(service, field), secret);
   return record;
 }
 
@@ -437,7 +438,7 @@ describe("assistant credentials CLI", () => {
       expect(parsed.credentialId).toBeTruthy();
 
       // Verify secret stored in mock map
-      expect(secureKeyStore.get("credential:twilio:account_sid")).toBe(
+      expect(secureKeyStore.get(credentialKey("twilio", "account_sid"))).toBe(
         "AC1234567890",
       );
 
@@ -546,7 +547,7 @@ describe("assistant credentials CLI", () => {
       expect(meta2!.updatedAt).toBeGreaterThan(firstUpdatedAt);
 
       // Verify secret is overwritten
-      expect(secureKeyStore.get("credential:twilio:account_sid")).toBe(
+      expect(secureKeyStore.get(credentialKey("twilio", "account_sid"))).toBe(
         "new_value",
       );
     });
@@ -568,7 +569,9 @@ describe("assistant credentials CLI", () => {
       expect(parsed.field).toBe("auth_token");
 
       // Verify both removed
-      expect(secureKeyStore.has("credential:twilio:auth_token")).toBe(false);
+      expect(secureKeyStore.has(credentialKey("twilio", "auth_token"))).toBe(
+        false,
+      );
       expect(
         metadataStore.find(
           (m) => m.service === "twilio" && m.field === "auth_token",
@@ -833,8 +836,10 @@ describe("assistant credentials CLI", () => {
       expect(parsed.value).toBe("instance_secret_abc123");
 
       // Verify the correct key was looked up in the secure store
-      expect(secureKeyStore.has("credential:twilio:auth_token")).toBe(true);
-      expect(secureKeyStore.get("credential:twilio:auth_token")).toBe(
+      expect(secureKeyStore.has(credentialKey("twilio", "auth_token"))).toBe(
+        true,
+      );
+      expect(secureKeyStore.get(credentialKey("twilio", "auth_token"))).toBe(
         "instance_secret_abc123",
       );
     });

package/src/__tests__/gateway-only-enforcement.test.ts

@@ -112,9 +112,11 @@ mock.module("../calls/twilio-provider.js", () => ({
   },
 }));
 
+import { credentialKey } from "../security/credential-key.js";
+
 const secureKeyStore: Record<string, string | undefined> = {
-  "credential:twilio:account_sid": "AC_test",
-  "credential:twilio:auth_token": "test_token",
+  [credentialKey("twilio", "account_sid")]: "AC_test",
+  [credentialKey("twilio", "auth_token")]: "test_token",
 };
 
 mock.module("../security/secure-keys.js", () => ({