@vellumai/vellum-gateway 0.6.2 → 0.6.4

package/src/__tests__/browser-relay-websocket.test.ts

@@ -3,6 +3,7 @@ import type { GatewayConfig } from "../config.js";
 import { initSigningKey, mintToken } from "../auth/token-service.js";
 import { CURRENT_POLICY_EPOCH } from "../auth/policy.js";
 import {
+  checkBrowserRelayAuth,
   createBrowserRelayWebsocketHandler,
   getBrowserRelayWebsocketHandlers,
   isLoopbackPeer,
@@ -11,17 +12,34 @@ import {
 const TEST_SIGNING_KEY = Buffer.from("test-signing-key-at-least-32-bytes-long");
 initSigningKey(TEST_SIGNING_KEY);
 
-/** Mint a valid edge JWT for browser relay auth. */
-function mintEdgeToken(): string {
+const TEST_ACTOR_PRINCIPAL = "guardian-actor-123";
+
+/** Mint a valid actor edge JWT for browser relay auth. */
+function mintEdgeToken(actorPrincipalId: string = "test-user"): string {
   return mintToken({
     aud: "vellum-gateway",
-    sub: "actor:test-assistant:test-user",
+    sub: `actor:test-assistant:${actorPrincipalId}`,
     scope_profile: "actor_client_v1",
     policy_epoch: CURRENT_POLICY_EPOCH,
     ttlSeconds: 300,
   });
 }
 
+/**
+ * Mint a service-style browser-relay edge token (svc:browser-relay:self).
+ * This token is valid for the gateway audience but carries no actor
+ * principal in its sub claim.
+ */
+function mintServiceEdgeToken(): string {
+  return mintToken({
+    aud: "vellum-gateway",
+    sub: "svc:browser-relay:self",
+    scope_profile: "gateway_service_v1",
+    policy_epoch: CURRENT_POLICY_EPOCH,
+    ttlSeconds: 300,
+  });
+}
+
 const WS_CONNECTING = WebSocket.CONNECTING; // 0
 const WS_OPEN = WebSocket.OPEN; // 1
 const WS_CLOSED = WebSocket.CLOSED; // 3
@@ -257,6 +275,7 @@ describe("getBrowserRelayWebsocketHandlers", () => {
       config: makeConfig({
         assistantRuntimeBaseUrl: "http://runtime.internal:7821",
       }),
+      auth: { authenticated: false, authBypassed: true },
     });
 
     handlers.open(ws as never);
@@ -275,11 +294,378 @@ describe("getBrowserRelayWebsocketHandlers", () => {
     fakeUpstream.emit("message", { data: "runtime-message" });
     expect(ws.sent).toEqual(["runtime-message"]);
   });
+
+  test("open appends guardianId query param when auth context carries one", () => {
+    const ws = createFakeDownstreamWs({
+      wsType: "browser-relay",
+      config: makeConfig({
+        assistantRuntimeBaseUrl: "http://runtime.internal:7821",
+      }),
+      auth: {
+        authenticated: true,
+        authBypassed: false,
+        guardianId: TEST_ACTOR_PRINCIPAL,
+      },
+    });
+
+    handlers.open(ws as never);
+
+    const MockWS = globalThis.WebSocket as unknown as ReturnType<typeof mock>;
+    const calledUrl = (MockWS.mock.calls[0] as unknown[])[0] as string;
+    const parsed = new URL(calledUrl);
+    expect(parsed.protocol).toBe("ws:");
+    expect(parsed.host).toBe("runtime.internal:7821");
+    expect(parsed.pathname).toBe("/v1/browser-relay");
+    expect(parsed.searchParams.get("token")).toMatch(/^ey/);
+    expect(parsed.searchParams.get("guardianId")).toBe(TEST_ACTOR_PRINCIPAL);
+  });
+
+  test("open omits guardianId query param when auth context has none (auth-bypass context)", () => {
+    const ws = createFakeDownstreamWs({
+      wsType: "browser-relay",
+      config: makeConfig({
+        assistantRuntimeBaseUrl: "http://runtime.internal:7821",
+      }),
+      auth: { authenticated: true, authBypassed: false },
+    });
+
+    handlers.open(ws as never);
+
+    const MockWS = globalThis.WebSocket as unknown as ReturnType<typeof mock>;
+    const calledUrl = (MockWS.mock.calls[0] as unknown[])[0] as string;
+    const parsed = new URL(calledUrl);
+    // Auth-bypass contexts may omit guardianId; open() should not
+    // synthesize one.
+    expect(parsed.searchParams.has("guardianId")).toBe(false);
+  });
+
+  test("open appends clientInstanceId query param when auth context carries one", () => {
+    const ws = createFakeDownstreamWs({
+      wsType: "browser-relay",
+      config: makeConfig({
+        assistantRuntimeBaseUrl: "http://runtime.internal:7821",
+      }),
+      auth: {
+        authenticated: true,
+        authBypassed: false,
+        guardianId: TEST_ACTOR_PRINCIPAL,
+        clientInstanceId: "install-ABC-123",
+      },
+    });
+
+    handlers.open(ws as never);
+
+    const MockWS = globalThis.WebSocket as unknown as ReturnType<typeof mock>;
+    const calledUrl = (MockWS.mock.calls[0] as unknown[])[0] as string;
+    const parsed = new URL(calledUrl);
+    expect(parsed.pathname).toBe("/v1/browser-relay");
+    // Both guardianId and clientInstanceId should ride alongside the
+    // upstream service token, matching the runtime's accepted query
+    // param names.
+    expect(parsed.searchParams.get("guardianId")).toBe(TEST_ACTOR_PRINCIPAL);
+    expect(parsed.searchParams.get("clientInstanceId")).toBe("install-ABC-123");
+  });
+
+  test("open forwards clientInstanceId even without a guardianId", () => {
+    // Auth-bypass paths can still carry a
+    // clientInstanceId lifted from the downstream handshake. The
+    // gateway must propagate it so the runtime's multi-instance
+    // registry keys the connection correctly.
+    const ws = createFakeDownstreamWs({
+      wsType: "browser-relay",
+      config: makeConfig({
+        assistantRuntimeBaseUrl: "http://runtime.internal:7821",
+      }),
+      auth: {
+        authenticated: false,
+        authBypassed: true,
+        clientInstanceId: "install-XYZ-789",
+      },
+    });
+
+    handlers.open(ws as never);
+
+    const MockWS = globalThis.WebSocket as unknown as ReturnType<typeof mock>;
+    const calledUrl = (MockWS.mock.calls[0] as unknown[])[0] as string;
+    const parsed = new URL(calledUrl);
+    expect(parsed.searchParams.has("guardianId")).toBe(false);
+    expect(parsed.searchParams.get("clientInstanceId")).toBe("install-XYZ-789");
+  });
+
+  test("open omits clientInstanceId query param when auth context has none", () => {
+    const ws = createFakeDownstreamWs({
+      wsType: "browser-relay",
+      config: makeConfig({
+        assistantRuntimeBaseUrl: "http://runtime.internal:7821",
+      }),
+      auth: {
+        authenticated: true,
+        authBypassed: false,
+        guardianId: TEST_ACTOR_PRINCIPAL,
+      },
+    });
+
+    handlers.open(ws as never);
+
+    const MockWS = globalThis.WebSocket as unknown as ReturnType<typeof mock>;
+    const calledUrl = (MockWS.mock.calls[0] as unknown[])[0] as string;
+    const parsed = new URL(calledUrl);
+    // Older extension builds (or dev bypass paths) do not emit a
+    // clientInstanceId — the gateway must not synthesize one or send
+    // an empty string, both of which would break the runtime's
+    // legacy-key fallback semantics.
+    expect(parsed.searchParams.has("clientInstanceId")).toBe(false);
+  });
+});
+
+describe("checkBrowserRelayAuth", () => {
+  test("returns structured auth context with guardianId for actor edge tokens", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.authenticated).toBe(true);
+    expect(result.context.authBypassed).toBe(false);
+    expect(result.context.guardianId).toBe(TEST_ACTOR_PRINCIPAL);
+  });
+
+  test("rejects service-style edge tokens with no actor principal", () => {
+    const token = mintServiceEdgeToken();
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(false);
+    if (result.ok) throw new Error("expected error");
+    expect(result.response.status).toBe(401);
+  });
+
+  test("returns error response when token is missing and auth is required", () => {
+    const config = makeConfig({});
+    const req = new Request("http://localhost:7830/v1/browser-relay", {
+      headers: { upgrade: "websocket" },
+    });
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(false);
+    if (result.ok) throw new Error("expected error");
+    expect(result.response.status).toBe(401);
+  });
+
+  test("returns bypassed context when runtimeProxyRequireAuth is disabled", () => {
+    const config = makeConfig({ runtimeProxyRequireAuth: false });
+    const req = new Request("http://localhost:7830/v1/browser-relay", {
+      headers: { upgrade: "websocket" },
+    });
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.authBypassed).toBe(true);
+    expect(result.context.authenticated).toBe(false);
+    expect(result.context.guardianId).toBeUndefined();
+  });
+
+  test("lifts clientInstanceId from the query param on the downstream handshake", () => {
+    // Primary path: the Chrome extension sets clientInstanceId as a
+    // query param via `buildUrl`. The gateway must lift it off the
+    // handshake and expose it on the auth context so the open()
+    // handler can forward it to the runtime.
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}&clientInstanceId=install-QUERY`,
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.clientInstanceId).toBe("install-QUERY");
+    expect(result.context.guardianId).toBe(TEST_ACTOR_PRINCIPAL);
+  });
+
+  test("lifts clientInstanceId from the x-client-instance-id header", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      {
+        headers: {
+          upgrade: "websocket",
+          "x-client-instance-id": "install-HEADER",
+        },
+      },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.clientInstanceId).toBe("install-HEADER");
+  });
+
+  test("prefers x-client-instance-id header over query param when both are present", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}&clientInstanceId=install-QUERY`,
+      {
+        headers: {
+          upgrade: "websocket",
+          "x-client-instance-id": "install-HEADER",
+        },
+      },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    // The header form is considered the more explicit signal when
+    // both are set, matching the runtime's own precedence.
+    expect(result.context.clientInstanceId).toBe("install-HEADER");
+  });
+
+  test("treats an empty clientInstanceId query param as absent", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}&clientInstanceId=`,
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.clientInstanceId).toBeUndefined();
+  });
+
+  test("returns undefined clientInstanceId when neither form is present", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.clientInstanceId).toBeUndefined();
+  });
+
+  test("lifts clientInstanceId even on the auth-bypass path", () => {
+    // When runtime proxy auth is disabled, the gateway still needs to
+    // propagate the per-install identifier so multi-instance routing
+    // continues to work in dev-bypass environments.
+    const config = makeConfig({ runtimeProxyRequireAuth: false });
+    const req = new Request(
+      "http://localhost:7830/v1/browser-relay?clientInstanceId=install-BYPASS",
+      { headers: { upgrade: "websocket" } },
+    );
+    const url = new URL(req.url);
+
+    const result = checkBrowserRelayAuth(req, url, config);
+
+    expect(result.ok).toBe(true);
+    if (!result.ok) throw new Error("expected ok");
+    expect(result.context.authBypassed).toBe(true);
+    expect(result.context.clientInstanceId).toBe("install-BYPASS");
+  });
+});
+
+describe("createBrowserRelayWebsocketHandler guardian propagation", () => {
+  test("upgrades with guardianId populated in auth context for actor edge tokens", () => {
+    const token = mintEdgeToken(TEST_ACTOR_PRINCIPAL);
+    const config = makeConfig({});
+    const handler = createBrowserRelayWebsocketHandler(config);
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      { headers: { upgrade: "websocket" } },
+    );
+
+    let capturedData: unknown = null;
+    const fakeServer = {
+      requestIP: mock(() => ({
+        address: "127.0.0.1",
+        family: "IPv4",
+        port: 54000,
+      })),
+      upgrade: mock((_req: Request, opts?: { data?: unknown }) => {
+        capturedData = opts?.data;
+        return true;
+      }),
+    } as unknown as import("bun").Server<any>;
+
+    const res = handler(req, fakeServer);
+
+    expect(res).toBeUndefined();
+    expect(fakeServer.upgrade).toHaveBeenCalledTimes(1);
+    expect(capturedData).toMatchObject({
+      wsType: "browser-relay",
+      auth: {
+        authenticated: true,
+        authBypassed: false,
+        guardianId: TEST_ACTOR_PRINCIPAL,
+      },
+    });
+  });
+
+  test("rejects service-style edge tokens before upgrade", () => {
+    const token = mintServiceEdgeToken();
+    const config = makeConfig({});
+    const handler = createBrowserRelayWebsocketHandler(config);
+    const req = new Request(
+      `http://localhost:7830/v1/browser-relay?token=${token}`,
+      { headers: { upgrade: "websocket" } },
+    );
+
+    const fakeServer = {
+      requestIP: mock(() => ({
+        address: "127.0.0.1",
+        family: "IPv4",
+        port: 54000,
+      })),
+      upgrade: mock(() => true),
+    } as unknown as import("bun").Server<any>;
+
+    const res = handler(req, fakeServer);
+
+    expect(res).toBeInstanceOf(Response);
+    expect(res!.status).toBe(401);
+    expect(fakeServer.upgrade).not.toHaveBeenCalled();
+  });
 });
 
 describe("isLoopbackPeer", () => {
   test("uses x-forwarded-for first hop when trustProxy is enabled", () => {
-    const req = new Request("http://localhost:7830/v1/browser-relay/token", {
+    const req = new Request("http://localhost:7830/v1/browser-relay", {
       headers: { "x-forwarded-for": "203.0.113.5, 127.0.0.1" },
     });
 
@@ -295,7 +681,7 @@ describe("isLoopbackPeer", () => {
   });
 
   test("falls back to peer IP when trustProxy is disabled", () => {
-    const req = new Request("http://localhost:7830/v1/browser-relay/token", {
+    const req = new Request("http://localhost:7830/v1/browser-relay", {
       headers: { "x-forwarded-for": "203.0.113.5, 127.0.0.1" },
     });
 

package/src/__tests__/config-file-cache.test.ts

@@ -13,26 +13,26 @@ import { ConfigFileCache } from "../config-file-cache.js";
 let testBaseDir: string;
 let workspaceDir: string;
 let configPath: string;
-let savedBaseDataDir: string | undefined;
+let savedWorkspaceDir: string | undefined;
 
 function writeConfig(data: Record<string, unknown>): void {
   writeFileSync(configPath, JSON.stringify(data));
 }
 
 beforeEach(() => {
-  savedBaseDataDir = process.env.BASE_DATA_DIR;
+  savedWorkspaceDir = process.env.VELLUM_WORKSPACE_DIR;
   testBaseDir = mkdtempSync(join(tmpdir(), "config-file-cache-test-"));
   workspaceDir = join(testBaseDir, ".vellum", "workspace");
   mkdirSync(workspaceDir, { recursive: true });
   configPath = join(workspaceDir, "config.json");
-  process.env.BASE_DATA_DIR = testBaseDir;
+  process.env.VELLUM_WORKSPACE_DIR = workspaceDir;
 });
 
 afterEach(() => {
-  if (savedBaseDataDir === undefined) {
-    delete process.env.BASE_DATA_DIR;
+  if (savedWorkspaceDir === undefined) {
+    delete process.env.VELLUM_WORKSPACE_DIR;
   } else {
-    process.env.BASE_DATA_DIR = savedBaseDataDir;
+    process.env.VELLUM_WORKSPACE_DIR = savedWorkspaceDir;
   }
   rmSync(testBaseDir, { recursive: true, force: true });
 });

package/src/__tests__/config.test.ts

@@ -23,7 +23,7 @@ describe("config: hardcoded defaults", () => {
     expect(config.unmappedPolicy).toBe("reject");
     expect(config.routingEntries).toEqual([]);
     expect(config.defaultAssistantId).toBeUndefined();
-    expect(config.logFile.dir).toMatch(/\.vellum\/logs$/);
+    expect(config.logFile.dir).toMatch(/logs$/);
     expect(config.logFile.retentionDays).toBe(30);
   });
 

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

@@ -39,7 +39,7 @@ const testDir = join(
 );
 
 function metadataDir(): string {
-  return join(testDir, ".vellum", "workspace", "data", "credentials");
+  return join(testDir, "data", "credentials");
 }
 
 function writeMetadata(
@@ -106,8 +106,8 @@ function encryptEntries(
 }
 
 function writeEncryptedStore(entries: Record<string, string>): void {
-  const storePath = join(testDir, ".vellum", "protected", "keys.enc");
-  mkdirSync(join(testDir, ".vellum", "protected"), { recursive: true });
+  mkdirSync(testDir, { recursive: true });
+  const storePath = join(testDir, "keys.enc");
 
   const salt = randomBytes(16);
   const key = pbkdf2Sync(
@@ -131,17 +131,16 @@ function writeEncryptedStore(entries: Record<string, string>): void {
  * The store.key is used directly as the AES-256-GCM key (no PBKDF2).
  */
 function writeEncryptedStoreV2(entries: Record<string, string>): void {
-  const protectedDir = join(testDir, ".vellum", "protected");
-  mkdirSync(protectedDir, { recursive: true });
+  mkdirSync(testDir, { recursive: true });
 
   const storeKey = randomBytes(KEY_LENGTH);
-  writeFileSync(join(protectedDir, "store.key"), storeKey);
+  writeFileSync(join(testDir, "store.key"), storeKey);
 
   const store = {
     version: 2,
     entries: encryptEntries(entries, storeKey),
   };
-  writeFileSync(join(protectedDir, "keys.enc"), JSON.stringify(store));
+  writeFileSync(join(testDir, "keys.enc"), JSON.stringify(store));
 }
 
 // ---------------------------------------------------------------------------
@@ -149,12 +148,14 @@ function writeEncryptedStoreV2(entries: Record<string, string>): void {
 // ---------------------------------------------------------------------------
 
 beforeEach(() => {
-  process.env.BASE_DATA_DIR = testDir;
+  process.env.GATEWAY_SECURITY_DIR = testDir;
+  process.env.VELLUM_WORKSPACE_DIR = testDir;
   logCalls.length = 0;
 });
 
 afterEach(() => {
-  delete process.env.BASE_DATA_DIR;
+  delete process.env.GATEWAY_SECURITY_DIR;
+  delete process.env.VELLUM_WORKSPACE_DIR;
   try {
     rmSync(testDir, { recursive: true, force: true });
   } catch {
@@ -178,8 +179,7 @@ describe("v2 encrypted store with store.key", () => {
 
   test("returns undefined for v2 store when store.key is missing", async () => {
     // Write a v2 store but without the store.key file
-    const protectedDir = join(testDir, ".vellum", "protected");
-    mkdirSync(protectedDir, { recursive: true });
+    mkdirSync(testDir, { recursive: true });
 
     const storeKey = randomBytes(KEY_LENGTH);
     const store = {
@@ -189,7 +189,7 @@ describe("v2 encrypted store with store.key", () => {
         storeKey,
       ),
     };
-    writeFileSync(join(protectedDir, "keys.enc"), JSON.stringify(store));
+    writeFileSync(join(testDir, "keys.enc"), JSON.stringify(store));
     // Deliberately do NOT write store.key
 
     const result = await readCredential(credentialKey("test", "key"));