@vellumai/assistant 0.12.2-staging.2 → 0.12.2-staging.3

package/src/mcp/__tests__/credential-target.test.ts

@@ -0,0 +1,59 @@
+import { createHash } from "node:crypto";
+import { describe, expect, test } from "bun:test";
+
+import {
+  mcpOAuthCredentialKey,
+  pluginMcpOAuthCredentialPrefix,
+  pluginMcpOAuthCredentialTarget,
+  workspaceMcpOAuthCredentialTarget,
+} from "../credential-target.js";
+
+describe("MCP OAuth credential identity", () => {
+  test("keeps workspace credential keys unchanged", () => {
+    expect(
+      mcpOAuthCredentialKey(
+        workspaceMcpOAuthCredentialTarget("server-1"),
+        "tokens",
+      ),
+    ).toBe("mcp:server-1:tokens");
+  });
+
+  test("isolates plugin credentials by owner, original key, transport, and URL", () => {
+    const target = pluginMcpOAuthCredentialTarget("plugin/name", "server:key", {
+      type: "streamable-http",
+      url: "https://example.com/mcp?b=2&a=1#ignored",
+    });
+    const digest = createHash("sha256")
+      .update("streamable-http\nhttps://example.com/mcp?b=2&a=1")
+      .digest("hex");
+
+    expect(mcpOAuthCredentialKey(target, "tokens")).toBe(
+      `mcp-plugin/v1/cGx1Z2luL25hbWU/c2VydmVyOmtleQ/${digest}/tokens`,
+    );
+  });
+
+  test("keeps query order and path changes isolated while ignoring fragments", () => {
+    const key = (url: string) =>
+      mcpOAuthCredentialKey(
+        pluginMcpOAuthCredentialTarget("plugin", "server", {
+          type: "sse",
+          url,
+        }),
+        "tokens",
+      );
+
+    expect(key("https://example.com/a?b=2&a=1#one")).toBe(
+      key("https://example.com/a?b=2&a=1#two"),
+    );
+    expect(key("https://example.com/a?b=2&a=1")).not.toBe(
+      key("https://example.com/a?a=1&b=2"),
+    );
+    expect(key("https://example.com/a")).not.toBe(key("https://example.com/b"));
+  });
+
+  test("builds a slash-delimited whole-plugin cleanup prefix", () => {
+    expect(pluginMcpOAuthCredentialPrefix("plugin/name")).toBe(
+      "mcp-plugin/v1/cGx1Z2luL25hbWU/",
+    );
+  });
+});

package/src/mcp/__tests__/effective-config.test.ts

@@ -24,6 +24,7 @@ import { getWorkspacePluginsDir } from "../../util/platform.js";
 import {
   buildEffectiveMcpConfig,
   pluginMcpServersChangedSinceLastBuild,
+  readEffectiveMcpConfig,
   resetEffectiveMcpConfigForTests,
 } from "../effective-config.js";
 
@@ -182,6 +183,16 @@ describe("pluginMcpServersChangedSinceLastBuild", () => {
     expect(pluginMcpServersChangedSinceLastBuild()).toBe(true);
   });
 
