@vellumai/assistant 0.12.2 → 0.12.3-staging.2

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

@@ -0,0 +1,156 @@
+/**
+ * McpServerManager.start() applies the shared fair cap after every
+ * configured server has been attempted in parallel. A slow earlier
+ * server must not prevent later ones from being listed, and later
+ * servers must keep tools when the global cap binds.
+ */
+
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type { ResolvedMcpConfig } from "../../config/schemas/mcp.js";
+
+const toolsByServer = new Map<string, Array<{ name: string }>>();
+const connectGates = new Map<string, Promise<void>>();
+const listedServers = new Set<string>();
+let mcpGlobalMaxTools: number | undefined;
+
+mock.module("../../config/loader.js", () => ({
+  getConfig: () => ({
+    tools: { exclude: [], mcpGlobalMaxTools },
+  }),
+}));
+
+mock.module("../client.js", () => ({
+  McpClient: class {
+    constructor(readonly serverId: string) {}
+    get isConnected() {
+      return true;
+    }
+    get lastError() {
+      return null;
+    }
+    async connect() {
+      const gate = connectGates.get(this.serverId);
+      if (gate) {
+        await gate;
+      }
+    }
+    async listTools() {
+      listedServers.add(this.serverId);
+      return (toolsByServer.get(this.serverId) ?? []).map((tool) => ({
+        name: tool.name,
+        description: `${this.serverId} ${tool.name}`,
+        inputSchema: { type: "object", properties: {} },
+      }));
+    }
+    async disconnect() {}
+  },
+}));
+
+const { McpServerManager } = await import("../manager.js");
+
+function httpServer(): ResolvedMcpConfig["servers"][string] {
+  return {
+    source: "workspace",
+    transport: {
+      type: "streamable-http",
+      url: "https://example.invalid/mcp",
+    },
+  };
+}
+
+function configWith(ids: string[]): ResolvedMcpConfig {
+  return {
+    servers: Object.fromEntries(ids.map((id) => [id, httpServer()])),
+  };
+}
+
+describe("McpServerManager tool selection", () => {
+  beforeEach(() => {
+    toolsByServer.clear();
+    connectGates.clear();
+    listedServers.clear();
+    mcpGlobalMaxTools = undefined;
+  });
+
+  test("later servers keep tools when eight servers exceed the global cap", async () => {
+    const ids = Array.from(
+      { length: 8 },
+      (_, i) => `server-${String(i + 1).padStart(2, "0")}`,
+    );
+    for (const id of ids) {
+      toolsByServer.set(
+        id,
+        Array.from({ length: 10 }, (_, i) => ({ name: `tool_${i}` })),
+      );
+    }
+
+    const manager = new McpServerManager();
+    const started = await manager.start(configWith(ids));
+
+    expect(started.discoveredToolCount).toBe(80);
+    expect(started.keptToolCount).toBe(50);
+    expect(started.truncatedServerIds.length).toBeGreaterThan(0);
+    expect(
+      started.servers.every((server) => server.tools.length > 0),
+      "Fair selection must not empty servers 6-8 merely because they were last.",
+    ).toBe(true);
+    await manager.stop();
+  });
+
+  test("a slow earlier server does not prevent later servers from connecting", async () => {
+    toolsByServer.set("slow", [{ name: "slow_tool" }]);
+    toolsByServer.set("fast", [{ name: "fast_tool" }]);
+    let releaseSlow!: () => void;
+    connectGates.set(
+      "slow",
+      new Promise<void>((resolve) => {
+        releaseSlow = resolve;
+      }),
+    );
+
+    const manager = new McpServerManager();
+    const starting = manager.start(configWith(["slow", "fast"]));
+    try {
+      await new Promise<void>((resolve) => setImmediate(resolve));
+      expect([...listedServers]).toEqual(["fast"]);
+
+      releaseSlow();
+      const started = await starting;
+      expect(started.connectedServerCount).toBe(2);
+      expect(started.servers.map((server) => server.serverId).sort()).toEqual([
+        "fast",
+        "slow",
+      ]);
+    } finally {
+      releaseSlow();
+      await starting;
+      await manager.stop();
+    }
+  });
+
+  test("a workspace global-max override raises how many tools are kept", async () => {
+    const ids = Array.from(
+      { length: 8 },
+      (_, i) => `server-${String(i + 1).padStart(2, "0")}`,
+    );
+    for (const id of ids) {
+      toolsByServer.set(
+        id,
+        Array.from({ length: 10 }, (_, i) => ({ name: `tool_${i}` })),
+      );
+    }
+
+    mcpGlobalMaxTools = 80;
+    const manager = new McpServerManager();
+    const started = await manager.start(configWith(ids));
+
+    expect(started.discoveredToolCount).toBe(80);
+    expect(started.keptToolCount).toBe(80);
+    expect(started.droppedToolCount).toBe(0);
+    expect(started.servers.every((server) => server.tools.length === 10)).toBe(
+      true,
+    );
+    await manager.stop();
+  });
+});

