@vellumai/assistant 0.12.1 → 0.12.2-dev.202609162214.308bda0

package/src/mcp/__tests__/tool-caps.test.ts

@@ -0,0 +1,107 @@
+/**
+ * The global MCP tool cap used to consume the budget in insertion order,
+ * so later servers could lose every tool when earlier ones filled the
+ * 50-tool ceiling. Selection is a deterministic round-robin by server id.
+ */
+
+import { describe, expect, test } from "bun:test";
+
+import {
+  MCP_GLOBAL_MAX_TOOLS,
+  MCP_MAX_TOOLS_PER_SERVER,
+} from "../../config/schemas/mcp.js";
+import { applyMcpToolCaps } from "../tool-caps.js";
+
+function toolsNamed(prefix: string, count: number): string[] {
+  return Array.from({ length: count }, (_, i) => `${prefix}-tool-${i}`);
+}
+
+function setsOf(servers: Array<{ serverId: string; tools: string[] }>) {
+  return new Map(servers.map((server) => [server.serverId, server.tools]));
+}
+
+describe("applyMcpToolCaps", () => {
+  test("keeps every tool when the discovered set is under both caps", () => {
+    const result = applyMcpToolCaps([
+      { serverId: "alpha", tools: toolsNamed("alpha", 3) },
+      { serverId: "beta", tools: toolsNamed("beta", 2) },
+    ]);
+
+    expect(result.keptToolCount).toBe(5);
+    expect(result.droppedToolCount).toBe(0);
+    expect(result.servers.map((server) => server.tools.length)).toEqual([3, 2]);
+  });
+
+  test("applies the per-server cap before the global budget", () => {
+    const result = applyMcpToolCaps(
+      [{ serverId: "busy", tools: toolsNamed("busy", 8) }],
+      { perServerMax: 3, globalMax: 50 },
+    );
+
+    expect(result.servers[0]?.tools).toEqual(toolsNamed("busy", 3));
+    expect(result.decisions[0]?.droppedByPerServerCap).toBe(5);
+    expect(result.decisions[0]?.droppedByGlobalCap).toBe(0);
+  });
+
+  test("eight servers totaling about 50 tools do not starve later servers", () => {
+    const servers = Array.from({ length: 8 }, (_, i) => ({
+      serverId: `server-${String(i + 1).padStart(2, "0")}`,
+      tools: toolsNamed(`s${i + 1}`, 10),
+    }));
+
+    const result = applyMcpToolCaps(servers, {
+      perServerMax: MCP_MAX_TOOLS_PER_SERVER,
+      globalMax: MCP_GLOBAL_MAX_TOOLS,
+    });
+
+    expect(result.discoveredToolCount).toBe(80);
+    expect(result.keptToolCount).toBe(MCP_GLOBAL_MAX_TOOLS);
+    expect(
+      result.servers.every((server) => server.tools.length > 0),
+      "Every connected server keeps at least one tool under a 50-tool global cap.",
+    ).toBe(true);
+    expect(result.servers[5]?.tools.length).toBeGreaterThan(0);
+    expect(result.servers[6]?.tools.length).toBeGreaterThan(0);
+    expect(result.servers[7]?.tools.length).toBeGreaterThan(0);
+  });
+
+  test("reordering servers does not change which tools are kept", () => {
+    const forward = Array.from({ length: 8 }, (_, i) => ({
+      serverId: `server-${String(i + 1).padStart(2, "0")}`,
+      tools: toolsNamed(`s${i + 1}`, 10),
+    }));
+    const reversed = [...forward].reverse();
+
+    const a = applyMcpToolCaps(forward, { globalMax: 50, perServerMax: 20 });
+    const b = applyMcpToolCaps(reversed, { globalMax: 50, perServerMax: 20 });
+
+    expect(setsOf(a.servers)).toEqual(setsOf(b.servers));
+    expect(
+      a.servers.every((server) => server.tools.length > 0),
+    ).toBe(true);
+    expect(
+      b.servers.every((server) => server.tools.length > 0),
+    ).toBe(true);
+  });
+
+  test("does not use insertion order to empty a later server", () => {
+    const result = applyMcpToolCaps(
+      [
+        { serverId: "early", tools: toolsNamed("early", 10) },
+        { serverId: "mid", tools: toolsNamed("mid", 10) },
+        { serverId: "late", tools: toolsNamed("late", 10) },
+      ],
+      { globalMax: 12, perServerMax: 20 },
+    );
+
+    expect(result.servers.find((s) => s.serverId === "late")?.tools.length).toBe(
+      4,
+    );
+    expect(result.servers.find((s) => s.serverId === "early")?.tools.length).toBe(
+      4,
+    );
+    expect(result.servers.find((s) => s.serverId === "mid")?.tools.length).toBe(
+      4,
+    );
+  });
+});

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