+  test("a read-only resolution does not record the changed plugin set as applied", () => {
+    buildEffectiveMcpConfig(workspaceConfig({}));
+    writePlugin("unabyss", UNABYSS);
+
+    expect(
+      readEffectiveMcpConfig(workspaceConfig({})).servers.unabyss,
+    ).toBeDefined();
+    expect(pluginMcpServersChangedSinceLastBuild()).toBe(true);
+  });
+
   test("true once a plugin is uninstalled", () => {
     writePlugin("unabyss", UNABYSS);
     buildEffectiveMcpConfig(workspaceConfig({}));

package/src/mcp/__tests__/manager-state.test.ts

@@ -0,0 +1,128 @@
+import { beforeEach, describe, expect, jest, mock, test } from "bun:test";
+
+let completeConnect: (() => void) | undefined;
+let completeToolDiscovery: (() => void) | undefined;
+let connected = true;
+let lastError: Error | null = null;
+let waitForConnect = false;
+let waitForToolDiscovery = false;
+const unexpectedCloseHandlers = new Map<string, () => void>();
+
+mock.module("../client.js", () => ({
+  McpClient: class {
+    constructor(
+      serverId: string,
+      _source: string,
+      onUnexpectedClose: () => void,
+    ) {
+      unexpectedCloseHandlers.set(serverId, onUnexpectedClose);
+    }
+    get isConnected() {
+      return connected;
+    }
+    get lastError() {
+      return lastError;
+    }
+    async connect() {
+      if (waitForConnect) {
+        await new Promise<void>((resolve) => {
+          completeConnect = resolve;
+        });
+      }
+    }
+    async listTools() {
+      if (waitForToolDiscovery) {
+        await new Promise<void>((resolve) => {
+          completeToolDiscovery = resolve;
+        });
+      }
+      return [];
+    }
+    async disconnect() {
+      connected = false;
+    }
+  },
+}));
+
+const { McpServerManager } = await import("../manager.js");
+const config = {
+  servers: {
+    example: {
+      source: "workspace" as const,
+      transport: {
+        type: "streamable-http" as const,
+        url: "https://example.com/mcp",
+      },
+    },
+  },
+};
+
+describe("MCP manager runtime state", () => {
+  beforeEach(() => {
+    connected = true;
+    lastError = null;
+    waitForConnect = false;
+    waitForToolDiscovery = false;
+    completeConnect = undefined;
+    completeToolDiscovery = undefined;
+    unexpectedCloseHandlers.clear();
+  });
+
+  test("retains connecting state until initialization and discovery finish", async () => {
+    waitForConnect = true;
+    const manager = new McpServerManager();
+    const starting = manager.start(config);
+    expect(manager.getServerState("example", "workspace")).toBe("connecting");
+    completeConnect!();
+    await starting;
+    expect(manager.getServerState("example", "workspace")).toBe("connected");
+    expect(manager.getServerState("example", "plugin")).toBeUndefined();
+    await manager.stop();
+    expect(manager.getServerState("example", "workspace")).toBeUndefined();
+  });
+
+  test("retains authentication failures without reporting a live client", async () => {
+    connected = false;
+    const manager = new McpServerManager();
+    await manager.start(config);
+    expect(manager.getServerState("example", "workspace")).toBe("needs-auth");
+    expect(manager.getClient("example")).toBeUndefined();
+  });
+
+  test("keeps transport failure distinct from required authentication", async () => {
+    connected = false;
+    lastError = new Error("Connection refused");
+    const manager = new McpServerManager();
+    await manager.start(config);
+    expect(manager.getServerState("example", "workspace")).toBe("error");
+  });
+
+  test("records and reports an unexpected close", async () => {
+    const onUnexpectedClose = jest.fn();
+    const manager = new McpServerManager(onUnexpectedClose);
+    await manager.start(config);
+
+    connected = false;
+    unexpectedCloseHandlers.get("example")!();
+
+    expect(manager.getServerState("example", "workspace")).toBe("error");
+    expect(manager.getClient("example")).toBeDefined();
+    expect(onUnexpectedClose).toHaveBeenCalledWith();
+  });
+
+  test("does not restore connected state when the server closes during discovery", async () => {
+    waitForToolDiscovery = true;
+    const manager = new McpServerManager();
+    const starting = manager.start(config);
+    while (!completeToolDiscovery) {
+      await Promise.resolve();
+    }
+
+    connected = false;
+    unexpectedCloseHandlers.get("example")!();
+    completeToolDiscovery!();
+    await starting;
+
+    expect(manager.getServerState("example", "workspace")).toBe("error");
+  });
+});

package/src/mcp/__tests__/mcp-auth-orchestrator.test.ts

@@ -35,24 +35,34 @@ mock.module("../mcp-oauth-provider.js", () => ({
 }));
 
 const mockSetMcpAuthPending = mock(
-  (_serverId: string, _authUrl: string, _attemptId: string) => {},
+  (
+    _serverId: string,
+    _authUrl: string,
+    _attemptId: string,
+    _targetKey: string,
+  ) => {},
 );
 // Default behavior: pretend the attempt still owns the slot (return true),
 // so completion writes are applied unless a test overrides this.
 const mockSetMcpAuthComplete = mock(
-  (_serverId: string, _attemptId: string): boolean => true,
+  (_serverId: string, _attemptId: string, _targetKey: string): boolean => true,
 );
 const mockSetMcpAuthError = mock(
-  (_serverId: string, _error: string, _attemptId: string): boolean => true,
+  (
+    _serverId: string,
+    _error: string,
+    _attemptId: string,
+    _targetKey: string,
+  ): boolean => true,
 );
 
 mock.module("../mcp-auth-state.js", () => ({
   setMcpAuthPending: (...args: unknown[]) =>
-    mockSetMcpAuthPending(...(args as [string, string, string])),
+    mockSetMcpAuthPending(...(args as [string, string, string, string])),
   setMcpAuthComplete: (...args: unknown[]) =>
-    mockSetMcpAuthComplete(...(args as [string, string])),
+    mockSetMcpAuthComplete(...(args as [string, string, string])),
   setMcpAuthError: (...args: unknown[]) =>
-    mockSetMcpAuthError(...(args as [string, string, string])),
+    mockSetMcpAuthError(...(args as [string, string, string, string])),
 }));
 
 const mockReloadMcpServers = mock(async () => ({
@@ -167,6 +177,7 @@ describe("orchestrateMcpOAuthConnect", () => {
       "test-server",
       "https://auth.example.com/oauth",
       expect.any(String) as unknown as string, // attemptId UUID
+      "mcp:test-server:tokens",
     ]);
     // Sanity-check the attemptId looks UUID-shaped
     expect(mockSetMcpAuthPending.mock.calls[0][2]).toMatch(
@@ -195,6 +206,7 @@ describe("orchestrateMcpOAuthConnect", () => {
     expect(mockSetMcpAuthComplete).toHaveBeenCalledWith(
       "test-server",
       expect.any(String) as unknown as string,
+      "mcp:test-server:tokens",
     );
     // Daemon-side reload should be triggered after a successful completion.
     expect(mockReloadMcpServers).toHaveBeenCalled();
@@ -218,6 +230,7 @@ describe("orchestrateMcpOAuthConnect", () => {
       "test-server",
       "exchange failed",
       expect.any(String) as unknown as string,
+      "mcp:test-server:tokens",
     );
     expect(mockSetMcpAuthComplete).not.toHaveBeenCalled();
   });
@@ -236,6 +249,7 @@ describe("orchestrateMcpOAuthConnect", () => {
       "test-server",
       "MCP OAuth callback timed out",
       expect.any(String) as unknown as string,
+      "mcp:test-server:tokens",
     );
     expect(mockSetMcpAuthComplete).not.toHaveBeenCalled();
   });

package/src/mcp/__tests__/mcp-auth-state-target.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, test } from "bun:test";
+
+import {
+  getMcpAuthState,
+  setMcpAuthComplete,
+  setMcpAuthError,
+  setMcpAuthPending,
+} from "../mcp-auth-state.js";
+
+describe("MCP auth state target identity", () => {
+  test("does not return state for a different endpoint target", () => {
+    const serverId = "target-filter-server";
+    setMcpAuthPending(
+      serverId,
+      "https://auth.example.com",
+      "attempt-a",
+      "key-a",
+    );
+
+    expect(getMcpAuthState(serverId, "key-b")).toBeNull();
+    expect(getMcpAuthState(serverId, "key-a")?.status).toBe("pending");
+  });
+
+  test("a stale endpoint tail cannot overwrite a newer target", () => {
+    const serverId = "target-supersede-server";
+    setMcpAuthPending(
+      serverId,
+      "https://auth-a.example.com",
+      "attempt-a",
+      "key-a",
+    );
+    setMcpAuthPending(
+      serverId,
+      "https://auth-b.example.com",
+      "attempt-b",
+      "key-b",
+    );
+
+    expect(setMcpAuthComplete(serverId, "attempt-a", "key-a")).toBe(false);
+    expect(
+      setMcpAuthError(serverId, "stale failure", "attempt-a", "key-a"),
+    ).toBe(false);
+    expect(getMcpAuthState(serverId, "key-b")).toMatchObject({
+      status: "pending",
+      attemptId: "attempt-b",
+      targetKey: "key-b",
+    });
+  });
+});

package/src/mcp/__tests__/plugin-mcp-oauth-cleanup.test.ts

@@ -0,0 +1,107 @@
+import { beforeEach, describe, expect, jest, mock, test } from "bun:test";
+
+const listSecureKeysAsync = jest.fn(async () => ({
+  accounts: [] as string[],
+  unreachable: false,
+}));
+const deleteSecureKeyAsync = jest.fn(
+  async (_key: string): Promise<"deleted" | "not-found" | "error"> => "deleted",
+);
+
+mock.module("../../security/secure-keys.js", () => ({
+  getSecureKeyAsync: jest.fn(async () => null),
+  setSecureKeyAsync: jest.fn(async () => true),
+  deleteSecureKeyAsync,
+  listSecureKeysAsync,
+}));
+
+const { deletePluginMcpOAuthCredentials } =
+  await import("../mcp-oauth-provider.js");
+
+describe("plugin MCP OAuth credential cleanup", () => {
+  beforeEach(() => {
+    listSecureKeysAsync.mockReset();
+    listSecureKeysAsync.mockImplementation(async () => ({
+      accounts: [],
+      unreachable: false,
+    }));
+    deleteSecureKeyAsync.mockReset();
+    deleteSecureKeyAsync.mockImplementation(async () => "deleted" as const);
+  });
+
+  test("deletes only the exact plugin owner prefix", async () => {
+    listSecureKeysAsync.mockImplementationOnce(async () => ({
+      accounts: [
+        "mcp-plugin/v1/cGx1Z2lu/client/digest/tokens",
+        "mcp-plugin/v1/cGx1Z2lu/client/digest/discovery",
+        "mcp-plugin/v1/cGx1Z2luMg/client/digest/tokens",
+        "mcp:plugin:tokens",
+      ],
+      unreachable: false,
+    }));
+
+    const result = await deletePluginMcpOAuthCredentials("plugin");
+
+    expect(deleteSecureKeyAsync.mock.calls.map(([key]) => key)).toEqual([
+      "mcp-plugin/v1/cGx1Z2lu/client/digest/tokens",
+      "mcp-plugin/v1/cGx1Z2lu/client/digest/discovery",
+    ]);
+    expect(result).toEqual({
+      ok: true,
+      unreachable: false,
+      matchedKeys: [
+        "mcp-plugin/v1/cGx1Z2lu/client/digest/tokens",
+        "mcp-plugin/v1/cGx1Z2lu/client/digest/discovery",
+      ],
+      failedKeys: [],
+    });
+  });
+
+  test("does not attempt deletion when credential storage is unreachable", async () => {
+    listSecureKeysAsync.mockImplementationOnce(async () => ({
+      accounts: [],
+      unreachable: true,
+    }));
+
+    const result = await deletePluginMcpOAuthCredentials("plugin");
+
+    expect(deleteSecureKeyAsync).not.toHaveBeenCalled();
+    expect(result).toEqual({
+      ok: false,
+      unreachable: true,
+      matchedKeys: [],
+      failedKeys: [],
+    });
+  });
+
+  test("reports any matched key that could not be deleted", async () => {
+    const key = "mcp-plugin/v1/cGx1Z2lu/client/digest/tokens";
+    listSecureKeysAsync.mockImplementationOnce(async () => ({
+      accounts: [key],
+      unreachable: false,
+    }));
+    deleteSecureKeyAsync.mockImplementationOnce(async () => "error" as const);
+
+    const result = await deletePluginMcpOAuthCredentials("plugin");
+
+    expect(result.ok).toBe(false);
+    expect(result.failedKeys).toEqual([key]);
+  });
+
+  test("reports a rejected matched-key deletion as a failure", async () => {
+    const key = "mcp-plugin/v1/cGx1Z2lu/client/digest/tokens";
+    listSecureKeysAsync.mockImplementationOnce(async () => ({
+      accounts: [key],
+      unreachable: false,
+    }));
+    deleteSecureKeyAsync.mockImplementationOnce(async () => {
+      throw new Error("backend failed");
+    });
+
+    const result = await deletePluginMcpOAuthCredentials("plugin");
+
+    expect(result.ok).toBe(false);
+    expect(result.unreachable).toBe(false);
+    expect(result.failedKeys).toEqual([key]);
+  });
+});

package/src/mcp/__tests__/plugin-mcp-oauth-provider-isolation.test.ts

@@ -0,0 +1,89 @@
+import { beforeEach, describe, expect, jest, mock, test } from "bun:test";
+
+import type { McpOAuthCredentialTarget } from "../credential-target.js";
+
+const store = new Map<string, string>();
+
+mock.module("../../security/secure-keys.js", () => ({
+  getSecureKeyAsync: jest.fn(async (key: string) => store.get(key)),
+  setSecureKeyAsync: jest.fn(async (key: string, value: string) => {
+    store.set(key, value);
+    return true;
+  }),
+  deleteSecureKeyAsync: jest.fn(async (key: string) =>
+    store.delete(key) ? "deleted" : "not-found",
+  ),
+  listSecureKeysAsync: jest.fn(async () => ({
+    accounts: [...store.keys()],
+    unreachable: false,
+  })),
+}));
+
+const { McpOAuthProvider } = await import("../mcp-oauth-provider.js");
+const { pluginMcpOAuthCredentialTarget, workspaceMcpOAuthCredentialTarget } =
+  await import("../credential-target.js");
+
+function provider(
+  serverId: string,
+  url: string,
+  credentialTarget: McpOAuthCredentialTarget,
+) {
+  return new McpOAuthProvider(serverId, url, false, { credentialTarget });
+}
+
+describe("McpOAuthProvider plugin credential isolation", () => {
+  beforeEach(() => {
+    store.clear();
+  });
+
+  test("round-trips across provider instances without crossing owners or workspace", async () => {
+    const url = "https://mcp.example.com/a";
+    const pluginA = pluginMcpOAuthCredentialTarget("plugin-a", "server", {
+      type: "streamable-http",
+      url,
+    });
+    const pluginB = pluginMcpOAuthCredentialTarget("plugin-b", "server", {
+      type: "streamable-http",
+      url,
+    });
+    const workspace = workspaceMcpOAuthCredentialTarget("plugin-a");
+
+    await provider("plugin-a", url, pluginA).saveTokens({
+      access_token: "plugin-a-token",
+      token_type: "bearer",
+    });
+
+    expect(await provider("plugin-a", url, pluginA).tokens()).toMatchObject({
+      access_token: "plugin-a-token",
+    });
+    expect(await provider("plugin-b", url, pluginB).tokens()).toBeUndefined();
+    expect(await provider("plugin-a", url, workspace).tokens()).toBeUndefined();
+  });
+
+  test("an endpoint change cannot read or invalidate the old endpoint tokens", async () => {
+    const oldUrl = "https://mcp.example.com/a";
+    const newUrl = "https://mcp.example.com/b";
+    const oldTarget = pluginMcpOAuthCredentialTarget("plugin-a", "server", {
+      type: "streamable-http",
+      url: oldUrl,
+    });
+    const newTarget = pluginMcpOAuthCredentialTarget("plugin-a", "server", {
+      type: "streamable-http",
+      url: newUrl,
+    });
+    await provider("plugin-a", oldUrl, oldTarget).saveTokens({
+      access_token: "old-token",
+      token_type: "bearer",
+    });
+
+    expect(
+      await provider("plugin-a", newUrl, newTarget).tokens(),
+    ).toBeUndefined();
+    await provider("plugin-a", newUrl, newTarget).invalidateCredentials(
+      "tokens",
+    );
+    expect(
+      await provider("plugin-a", oldUrl, oldTarget).tokens(),
+    ).toMatchObject({ access_token: "old-token" });
+  });
+});

package/src/mcp/__tests__/plugin-server-credential-isolation.test.ts

@@ -1,6 +1,7 @@
 /**
- * A plugin-declared server must connect without resolving
- * `mcp:<serverId>:*` from the workspace credential store.
+ * A plugin-declared server must connect with its plugin and endpoint-scoped
+ * OAuth credentials without resolving `mcp:<serverId>:*` from the workspace
+ * credential store.
  *
  * A plugin controls both its server key and its URL, so a stored
  * credential whose id happens to match would otherwise be sent to an
@@ -12,7 +13,9 @@
 
 import { beforeEach, describe, expect, jest, mock, test } from "bun:test";
 
-const getSecureKeyAsync = jest.fn(async (_key: string) => "stored-oauth-token");
+const getSecureKeyAsync = jest.fn(
+  async (_key: string): Promise<string | null> => "stored-oauth-token",
+);
 const getMcpHeaders = jest.fn(async (_serverId: string) => ({
   Authorization: "Bearer workspace-secret",
 }));
@@ -36,10 +39,19 @@ const { McpServerManager } = await import("../manager.js");
 /** Refused immediately, so `connect` fails after the credential lookups. */
 const UNREACHABLE = "http://127.0.0.1:1/mcp";
 
-function httpServer(source: "workspace" | "plugin") {
+function workspaceHttpServer() {
   return {
     transport: { type: "streamable-http" as const, url: UNREACHABLE },
-    source,
+    source: "workspace" as const,
+  };
+}
+
+function pluginHttpServer(url = UNREACHABLE) {
+  return {
+    transport: { type: "streamable-http" as const, url },
+    source: "plugin" as const,
+    pluginName: "unabyss",
+    serverKey: "unabyss",
   };
 }
 
@@ -49,22 +61,34 @@ describe("plugin-declared MCP servers", () => {
     getMcpHeaders.mockClear();
   });
 
-  test("connect without resolving workspace credentials", async () => {
+  test("connect with plugin-scoped OAuth credentials", async () => {
     const manager = new McpServerManager();
 
     await manager.start({
-      servers: { unabyss: httpServer("plugin") },
+      servers: { unabyss: pluginHttpServer() },
     });
 
     expect(getMcpHeaders).not.toHaveBeenCalled();
-    expect(getSecureKeyAsync).not.toHaveBeenCalled();
+    expect(getSecureKeyAsync.mock.calls.length).toBeGreaterThan(0);
+    expect(getSecureKeyAsync.mock.calls.map(([key]) => key)).toEqual(
+      expect.arrayContaining([
+        expect.stringMatching(
+          /^mcp-plugin\/v1\/dW5hYnlzcw\/dW5hYnlzcw\/[a-f0-9]{64}\/tokens$/,
+        ),
+      ]),
+    );
+    expect(
+      getSecureKeyAsync.mock.calls.some(
+        ([key]) => key === "mcp:unabyss:tokens",
+      ),
+    ).toBe(false);
   });
 
   test("workspace servers keep resolving theirs", async () => {
     const manager = new McpServerManager();
 
     await manager.start({
-      servers: { "from-workspace": httpServer("workspace") },
+      servers: { "from-workspace": workspaceHttpServer() },
     });
 
     expect(getMcpHeaders).toHaveBeenCalledWith("from-workspace");
@@ -78,12 +102,40 @@ describe("plugin-declared MCP servers", () => {
 
     await manager.start({
       servers: {
-        unabyss: httpServer("plugin"),
-        "from-workspace": httpServer("workspace"),
+        unabyss: pluginHttpServer(),
+        "from-workspace": workspaceHttpServer(),
       },
     });
 
     const lookedUpIds = getMcpHeaders.mock.calls.map(([id]) => id);
     expect(lookedUpIds).toEqual(["from-workspace"]);
   });
+
+  test("does not reuse an OAuth provider after the endpoint changes", async () => {
+    const { McpClient } = await import("../client.js");
+    const endpointA = "https://mcp.example.com/a";
+    const endpointB = "https://mcp.example.com/b";
+    getSecureKeyAsync.mockImplementation(async (key: string) =>
+      key.endsWith("/tokens") &&
+      key.includes("mcp-plugin/") &&
+      getSecureKeyAsync.mock.calls.length === 1
+        ? "stored-oauth-token"
+        : null,
+    );
+    const client = new McpClient("unabyss", pluginHttpServer(endpointA));
+    const providers: unknown[] = [];
+    (client as any).createTransport = () => {
+      providers.push((client as any).oauthProvider);
+      return {};
+    };
+    (client as any).client.connect = async () => {};
+    (client as any).client.close = async () => {};
+
+    await client.connect({ type: "streamable-http", url: endpointA });
+    await client.disconnect();
+    await client.connect({ type: "streamable-http", url: endpointB });
+
+    expect(providers[0]).not.toBeNull();
+    expect(providers[1]).toBeNull();
+  });
 });