package/src/mcp/__tests__/startup.test.ts

@@ -15,7 +15,21 @@
 
 import { afterEach, describe, expect, jest, mock, test } from "bun:test";
 
-const start = jest.fn(async () => [] as unknown[]);
+function startResult(servers: unknown[] = []) {
+  return {
+    servers,
+    configuredServerCount: servers.length,
+    connectedServerCount: servers.length,
+    errorServerCount: 0,
+    needsAuthServerCount: 0,
+    discoveredToolCount: servers.length,
+    keptToolCount: servers.length,
+    droppedToolCount: 0,
+    truncatedServerIds: [] as string[],
+  };
+}
+
+const start = jest.fn(async () => startResult());
 
 mock.module("../manager.js", () => ({
   getMcpServerManager: () => ({
@@ -88,20 +102,24 @@ afterEach(() => {
 
 describe("startConfiguredMcpServers", () => {
   test("registers the tools a server reports, namespaced by server id", async () => {
-    start.mockResolvedValueOnce([serverReporting("fastmail", "search_email")]);
+    start.mockResolvedValueOnce(
+      startResult([serverReporting("fastmail", "search_email")]),
+    );
 
     const registered = await startConfiguredMcpServers(
       workspaceConfigWith("fastmail") as never,
     );
 
-    expect(registered).toBe(1);
+    expect(registered.registeredToolCount).toBe(1);
+    expect(registered.connectedServerCount).toBe(1);
     expect(getTool("mcp__fastmail__search_email")).toBeDefined();
   });
 
   test("does nothing when no servers are configured", async () => {
     const registered = await startConfiguredMcpServers(undefined);
 
-    expect(registered).toBe(0);
+    expect(registered.registeredToolCount).toBe(0);
+    expect(registered.configuredServerCount).toBe(0);
     expect(start).not.toHaveBeenCalled();
   });
 
@@ -112,13 +130,16 @@ describe("startConfiguredMcpServers", () => {
       workspaceConfigWith("fastmail") as never,
     );
 
-    expect(registered).toBe(0);
+    expect(registered.registeredToolCount).toBe(0);
+    expect(registered.configuredServerCount).toBe(1);
   });
 });
 
 describe("restartConfiguredMcpServers", () => {
   test("drops tools the reconnected set no longer reports", async () => {
-    start.mockResolvedValueOnce([serverReporting("fastmail", "search_email")]);
+    start.mockResolvedValueOnce(
+      startResult([serverReporting("fastmail", "search_email")]),
+    );
     await startConfiguredMcpServers(workspaceConfigWith("fastmail") as never);
     expect(getTool("mcp__fastmail__search_email")).toBeDefined();
 
@@ -135,7 +156,9 @@ describe("restartConfiguredMcpServers", () => {
   });
 
   test("reconnects from the config on disk rather than the previous set", async () => {
-    start.mockResolvedValueOnce([serverReporting("fastmail", "search_email")]);
+    start.mockResolvedValueOnce(
+      startResult([serverReporting("fastmail", "search_email")]),
+    );
     await startConfiguredMcpServers(workspaceConfigWith("fastmail") as never);
     start.mockClear();
 

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

@@ -0,0 +1,134 @@
+/**
+ * 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, resolveMcpGlobalMaxTools } 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,
+    );
+  });
+
+  test("honors a workspace global-max override", () => {
+    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, { globalMax: 80, perServerMax: 20 });
+
+    expect(result.keptToolCount).toBe(80);
+    expect(result.droppedToolCount).toBe(0);
+    expect(result.servers.every((server) => server.tools.length === 10)).toBe(
+      true,
+    );
+  });
+});
+
+describe("resolveMcpGlobalMaxTools", () => {
+  test("uses the shipped default when the workspace omits the override", () => {
+    expect(resolveMcpGlobalMaxTools()).toBe(MCP_GLOBAL_MAX_TOOLS);
+    expect(resolveMcpGlobalMaxTools({})).toBe(MCP_GLOBAL_MAX_TOOLS);
+  });
+
+  test("uses tools.mcpGlobalMaxTools from config.json when set", () => {
+    expect(resolveMcpGlobalMaxTools({ mcpGlobalMaxTools: 80 })).toBe(80);
+    expect(resolveMcpGlobalMaxTools({ mcpGlobalMaxTools: 1 })).toBe(1);
+  });
+});

package/src/mcp/manager.ts

@@ -1,11 +1,15 @@
-import {
-  MCP_GLOBAL_MAX_TOOLS,
-  MCP_MAX_TOOLS_PER_SERVER,
-  type ResolvedMcpConfig,
-  type ResolvedMcpServerConfig,
+import { getConfig } from "../config/loader.js";
+import type {
+  ResolvedMcpConfig,
+  ResolvedMcpServerConfig,
 } from "../config/schemas/mcp.js";
 import { getLogger } from "../util/logger.js";
 import { McpClient, type McpToolInfo } from "./client.js";
+import {
+  applyMcpToolCaps,
+  resolveMcpGlobalMaxTools,
+  truncatedServerIdsFromCaps,
+} from "./tool-caps.js";
 
 const log = getLogger("mcp-manager");
 
@@ -15,6 +19,18 @@ export interface McpServerToolInfo {
   tools: McpToolInfo[];
 }
 
+export interface McpStartResult {
+  servers: McpServerToolInfo[];
+  configuredServerCount: number;
+  connectedServerCount: number;
+  errorServerCount: number;
+  needsAuthServerCount: number;
+  discoveredToolCount: number;
+  keptToolCount: number;
+  droppedToolCount: number;
+  truncatedServerIds: string[];
+}
+
 export type McpConnectionState =
   | "connecting"
   | "connected"
@@ -36,121 +52,175 @@ export class McpServerManager {
     this.onUnexpectedClose = handler;
   }
 
-  async start(config: ResolvedMcpConfig): Promise<McpServerToolInfo[]> {
-    const results: McpServerToolInfo[] = [];
+  async start(config: ResolvedMcpConfig): Promise<McpStartResult> {
+    const entries = Object.entries(config.servers);
+    log.info(
+      { configuredServerCount: entries.length },
+      "Starting configured MCP servers",
+    );
+    console.log(`[MCP] Starting ${entries.length} server(s)...`);
 
-    console.log(
-      `[MCP] Starting ${Object.keys(config.servers).length} server(s)...`,
+    const started = await Promise.all(
+      entries.map(([serverId, serverConfig]) =>
+        this.startOne(serverId, serverConfig),
+      ),
     );
-    for (const [serverId, serverConfig] of Object.entries(config.servers)) {
-      this.connectionStates.set(serverId, {
-        source: serverConfig.source,
-        state: "connecting",
-      });
-      try {
-        console.log(
-          `[MCP] Starting server "${serverId}" (transport: ${serverConfig.transport.type})`,
-        );
-        if (
-          serverConfig.transport.type === "sse" ||
-          serverConfig.transport.type === "streamable-http"
-        ) {
-          log.debug(
-            { serverId },
-            "HTTP transport: OAuth provider will be available if server requires authentication",
-          );
-        }
-        const client = new McpClient(serverId, serverConfig, () => {
-          if (this.clients.get(serverId) !== client) {
-            return;
-          }
-          this.connectionStates.set(serverId, {
-            source: serverConfig.source,
-            state: "error",
-          });
-          this.onUnexpectedClose?.();
-        });
-        await client.connect(serverConfig.transport);
-
-        if (!client.isConnected) {
-          this.connectionStates.set(serverId, {
-            source: serverConfig.source,
-            state: client.lastError ? "error" : "needs-auth",
-          });
-          continue;
-        }
+    const connected = started.filter(
+      (result): result is McpServerToolInfo => result != null,
+    );
+    const errorServerCount = [...this.connectionStates.values()].filter(
+      (entry) => entry.state === "error",
+    ).length;
+    const needsAuthServerCount = [...this.connectionStates.values()].filter(
+      (entry) => entry.state === "needs-auth",
+    ).length;
 
-        this.clients.set(serverId, client);
+    const capped = applyMcpToolCaps(
+      connected.map((result) => ({
+        serverId: result.serverId,
+        tools: result.tools,
+      })),
+      { globalMax: resolveMcpGlobalMaxTools(getConfig().tools) },
+    );
+    const keptByServer = new Map(
+      capped.servers.map((server) => [server.serverId, server.tools]),
+    );
+    const results = connected.map((result) => ({
+      ...result,
+      tools: keptByServer.get(result.serverId) ?? [],
+    }));
+    const truncatedServerIds = truncatedServerIdsFromCaps(capped.decisions);
 
-        let tools = await client.listTools();
-        if (!client.isConnected || this.clients.get(serverId) !== client) {
-          this.clients.delete(serverId);
-          this.connectionStates.set(serverId, {
-            source: serverConfig.source,
-            state: "error",
-          });
-          continue;
-        }
-        log.info(
-          { serverId, toolCount: tools.length },
-          "MCP server tools discovered",
-        );
+    if (capped.droppedToolCount > 0) {
+      log.warn(
+        {
+          configuredServerCount: entries.length,
+          connectedServerCount: connected.length,
+          errorServerCount,
+          needsAuthServerCount,
+          discoveredToolCount: capped.discoveredToolCount,
+          keptToolCount: capped.keptToolCount,
+          droppedToolCount: capped.droppedToolCount,
+          globalCap: capped.globalCap,
+          perServerCap: capped.perServerCap,
+          truncatedServerIds,
+          decisions: capped.decisions.filter(
+            (decision) =>
+              decision.droppedByPerServerCap > 0 ||
+              decision.droppedByGlobalCap > 0,
+          ),
+        },
+        "MCP tool caps dropped tools using fair per-server selection",
+      );
+    } else {
+      log.info(
+        {
+          configuredServerCount: entries.length,
+          connectedServerCount: connected.length,
+          errorServerCount,
+          needsAuthServerCount,
+          discoveredToolCount: capped.discoveredToolCount,
+          keptToolCount: capped.keptToolCount,
+          globalCap: capped.globalCap,
+          perServerCap: capped.perServerCap,
+        },
+        "MCP servers connected",
+      );
+    }
 
-        if (tools.length > MCP_MAX_TOOLS_PER_SERVER) {
-          log.warn(
-            {
-              serverId,
-              discovered: tools.length,
-              max: MCP_MAX_TOOLS_PER_SERVER,
-            },
-            "MCP server exceeded per-server tool cap, truncating",
-          );
-          tools = tools.slice(0, MCP_MAX_TOOLS_PER_SERVER);
+    return {
+      servers: results,
+      configuredServerCount: entries.length,
+      connectedServerCount: connected.length,
+      errorServerCount,
+      needsAuthServerCount,
+      discoveredToolCount: capped.discoveredToolCount,
+      keptToolCount: capped.keptToolCount,
+      droppedToolCount: capped.droppedToolCount,
+      truncatedServerIds,
+    };
+  }
+
+  private async startOne(
+    serverId: string,
+    serverConfig: ResolvedMcpServerConfig,
+  ): Promise<McpServerToolInfo | null> {
+    this.connectionStates.set(serverId, {
+      source: serverConfig.source,
+      state: "connecting",
+    });
+    try {
+      console.log(
+        `[MCP] Starting server "${serverId}" (transport: ${serverConfig.transport.type})`,
+      );
+      if (
+        serverConfig.transport.type === "sse" ||
+        serverConfig.transport.type === "streamable-http"
+      ) {
+        log.debug(
+          { serverId },
+          "HTTP transport: OAuth provider will be available if server requires authentication",
+        );
+      }
+      const client = new McpClient(serverId, serverConfig, () => {
+        if (this.clients.get(serverId) !== client) {
+          return;
         }
+        this.connectionStates.set(serverId, {
+          source: serverConfig.source,
+          state: "error",
+        });
+        this.onUnexpectedClose?.();
+      });
+      await client.connect(serverConfig.transport);
 
-        results.push({ serverId, serverConfig, tools });
+      if (!client.isConnected) {
         this.connectionStates.set(serverId, {
           source: serverConfig.source,
-          state: "connected",
+          state: client.lastError ? "error" : "needs-auth",
         });
-      } catch (err) {
+        return null;
+      }
+
+      this.clients.set(serverId, client);
+
+      const tools = await client.listTools();
+      if (!client.isConnected || this.clients.get(serverId) !== client) {
+        this.clients.delete(serverId);
         this.connectionStates.set(serverId, {
           source: serverConfig.source,
           state: "error",
         });
-        console.error(`[MCP] Failed to connect to server "${serverId}":`, err);
-        log.error({ err, serverId }, "Failed to connect to MCP server");
-        // Clean up any partially-connected client
-        const staleClient = this.clients.get(serverId);
-        if (staleClient) {
-          try {
-            await staleClient.disconnect();
-          } catch {
-            /* ignore */
-          }
-          this.clients.delete(serverId);
-        }
+        return null;
       }
-    }
-
-    const totalTools = results.reduce((sum, r) => sum + r.tools.length, 0);
-    if (totalTools > MCP_GLOBAL_MAX_TOOLS) {
-      log.warn(
-        { totalTools, globalMax: MCP_GLOBAL_MAX_TOOLS },
-        "Total MCP tools exceed the global cap, truncating",
+      log.info(
+        { serverId, toolCount: tools.length },
+        "MCP server tools discovered",
       );
-      let remaining = MCP_GLOBAL_MAX_TOOLS;
-      for (const result of results) {
-        if (remaining <= 0) {
-          result.tools = [];
-        } else if (result.tools.length > remaining) {
-          result.tools = result.tools.slice(0, remaining);
+
+      this.connectionStates.set(serverId, {
+        source: serverConfig.source,
+        state: "connected",
+      });
+      return { serverId, serverConfig, tools };
+    } catch (err) {
+      this.connectionStates.set(serverId, {
+        source: serverConfig.source,
+        state: "error",
+      });
+      console.error(`[MCP] Failed to connect to server "${serverId}":`, err);
+      log.error({ err, serverId }, "Failed to connect to MCP server");
+      const staleClient = this.clients.get(serverId);
+      if (staleClient) {
+        try {
+          await staleClient.disconnect();
+        } catch {
+          /* ignore */
         }
-        remaining -= result.tools.length;
+        this.clients.delete(serverId);
       }
+      return null;
     }
-
-    return results;
   }
 
   async stop(): Promise<void> {