@@ -0,0 +1,135 @@
+import { mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { beforeEach, describe, expect, test } from "bun:test";
+
+import { AGENT_PLUGINS_MCP_SCHEMA_URL } from "../spec-schema.js";
+import {
+  loadWorkspaceMcpConfig,
+  overlayWorkspaceMcpForConfigRead,
+  readWorkspaceMcpFile,
+  saveWorkspaceMcpConfig,
+} from "../workspace-mcp-config.js";
+
+function mcpPath(): string {
+  return join(process.env.VELLUM_WORKSPACE_DIR!, "mcp.json");
+}
+
+describe("workspace mcp.json", () => {
+  beforeEach(() => {
+    mkdirSync(process.env.VELLUM_WORKSPACE_DIR!, { recursive: true });
+    rmSync(mcpPath(), { force: true });
+  });
+
+  test("missing file is an empty server map", () => {
+    expect(loadWorkspaceMcpConfig()).toEqual({ servers: {} });
+    expect(readWorkspaceMcpFile().fileState).toBe("missing");
+  });
+
+  test("round-trips a spec document without wrapping transport", () => {
+    saveWorkspaceMcpConfig({
+      servers: {
+        remote: {
+          transport: {
+            type: "streamable-http",
+            url: "https://example.com/mcp",
+          },
+        },
+        local: {
+          transport: {
+            type: "stdio",
+            command: "npx",
+            args: ["-y", "srv"],
+          },
+        },
+      },
+    });
+
+    expect(JSON.parse(readFileSync(mcpPath(), "utf8"))).toEqual({
+      $schema: AGENT_PLUGINS_MCP_SCHEMA_URL,
+      mcpServers: {
+        remote: {
+          type: "streamable-http",
+          url: "https://example.com/mcp",
+        },
+        local: {
+          type: "stdio",
+          command: "npx",
+          args: ["-y", "srv"],
+        },
+      },
+    });
+    expect(loadWorkspaceMcpConfig()).toEqual({
+      servers: {
+        remote: {
+          transport: {
+            type: "streamable-http",
+            url: "https://example.com/mcp",
+          },
+        },
+        local: {
+          transport: {
+            type: "stdio",
+            command: "npx",
+            args: ["-y", "srv"],
+          },
+        },
+      },
+    });
+  });
+
+  test("does not persist transport headers", () => {
+    saveWorkspaceMcpConfig({
+      servers: {
+        remote: {
+          transport: {
+            type: "streamable-http",
+            url: "https://example.com/mcp",
+            headers: { Authorization: "Bearer secret" },
+          },
+        },
+      },
+    });
+
+    const onDisk = readFileSync(mcpPath(), "utf8");
+    expect(onDisk).not.toContain("secret");
+    expect(loadWorkspaceMcpConfig().servers.remote?.transport).toEqual({
+      type: "streamable-http",
+      url: "https://example.com/mcp",
+    });
+  });
+
+  test("invalid file does not throw on read and refuses a later write", () => {
+    mkdirSync(process.env.VELLUM_WORKSPACE_DIR!, { recursive: true });
+    writeFileSync(mcpPath(), "{ not json");
+
+    expect(loadWorkspaceMcpConfig()).toEqual({ servers: {} });
+    expect(readWorkspaceMcpFile().fileState).toBe("invalid");
+    expect(() => saveWorkspaceMcpConfig({ servers: {} })).toThrow(/invalid JSON/);
+  });
+
+  test("overlay projects transport-only servers for config reads", () => {
+    saveWorkspaceMcpConfig({
+      servers: {
+        remote: {
+          transport: {
+            type: "streamable-http",
+            url: "https://example.com/mcp",
+          },
+        },
+      },
+    });
+
+    const config: Record<string, unknown> = { llm: { activeProfile: "balanced" } };
+    overlayWorkspaceMcpForConfigRead(config);
+    expect(config.mcp).toEqual({
+      servers: {
+        remote: {
+          transport: {
+            type: "streamable-http",
+            url: "https://example.com/mcp",
+          },
+        },
+      },
+    });
+  });
+});

package/src/mcp/client.ts

@@ -5,11 +5,18 @@ import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import type { ToolAnnotations } from "@modelcontextprotocol/sdk/types.js";
 
-import type { McpServerSource, McpTransport } from "../config/schemas/mcp.js";
-import { getSecureKeyAsync } from "../security/secure-keys.js";
+import type {
+  McpServerSource,
+  McpTransport,
+  ResolvedMcpServerConfig,
+} from "../config/schemas/mcp.js";
 import { getLogger } from "../util/logger.js";
+import {
+  resolveMcpOAuthCredentialTarget,
+  workspaceMcpOAuthCredentialTarget,
+} from "./credential-target.js";
 import { getMcpHeaders } from "./mcp-header-store.js";
-import { McpOAuthProvider } from "./mcp-oauth-provider.js";
+import { hasMcpOAuthTokens, McpOAuthProvider } from "./mcp-oauth-provider.js";
 
 const log = getLogger("mcp-client");
 
@@ -55,11 +62,12 @@ export interface McpCallResult {
 export class McpClient {
   readonly serverId: string;
   /**
-   * Where this server was declared. Only a `workspace` server resolves
-   * `mcp:<serverId>:tokens` / `mcp:<serverId>:headers` from the credential
-   * store — see {@link McpServerSource} for why a plugin server must not.
+   * Where this server was declared. Workspace servers use their public id for
+   * OAuth and static-header identity. Plugin servers use their declaring owner,
+   * original key, and endpoint for OAuth and never read workspace headers.
    */
   readonly source: McpServerSource;
+  private readonly serverConfig: ResolvedMcpServerConfig | null;
   private client: Client;
   private transport:
     | StdioClientTransport
@@ -79,9 +87,18 @@ export class McpClient {
     return this.connected;
   }
 
-  constructor(serverId: string, source: McpServerSource = "workspace") {
+  constructor(
+    serverId: string,
+    sourceOrConfig: McpServerSource | ResolvedMcpServerConfig = "workspace",
+    private readonly onUnexpectedClose?: () => void,
+  ) {
     this.serverId = serverId;
-    this.source = source;
+    this.source =
+      typeof sourceOrConfig === "string"
+        ? sourceOrConfig
+        : sourceOrConfig.source;
+    this.serverConfig =
+      typeof sourceOrConfig === "string" ? null : sourceOrConfig;
     this.client = new Client({
       name: "vellum-assistant",
       version: "1.0.0",
@@ -96,6 +113,14 @@ export class McpClient {
         "MCP SDK transport error (non-fatal)",
       );
     };
+    this.client.onclose = () => {
+      const wasConnected = this.connected;
+      this.connected = false;
+      this.transport = null;
+      if (wasConnected) {
+        this.onUnexpectedClose?.();
+      }
+    };
   }
 
   async connect(transportConfig: McpTransport): Promise<void> {
@@ -106,21 +131,20 @@ export class McpClient {
     const isHttpTransport =
       transportConfig.type === "sse" ||
       transportConfig.type === "streamable-http";
-    const usesStoredCredentials = this.source === "workspace";
+    const credentialTarget = this.resolveCredentialTarget(transportConfig);
+    this.oauthProvider = null;
 
     // For HTTP transports, only attach an OAuth provider if cached tokens
     // exist. This avoids triggering client registration during daemon
     // startup. If no tokens, try without auth: if the server requires it,
     // skip silently.
-    if (isHttpTransport && usesStoredCredentials) {
-      const cachedTokens = await getSecureKeyAsync(
-        `mcp:${this.serverId}:tokens`,
-      );
-      if (cachedTokens) {
+    if (isHttpTransport && credentialTarget) {
+      if (await hasMcpOAuthTokens(credentialTarget)) {
         this.oauthProvider = new McpOAuthProvider(
           this.serverId,
           transportConfig.url,
           /* interactive */ false,
+          { credentialTarget },
         );
       }
     }
@@ -128,7 +152,7 @@ export class McpClient {
     // Resolve static auth headers from credential store, falling back to
     // any legacy headers in the transport config for backward compatibility.
     let effectiveConfig = transportConfig;
-    if (isHttpTransport && usesStoredCredentials) {
+    if (isHttpTransport && this.source === "workspace") {
       const storedHeaders = await getMcpHeaders(this.serverId);
       if (storedHeaders) {
         effectiveConfig = {
@@ -184,10 +208,34 @@ export class McpClient {
       return;
     }
 
+    if (this.transport === null) {
+      this._lastError = new Error(
+        `MCP server "${this.serverId}" closed during initialization`,
+      );
+      return;
+    }
+
     this.connected = true;
     log.info({ serverId: this.serverId }, "MCP client connected");
   }
 
+  private resolveCredentialTarget(
+    transport: McpTransport,
+  ): ReturnType<typeof resolveMcpOAuthCredentialTarget> {
+    if (this.source === "workspace") {
+      return transport.type === "stdio"
+        ? null
+        : workspaceMcpOAuthCredentialTarget(this.serverId);
+    }
+    if (!this.serverConfig) {
+      return null;
+    }
+    return resolveMcpOAuthCredentialTarget(this.serverId, {
+      ...this.serverConfig,
+      transport,
+    });
+  }
+
   async listTools(): Promise<McpToolInfo[]> {
     if (!this.connected) {
       throw new Error(`MCP client "${this.serverId}" is not connected`);
@@ -278,12 +326,12 @@ export class McpClient {
       return;
     }
 
+    this.connected = false;
     try {
       await this.client.close();
     } catch (err) {
       log.warn({ err, serverId: this.serverId }, "Error closing MCP client");
     }
-    this.connected = false;
     this.transport = null;
     log.info({ serverId: this.serverId }, "MCP client disconnected");
   }

package/src/mcp/credential-target.ts

@@ -0,0 +1,105 @@
+import { createHash } from "node:crypto";
+
+import type {
+  McpTransport,
+  ResolvedMcpServerConfig,
+} from "../config/schemas/mcp.js";
+
+export const MCP_OAUTH_CREDENTIAL_LEAVES = [
+  "tokens",
+  "client_info",
+  "client_binding",
+  "discovery",
+] as const;
+
+export type McpOAuthCredentialLeaf =
+  (typeof MCP_OAUTH_CREDENTIAL_LEAVES)[number];
+
+export type McpOAuthCredentialTarget =
+  | {
+      readonly source: "workspace";
+      readonly serverId: string;
+    }
+  | {
+      readonly source: "plugin";
+      readonly pluginName: string;
+      readonly serverKey: string;
+      readonly transportType: "sse" | "streamable-http";
+      readonly url: string;
+    };
+
+export function workspaceMcpOAuthCredentialTarget(
+  serverId: string,
+): McpOAuthCredentialTarget {
+  return { source: "workspace", serverId };
+}
+
+export function resolveMcpOAuthCredentialTarget(
+  serverId: string,
+  config: ResolvedMcpServerConfig,
+): McpOAuthCredentialTarget | null {
+  const transport = config.transport;
+  if (transport.type === "stdio") {
+    return null;
+  }
+  if (config.source === "workspace") {
+    return workspaceMcpOAuthCredentialTarget(serverId);
+  }
+  return pluginMcpOAuthCredentialTarget(
+    config.pluginName,
+    config.serverKey,
+    transport,
+  );
+}
+
+export function pluginMcpOAuthCredentialTarget(
+  pluginName: string,
+  serverKey: string,
+  transport: Extract<McpTransport, { type: "sse" | "streamable-http" }>,
+): McpOAuthCredentialTarget {
+  return {
+    source: "plugin",
+    pluginName,
+    serverKey,
+    transportType: transport.type,
+    url: transport.url,
+  };
+}
+
+export function mcpOAuthCredentialKey(
+  target: McpOAuthCredentialTarget,
+  leaf: McpOAuthCredentialLeaf,
+): string {
+  if (target.source === "workspace") {
+    return `mcp:${target.serverId}:${leaf}`;
+  }
+  return `${pluginMcpOAuthCredentialPrefix(target.pluginName, target.serverKey, target.transportType, target.url)}${leaf}`;
+}
+
+export function pluginMcpOAuthCredentialPrefix(
+  pluginName: string,
+  serverKey?: string,
+  transportType?: "sse" | "streamable-http",
+  rawUrl?: string,
+): string {
+  const pluginPrefix = `mcp-plugin/v1/${encodeSegment(pluginName)}/`;
+  if (serverKey === undefined) {
+    return pluginPrefix;
+  }
+  if (transportType === undefined || rawUrl === undefined) {
+    throw new Error("Plugin MCP credential endpoint is incomplete");
+  }
+  const endpoint = `${transportType}\n${canonicalUrlWithoutFragment(rawUrl)}`;
+  const endpointDigest = createHash("sha256").update(endpoint).digest("hex");
+  return `${pluginPrefix}${encodeSegment(serverKey)}/${endpointDigest}/`;
+}
+
+function encodeSegment(value: string): string {
+  return Buffer.from(value, "utf8").toString("base64url");
+}
+
+function canonicalUrlWithoutFragment(rawUrl: string): string {
+  const url = new URL(rawUrl);
+  url.hash = "";
+  return url.toString();
+}

package/src/mcp/effective-config.ts

@@ -1,11 +1,11 @@
 /**
  * The MCP configuration the daemon actually runs.
  *
- * Two sources contribute servers: the workspace `config.json`, which the
+ * Two sources contribute servers: the workspace `mcp.json`, which the
  * user owns, and each installed plugin's `mcp.json`, which its author
  * owns. Both must reach the MCP manager, so every consumer that starts
- * servers builds its config through here rather than reading
- * `config.mcp` directly — a path that only reads the workspace half
+ * servers builds its config through here rather than reading the
+ * workspace file directly — a path that only reads the workspace half
  * leaves a plugin's tools silently missing.
  *
  * Every server comes out attributed with its {@link McpServerSource}, so
@@ -42,6 +42,25 @@ let lastBuiltPluginFingerprint: string | null = null;
 export function buildEffectiveMcpConfig(
   workspaceConfig?: McpConfig,
 ): ResolvedMcpConfig {
+  const result = resolveEffectiveMcpConfig(workspaceConfig, true);
+  lastBuiltPluginFingerprint = fingerprintServers(result.pluginServers);
+  return result.config;
+}
+
+/** Resolve the effective config without recording it as applied by runtime. */
+export function readEffectiveMcpConfig(
+  workspaceConfig?: McpConfig,
+): ResolvedMcpConfig {
+  return resolveEffectiveMcpConfig(workspaceConfig, false).config;
+}
+
+function resolveEffectiveMcpConfig(
+  workspaceConfig: McpConfig | undefined,
+  reportIssues: boolean,
+): {
+  config: ResolvedMcpConfig;
+  pluginServers: ReturnType<typeof readPluginMcpServers>["servers"];
+} {
   const base = workspaceConfig ?? McpConfigSchema.parse({});
   const servers: Record<string, ResolvedMcpServerConfig> = {};
   for (const [id, config] of Object.entries(base.servers)) {
@@ -49,30 +68,32 @@ export function buildEffectiveMcpConfig(
   }
 
   const { servers: pluginServers, issues } = readPluginMcpServers();
-  for (const issue of issues) {
-    log.warn(
-      {
-        plugin: issue.pluginName,
-        ...(issue.serverKey && { serverKey: issue.serverKey }),
-      },
-      `Plugin MCP declaration problem: ${issue.message}`,
-    );
+  if (reportIssues) {
+    for (const issue of issues) {
+      log.warn(
+        {
+          plugin: issue.pluginName,
+          ...(issue.serverKey && { serverKey: issue.serverKey }),
+        },
+        `Plugin MCP declaration problem: ${issue.message}`,
+      );
+    }
   }
 
   for (const server of pluginServers) {
     if (Object.hasOwn(servers, server.id)) {
-      log.warn(
-        { plugin: server.pluginName, serverId: server.id },
-        "Plugin MCP server shadowed by a workspace server of the same id; skipping",
-      );
+      if (reportIssues) {
+        log.warn(
+          { plugin: server.pluginName, serverId: server.id },
+          "Plugin MCP server shadowed by a workspace server of the same id; skipping",
+        );
+      }
       continue;
     }
     servers[server.id] = server.config;
   }
 
-  lastBuiltPluginFingerprint = fingerprintPluginServers();
-
-  return { ...base, servers };
+  return { config: { ...base, servers }, pluginServers };
 }
 
 /**
@@ -98,6 +119,12 @@ export function pluginMcpServersChangedSinceLastBuild(): boolean {
  */
 function fingerprintPluginServers(): string {
   const { servers } = readPluginMcpServers();
+  return fingerprintServers(servers);
+}
+
+function fingerprintServers(
+  servers: ReturnType<typeof readPluginMcpServers>["servers"],
+): string {
   return JSON.stringify(
     [...servers]
       .sort((a, b) => a.id.localeCompare(b.id))