@vellumai/assistant 0.11.6 → 0.11.7-staging.2

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

@@ -1,5 +1,5 @@
 import { randomBytes } from "node:crypto";
-import { existsSync, mkdirSync, mkdtempSync, rmSync } from "node:fs";
+import { existsSync, mkdirSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { afterEach, beforeEach, describe, expect, test } from "bun:test";
@@ -20,231 +20,6 @@ const TEST_DIR = join(
 );
 const STORE_PATH = join(TEST_DIR, "keys.enc");
 
-describe("CredentialBroker", () => {
-  let broker: CredentialBroker;
-  let tmpDir: string;
-
-  beforeEach(() => {
-    tmpDir = mkdtempSync(join(tmpdir(), "broker-test-"));
-    _setMetadataPath(join(tmpDir, "metadata.json"));
-    broker = new CredentialBroker();
-  });
-
-  afterEach(() => {
-    _setMetadataPath(null);
-    rmSync(tmpDir, { recursive: true, force: true });
-  });
-
-  describe("authorize", () => {
-    test("denies when no credential metadata exists", () => {
-      const result = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      expect(result.authorized).toBe(false);
-      if (!result.authorized) {
-        expect(result.reason).toContain("No credential found");
-      }
-    });
-
-    test("authorizes when tool is in allowedTools", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["browser_fill_credential"],
-      });
-      const result = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      expect(result.authorized).toBe(true);
-      if (result.authorized) {
-        expect(result.token.service).toBe("github");
-        expect(result.token.field).toBe("token");
-        expect(result.token.toolName).toBe("browser_fill_credential");
-        expect(result.token.consumed).toBe(false);
-        expect(result.token.tokenId).toBeTruthy();
-      }
-    });
-
-    test("denies when tool is not in allowedTools", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["other_tool"],
-      });
-      const result = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      expect(result.authorized).toBe(false);
-      if (!result.authorized) {
-        expect(result.reason).toContain("not allowed");
-        expect(result.reason).toContain("browser_fill_credential");
-        expect(result.reason).toContain("Allowed tools: other_tool");
-      }
-    });
-
-    test("denies when allowedTools is empty (fail-closed)", () => {
-      upsertCredentialMetadata("github", "token");
-      const result = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      expect(result.authorized).toBe(false);
-      if (!result.authorized) {
-        expect(result.reason).toContain("not allowed");
-        expect(result.reason).toContain("No tools are currently allowed");
-        expect(result.reason).toContain("assistant credentials prompt");
-      }
-    });
-
-    test("issues unique token IDs", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["tool1", "tool2"],
-      });
-      const r1 = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool1",
-      });
-      const r2 = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool2",
-      });
-      expect(r1.authorized).toBe(true);
-      expect(r2.authorized).toBe(true);
-      if (r1.authorized && r2.authorized) {
-        expect(r1.token.tokenId).not.toBe(r2.token.tokenId);
-      }
-    });
-  });
-
-  describe("consume", () => {
-    test("returns storage key on first consumption", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["browser_fill_credential"],
-      });
-      const auth = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      expect(auth.authorized).toBe(true);
-      if (!auth.authorized) {
-        return;
-      }
-
-      const result = broker.consume(auth.token.tokenId);
-      expect(result.success).toBe(true);
-      expect(result.storageKey).toBe(credentialKey("github", "token"));
-    });
-
-    test("rejects double consumption", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["browser_fill_credential"],
-      });
-      const auth = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      if (!auth.authorized) {
-        return;
-      }
-
-      broker.consume(auth.token.tokenId);
-      const result = broker.consume(auth.token.tokenId);
-      expect(result.success).toBe(false);
-      expect(result.reason).toContain("already consumed");
-    });
-
-    test("rejects unknown token ID", () => {
-      const result = broker.consume("nonexistent-token");
-      expect(result.success).toBe(false);
-      expect(result.reason).toContain("not found");
-    });
-  });
-
-  describe("revoke", () => {
-    test("revokes existing token", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["browser_fill_credential"],
-      });
-      const auth = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "browser_fill_credential",
-      });
-      if (!auth.authorized) {
-        return;
-      }
-
-      expect(broker.revoke(auth.token.tokenId)).toBe(true);
-      // After revocation, consume should fail
-      const result = broker.consume(auth.token.tokenId);
-      expect(result.success).toBe(false);
-    });
-
-    test("returns false for unknown token", () => {
-      expect(broker.revoke("nonexistent")).toBe(false);
-    });
-  });
-
-  describe("revokeAll", () => {
-    test("clears all active tokens", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["tool1", "tool2"],
-      });
-      broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool1",
-      });
-      broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool2",
-      });
-      expect(broker.activeTokenCount).toBe(2);
-
-      broker.revokeAll();
-      expect(broker.activeTokenCount).toBe(0);
-    });
-  });
-
-  describe("activeTokenCount", () => {
-    test("counts only unconsumed tokens", () => {
-      upsertCredentialMetadata("github", "token", {
-        allowedTools: ["tool1", "tool2"],
-      });
-      const auth1 = broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool1",
-      });
-      broker.authorize({
-        service: "github",
-        field: "token",
-        toolName: "tool2",
-      });
-      expect(broker.activeTokenCount).toBe(2);
-
-      if (auth1.authorized) {
-        broker.consume(auth1.token.tokenId);
-      }
-      expect(broker.activeTokenCount).toBe(1);
-    });
-  });
-});
-
-// ---------------------------------------------------------------------------
-// Behavior that reads stored secret values, exercised against the encrypted
-// store backend (transient injection, multi-template serverUseById, revokeAll,
-// and canonical/legacy capability-key resolution).
-// ---------------------------------------------------------------------------
-
 describe("CredentialBroker (encrypted store backend)", () => {
   let broker: CredentialBroker;
 
@@ -267,51 +42,12 @@ describe("CredentialBroker (encrypted store backend)", () => {
   });
 
   describe("transient credentials", () => {
-    test("consume returns transient value and deletes it", () => {
-      // GIVEN a credential with metadata and an injected one-time transient value
-      upsertCredentialMetadata("svc", "key", { allowedTools: ["tool1"] });
-      broker.injectTransient("svc", "key", "one-time-secret");
-
-      // WHEN the credential is authorized and consumed
-      const auth = broker.authorize({
-        service: "svc",
-        field: "key",
-        toolName: "tool1",
-      });
-      expect(auth.authorized).toBe(true);
-      if (!auth.authorized) {
-        return;
-      }
-      const result = broker.consume(auth.token.tokenId);
-
-      // THEN the transient value is returned alongside the storage key
-      expect(result.success).toBe(true);
-      expect(result.value).toBe("one-time-secret");
-      expect(result.storageKey).toBe(credentialKey("svc", "key"));
-
-      // AND a subsequent authorize+consume no longer has the transient value
-      const auth2 = broker.authorize({
-        service: "svc",
-        field: "key",
-        toolName: "tool1",
-      });
-      expect(auth2.authorized).toBe(true);
-      if (!auth2.authorized) {
-        return;
-      }
-      const result2 = broker.consume(auth2.token.tokenId);
-      expect(result2.success).toBe(true);
-      expect(result2.value).toBeUndefined();
-    });
-
     test("browserFill uses transient value when available", async () => {
-      // GIVEN a credential with a transient value and no stored value
       upsertCredentialMetadata("github", "token", {
         allowedTools: ["browser_fill_credential"],
       });
       broker.injectTransient("github", "token", "transient-ghp-123");
 
-      // WHEN browserFill runs
       let filledValue: string | undefined;
       const result = await broker.browserFill({
         service: "github",
@@ -322,20 +58,17 @@ describe("CredentialBroker (encrypted store backend)", () => {
         },
       });
 
-      // THEN the transient value is filled
       expect(result.success).toBe(true);
       expect(filledValue).toBe("transient-ghp-123");
     });
 
     test("browserFill consumes transient value — second fill falls back to stored", async () => {
-      // GIVEN a credential with both a stored value and a transient value
       upsertCredentialMetadata("github", "token", {
         allowedTools: ["browser_fill_credential"],
       });
       await setSecureKeyAsync(credentialKey("github", "token"), "stored-value");
       broker.injectTransient("github", "token", "transient-value");
 
-      // WHEN browserFill runs twice
       let filled1: string | undefined;
       await broker.browserFill({
         service: "github",
@@ -355,19 +88,16 @@ describe("CredentialBroker (encrypted store backend)", () => {
         },
       });
 
-      // THEN the first fill uses the transient and the second falls back to stored
       expect(filled1).toBe("transient-value");
       expect(filled2).toBe("stored-value");
     });
 
     test("browserFill preserves transient value on fill failure", async () => {
-      // GIVEN a credential with a transient value
       upsertCredentialMetadata("github", "token", {
         allowedTools: ["browser_fill_credential"],
       });
       broker.injectTransient("github", "token", "transient-preserved");
 
-      // WHEN the first fill throws
       const result1 = await broker.browserFill({
         service: "github",
         field: "token",
@@ -387,19 +117,16 @@ describe("CredentialBroker (encrypted store backend)", () => {
         },
       });
 
-      // THEN the transient value survives for the second fill
       expect(result2.success).toBe(true);
       expect(filled).toBe("transient-preserved");
     });
 
     test("serverUse uses transient value when available", async () => {
-      // GIVEN a credential with a transient value
       upsertCredentialMetadata("vercel", "api_token", {
         allowedTools: ["deploy"],
       });
       broker.injectTransient("vercel", "api_token", "transient-vercel-tok");
 
-      // WHEN serverUse runs
       const result = await broker.serverUse({
         service: "vercel",
         field: "api_token",
@@ -410,19 +137,16 @@ describe("CredentialBroker (encrypted store backend)", () => {
         },
       });
 
-      // THEN the transient value is used for execution
       expect(result.success).toBe(true);
       expect(result.result).toBe("deployed");
     });
 
     test("serverUse consumes transient — subsequent call has no value without stored key", async () => {
-      // GIVEN a credential with only a transient value (no stored value)
       upsertCredentialMetadata("vercel", "api_token", {
         allowedTools: ["deploy"],
       });
       broker.injectTransient("vercel", "api_token", "transient-only");
 
-      // WHEN serverUse runs twice
       await broker.serverUse({
         service: "vercel",
         field: "api_token",
@@ -438,202 +162,102 @@ describe("CredentialBroker (encrypted store backend)", () => {
         },
       });
 
-      // THEN the second call fails because no stored value remains
       expect(result.success).toBe(false);
       expect(result.reason).toContain("no stored value");
     });
 
-    test("injectTransient replaces previous transient for same key", () => {
-      // GIVEN two transient values injected for the same credential
+    test("injectTransient replaces previous transient for same key", async () => {
       upsertCredentialMetadata("svc", "key", { allowedTools: ["t"] });
       broker.injectTransient("svc", "key", "first");
       broker.injectTransient("svc", "key", "second");
 
-      // WHEN the credential is authorized and consumed
-      const auth = broker.authorize({
+      const result = await broker.serverUse({
         service: "svc",
         field: "key",
         toolName: "t",
+        execute: async (v) => v,
       });
-      if (!auth.authorized) {
-        return;
-      }
-      const result = broker.consume(auth.token.tokenId);
 
-      // THEN the most recent transient value wins
-      expect(result.value).toBe("second");
+      expect(result.success).toBe(true);
+      expect(result.result).toBe("second");
     });
 
-    test("transient value for one credential does not affect another", () => {
-      // GIVEN a transient value injected for svcA only
+    test("transient value for one credential does not affect another", async () => {
       upsertCredentialMetadata("svcA", "key", { allowedTools: ["t"] });
       upsertCredentialMetadata("svcB", "key", { allowedTools: ["t"] });
+      await setSecureKeyAsync(credentialKey("svcB", "key"), "stored-b");
       broker.injectTransient("svcA", "key", "val-a");
 
-      // WHEN both credentials are authorized and consumed
-      const authB = broker.authorize({
+      const resultB = await broker.serverUse({
         service: "svcB",
         field: "key",
         toolName: "t",
+        execute: async (v) => v,
       });
-      if (!authB.authorized) {
-        return;
-      }
-      const resultB = broker.consume(authB.token.tokenId);
-      const authA = broker.authorize({
+      const resultA = await broker.serverUse({
         service: "svcA",
         field: "key",
         toolName: "t",
+        execute: async (v) => v,
       });
-      if (!authA.authorized) {
-        return;
-      }
-      const resultA = broker.consume(authA.token.tokenId);
 
-      // THEN only svcA carries the transient value
       expect(resultB.success).toBe(true);
-      expect(resultB.value).toBeUndefined();
-      expect(resultA.value).toBe("val-a");
-    });
-  });
-
-  describe("serverUseById edge cases", () => {
-    test("serverUseById with multiple injection templates returns all", async () => {
-      // GIVEN a credential with two injection templates and a stored value
-      const meta = upsertCredentialMetadata("multi", "api_key", {
-        allowedTools: ["proxy"],
-        injectionTemplates: [
-          {
-            hostPattern: "*.fal.ai",
-            injectionType: "header",
-            headerName: "Authorization",
-            valuePrefix: "Key ",
-          },
-          {
-            hostPattern: "gateway.fal.ai",
-            injectionType: "header",
-            headerName: "X-Fal-Key",
-          },
-        ],
-      });
-      await setSecureKeyAsync(
-        credentialKey("multi", "api_key"),
-        "multi-secret",
-      );
-
-      // WHEN serverUseById resolves the credential by id
-      const result = await broker.serverUseById({
-        credentialId: meta.credentialId,
-        requestingTool: "proxy",
-      });
-
-      // THEN both injection templates are returned without leaking the secret
-      expect(result.success).toBe(true);
-      if (!result.success) {
-        return;
-      }
-      expect(result.injectionTemplates).toHaveLength(2);
-      expect(result.injectionTemplates[0].hostPattern).toBe("*.fal.ai");
-      expect(result.injectionTemplates[1].hostPattern).toBe("gateway.fal.ai");
-      expect(JSON.stringify(result)).not.toContain("multi-secret");
-    });
-  });
-
-  describe("revokeAll", () => {
-    test("revokeAll clears all tokens and subsequent consume fails", () => {
-      // GIVEN two active tokens for a credential
-      upsertCredentialMetadata("svc", "key", { allowedTools: ["t1", "t2"] });
-      const a1 = broker.authorize({
-        service: "svc",
-        field: "key",
-        toolName: "t1",
-      });
-      const a2 = broker.authorize({
-        service: "svc",
-        field: "key",
-        toolName: "t2",
-      });
-      expect(broker.activeTokenCount).toBe(2);
-
-      // WHEN revokeAll is called
-      broker.revokeAll();
-
-      // THEN the token count is zero and consuming either token fails
-      expect(broker.activeTokenCount).toBe(0);
-      if (a1.authorized) {
-        expect(broker.consume(a1.token.tokenId).success).toBe(false);
-      }
-      if (a2.authorized) {
-        expect(broker.consume(a2.token.tokenId).success).toBe(false);
-      }
-    });
-
-    test("revokeAll on empty broker is a no-op", () => {
-      // GIVEN a broker with no active tokens
-      expect(broker.activeTokenCount).toBe(0);
-
-      // WHEN revokeAll is called
-      broker.revokeAll();
-
-      // THEN the token count remains zero
-      expect(broker.activeTokenCount).toBe(0);
+      expect(resultB.result).toBe("stored-b");
+      expect(resultA.success).toBe(true);
+      expect(resultA.result).toBe("val-a");
     });
   });
 
   describe("canonical capability key", () => {
-    test("authorize succeeds with canonical key when metadata has canonical key", () => {
-      // GIVEN metadata allowing the canonical browser-fill capability key
+    test("browserFill succeeds with canonical key when metadata has canonical key", async () => {
       upsertCredentialMetadata("github", "token", {
         allowedTools: [BROWSER_FILL_CAPABILITY],
       });
+      await setSecureKeyAsync(credentialKey("github", "token"), "ghp_secret");
 
-      // WHEN authorizing with the canonical key
-      const result = broker.authorize({
+      const result = await broker.browserFill({
         service: "github",
         field: "token",
         toolName: BROWSER_FILL_CAPABILITY,
+        fill: async () => {},
       });
 
-      // THEN authorization succeeds
-      expect(result.authorized).toBe(true);
+      expect(result.success).toBe(true);
     });
 
-    test("authorize succeeds with canonical key when metadata has legacy alias", () => {
-      // GIVEN metadata stored under the legacy browser_fill_credential alias
+    test("browserFill succeeds with canonical key when metadata has legacy alias", async () => {
       upsertCredentialMetadata("github", "token", {
         allowedTools: ["browser_fill_credential"],
       });
+      await setSecureKeyAsync(credentialKey("github", "token"), "ghp_secret");
 
-      // WHEN authorizing with the canonical key
-      const result = broker.authorize({
+      const result = await broker.browserFill({
         service: "github",
         field: "token",
         toolName: BROWSER_FILL_CAPABILITY,
+        fill: async () => {},
       });
 
-      // THEN authorization succeeds
-      expect(result.authorized).toBe(true);
+      expect(result.success).toBe(true);
     });
 
-    test("authorize succeeds with legacy alias when metadata has canonical key", () => {
-      // GIVEN metadata allowing the canonical capability key
+    test("browserFill succeeds with legacy alias when metadata has canonical key", async () => {
       upsertCredentialMetadata("github", "token", {
         allowedTools: [BROWSER_FILL_CAPABILITY],
       });
+      await setSecureKeyAsync(credentialKey("github", "token"), "ghp_secret");
 
-      // WHEN authorizing with the legacy alias
-      const result = broker.authorize({
+      const result = await broker.browserFill({
         service: "github",
         field: "token",
         toolName: "browser_fill_credential",
+        fill: async () => {},
       });
 
-      // THEN authorization succeeds
-      expect(result.authorized).toBe(true);
+      expect(result.success).toBe(true);
     });
 
     test("serverUse with canonical key works when metadata has legacy alias", async () => {
-      // GIVEN metadata stored under the legacy alias with a stored value
       upsertCredentialMetadata("vercel", "api_token", {
         allowedTools: ["browser_fill_credential"],
       });
@@ -642,7 +266,6 @@ describe("CredentialBroker (encrypted store backend)", () => {
         "vercel-tok",
       );
 
-      // WHEN serverUse runs with the canonical key
       const result = await broker.serverUse({
         service: "vercel",
         field: "api_token",
@@ -650,43 +273,43 @@ describe("CredentialBroker (encrypted store backend)", () => {
         execute: async (v) => v,
       });
 
-      // THEN execution succeeds with the stored value
       expect(result.success).toBe(true);
       expect(result.result).toBe("vercel-tok");
     });
 
-    test("non-aliased tool names are unaffected by alias resolution", () => {
-      // GIVEN metadata allowing a non-aliased custom tool
+    test("non-aliased tool names are unaffected by alias resolution", async () => {
       upsertCredentialMetadata("svc", "key", {
         allowedTools: ["custom_tool"],
       });
+      await setSecureKeyAsync(credentialKey("svc", "key"), "secret");
 
-      // WHEN authorizing with that custom tool
-      const result = broker.authorize({
+      const result = await broker.serverUse({
         service: "svc",
         field: "key",
         toolName: "custom_tool",
+        execute: async (v) => v,
       });
 
-      // THEN authorization succeeds
-      expect(result.authorized).toBe(true);
+      expect(result.success).toBe(true);
+      expect(result.result).toBe("secret");
     });
 
-    test("non-aliased tool denied when only canonical key is allowed", () => {
-      // GIVEN metadata allowing only the canonical capability key
+    test("non-aliased tool denied when only canonical key is allowed", async () => {
       upsertCredentialMetadata("svc", "key", {
         allowedTools: [BROWSER_FILL_CAPABILITY],
       });
+      await setSecureKeyAsync(credentialKey("svc", "key"), "secret");
 
-      // WHEN authorizing with an unrelated tool
-      const result = broker.authorize({
+      const result = await broker.serverUse({
         service: "svc",
         field: "key",
         toolName: "unrelated_tool",
+        execute: async () => {
+          throw new Error("should not be called");
+        },
       });
 
-      // THEN authorization is denied
-      expect(result.authorized).toBe(false);
+      expect(result.success).toBe(false);
     });
   });
 });