@vellumai/assistant 0.12.2 → 0.12.3-staging.1

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,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/manager.ts

@@ -1,11 +1,13 @@
-import {
-  MCP_GLOBAL_MAX_TOOLS,
-  MCP_MAX_TOOLS_PER_SERVER,
-  type ResolvedMcpConfig,
-  type ResolvedMcpServerConfig,
+import type {
+  ResolvedMcpConfig,
+  ResolvedMcpServerConfig,
 } from "../config/schemas/mcp.js";
 import { getLogger } from "../util/logger.js";
 import { McpClient, type McpToolInfo } from "./client.js";
+import {
+  applyMcpToolCaps,
+  truncatedServerIdsFromCaps,
+} from "./tool-caps.js";
 
 const log = getLogger("mcp-manager");
 
@@ -15,6 +17,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 +50,174 @@ 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,
+      })),
+    );
+    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> {

package/src/mcp/startup.ts

@@ -20,8 +20,14 @@ import { registerMcpTools, unregisterAllMcpTools } from "../tools/registry.js";
 import { getLogger } from "../util/logger.js";
 import { buildEffectiveMcpConfig } from "./effective-config.js";
 import { getMcpServerManager, stopMcpServerManager } from "./manager.js";
+import {
+  EMPTY_MCP_STARTUP_SNAPSHOT,
+  type McpStartupSnapshot,
+} from "./tool-caps.js";
 import { loadWorkspaceMcpConfig } from "./workspace-mcp-config.js";
 
+export type { McpStartupSnapshot } from "./tool-caps.js";
+
 const log = getLogger("mcp-startup");
 
 /**
@@ -39,23 +45,27 @@ const log = getLogger("mcp-startup");
  *
  * @param workspaceMcpConfig The projected workspace MCP config. Callers
  * that omit it read `/workspace/mcp.json`.
- * @returns The number of tools registered across all servers.
+ * @returns A count snapshot of configured/connected servers and the tools
+ * that reached this process's registry after caps.
  */
 export async function startConfiguredMcpServers(
   workspaceMcpConfig?: McpConfig,
-): Promise<number> {
-  let registered = 0;
+): Promise<McpStartupSnapshot> {
+  let snapshot: McpStartupSnapshot = EMPTY_MCP_STARTUP_SNAPSHOT;
   try {
     const mcpConfig = buildEffectiveMcpConfig(
       workspaceMcpConfig ?? loadWorkspaceMcpConfig(),
     );
-    if (Object.keys(mcpConfig.servers).length === 0) {
-      return 0;
+    const configuredServerCount = Object.keys(mcpConfig.servers).length;
+    snapshot = { ...EMPTY_MCP_STARTUP_SNAPSHOT, configuredServerCount };
+    if (configuredServerCount === 0) {
+      return snapshot;
     }
 
     const manager = getMcpServerManager();
-    const serverToolInfos = await manager.start(mcpConfig);
-    for (const { serverId, serverConfig, tools } of serverToolInfos) {
+    const started = await manager.start(mcpConfig);
+    let registered = 0;
+    for (const { serverId, serverConfig, tools } of started.servers) {
       const mcpTools = createMcpToolsFromServer(
         tools,
         serverId,
@@ -64,16 +74,26 @@ export async function startConfiguredMcpServers(
       );
       registered += registerMcpTools(serverId, mcpTools).length;
     }
+    snapshot = {
+      configuredServerCount: started.configuredServerCount,
+      connectedServerCount: started.connectedServerCount,
+      discoveredToolCount: started.discoveredToolCount,
+      registeredToolCount: registered,
+      droppedToolCount: started.droppedToolCount,
+      truncatedServerIds: started.truncatedServerIds,
+      errorServerCount: started.errorServerCount,
+      needsAuthServerCount: started.needsAuthServerCount,
+    };
   } catch (err) {
     log.error(
-      { err },
+      { err, ...snapshot },
       "MCP server initialization failed, continuing without MCP tools",
     );
-    return registered;
+    return snapshot;
   }
 
-  log.info({ toolCount: registered }, "MCP tools registered");
-  return registered;
+  log.info(snapshot, "MCP tools registered");
+  return snapshot;
 }
 
 /**
@@ -91,9 +111,9 @@ export async function startConfiguredMcpServers(
  * The window between the teardown and the reconnect has no MCP tools
  * registered, so an `mcp__*` call landing inside it fails as an unknown tool.
  *
- * @returns The number of tools registered after the reconnect.
+ * @returns A count snapshot of the tools registered after the reconnect.
  */
-export async function restartConfiguredMcpServers(): Promise<number> {
+export async function restartConfiguredMcpServers(): Promise<McpStartupSnapshot> {
   try {
     await stopMcpServerManager();
   } catch (err) {