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

package/src/mcp/__tests__/reload-signal-emission.test.ts

@@ -11,16 +11,27 @@ import { existsSync, rmSync } from "node:fs";
 import { join } from "node:path";
 import { beforeEach, describe, expect, jest, mock, test } from "bun:test";
 
+const start = jest.fn(async () => []);
+const stop = jest.fn(async () => {});
+const publishMcpChanged = jest.fn();
+const realResourceSync =
+  await import("../../runtime/sync/resource-sync-events.js");
+
 mock.module("../manager.js", () => ({
   getMcpServerManager: () => ({
-    start: jest.fn(async () => []),
-    stop: jest.fn(async () => {}),
+    start,
+    stop,
     callTool: jest.fn(),
     getClient: () => undefined,
   }),
   stopMcpServerManager: jest.fn(async () => {}),
 }));
 
+mock.module("../../runtime/sync/resource-sync-events.js", () => ({
+  ...realResourceSync,
+  publishMcpChanged,
+}));
+
 mock.module("../mcp-header-store.js", () => ({
   migrateLegacyMcpHeaders: async () => {},
 }));
@@ -36,6 +47,11 @@ const { reloadMcpServers } = await import("../../daemon/mcp-reload-service.js");
 const signalPath = () => join(getSignalsDir(), MCP_RELOAD_SIGNAL_FILE);
 
 beforeEach(() => {
+  start.mockReset();
+  start.mockResolvedValue([]);
+  stop.mockReset();
+  stop.mockResolvedValue(undefined);
+  publishMcpChanged.mockReset();
   rmSync(signalPath(), { force: true });
 });
 
@@ -51,4 +67,25 @@ describe("reloadMcpServers", () => {
         "servers the config no longer lists.",
     ).toBe(true);
   });
+
+  test("publishes status only after a slow reload settles", async () => {
+    let finishStop: (() => void) | undefined;
+    stop.mockImplementationOnce(
+      () =>
+        new Promise<void>((resolve) => {
+          finishStop = resolve;
+        }),
+    );
+
+    const reload = reloadMcpServers();
+    while (!finishStop) {
+      await Promise.resolve();
+    }
+    expect(publishMcpChanged).not.toHaveBeenCalled();
+
+    finishStop!();
+    await reload;
+
+    expect(publishMcpChanged).toHaveBeenCalledTimes(1);
+  });
 });

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

@@ -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))

package/src/mcp/manager.ts

@@ -15,9 +15,26 @@ export interface McpServerToolInfo {
   tools: McpToolInfo[];
 }
 
+export type McpConnectionState =
+  | "connecting"
+  | "connected"
+  | "needs-auth"
+  | "error";
+
+export type McpUnexpectedCloseHandler = () => void;
+
 export class McpServerManager {
   private clients = new Map<string, McpClient>();
-  private serverConfigs = new Map<string, ResolvedMcpServerConfig>();
+  private connectionStates = new Map<
+    string,
+    { source: ResolvedMcpServerConfig["source"]; state: McpConnectionState }
+  >();
+
+  constructor(private onUnexpectedClose?: McpUnexpectedCloseHandler) {}
+
+  setUnexpectedCloseHandler(handler: McpUnexpectedCloseHandler): void {
+    this.onUnexpectedClose = handler;
+  }
 
   async start(config: ResolvedMcpConfig): Promise<McpServerToolInfo[]> {
     const results: McpServerToolInfo[] = [];
@@ -26,6 +43,10 @@ export class McpServerManager {
       `[MCP] Starting ${Object.keys(config.servers).length} server(s)...`,
     );
     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})`,
@@ -39,20 +60,37 @@ export class McpServerManager {
             "HTTP transport: OAuth provider will be available if server requires authentication",
           );
         }
-        // The server's own origin decides whether it may resolve
-        // `mcp:<serverId>:*` from the credential store.
-        const client = new McpClient(serverId, serverConfig.source);
+        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) {
-          // Server requires authentication — connect() logged guidance
+          this.connectionStates.set(serverId, {
+            source: serverConfig.source,
+            state: client.lastError ? "error" : "needs-auth",
+          });
           continue;
         }
 
         this.clients.set(serverId, client);
-        this.serverConfigs.set(serverId, serverConfig);
 
         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",
@@ -71,7 +109,15 @@ export class McpServerManager {
         }
 
         results.push({ serverId, serverConfig, tools });
+        this.connectionStates.set(serverId, {
+          source: serverConfig.source,
+          state: "connected",
+        });
       } 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");
         // Clean up any partially-connected client
@@ -83,7 +129,6 @@ export class McpServerManager {
             /* ignore */
           }
           this.clients.delete(serverId);
-          this.serverConfigs.delete(serverId);
         }
       }
     }
@@ -119,7 +164,7 @@ export class McpServerManager {
     );
     await Promise.all(disconnects);
     this.clients.clear();
-    this.serverConfigs.clear();
+    this.connectionStates.clear();
     log.info("All MCP servers disconnected");
   }
 
@@ -139,6 +184,14 @@ export class McpServerManager {
   getClient(serverId: string): McpClient | undefined {
     return this.clients.get(serverId);
   }
+
+  getServerState(
+    serverId: string,
+    source: ResolvedMcpServerConfig["source"],
+  ): McpConnectionState | undefined {
+    const entry = this.connectionStates.get(serverId);
+    return entry?.source === source ? entry.state : undefined;
+  }
 }
 
 // Singleton instance

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

@@ -15,6 +15,11 @@ import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/
 
 import { reloadMcpServers } from "../daemon/mcp-reload-service.js";
 import { getLogger } from "../util/logger.js";
+import {
+  mcpOAuthCredentialKey,
+  type McpOAuthCredentialTarget,
+  workspaceMcpOAuthCredentialTarget,
+} from "./credential-target.js";
 import {
   setMcpAuthComplete,
   setMcpAuthError,
@@ -46,8 +51,12 @@ export interface OrchestrateMcpOAuthConnectResult {
 export async function orchestrateMcpOAuthConnect(args: {
   serverId: string;
   transport: McpAuthTransportConfig;
+  credentialTarget?: McpOAuthCredentialTarget;
 }): Promise<OrchestrateMcpOAuthConnectResult> {
   const { serverId, transport } = args;
+  const credentialTarget =
+    args.credentialTarget ?? workspaceMcpOAuthCredentialTarget(serverId);
+  const targetKey = mcpOAuthCredentialKey(credentialTarget, "tokens");
 
   let capturedAuthUrl: string | undefined;
   const provider = new McpOAuthProvider(
@@ -55,6 +64,7 @@ export async function orchestrateMcpOAuthConnect(args: {
     transport.url,
     /* interactive */ false,
     {
+      credentialTarget,
       onAuthorizationUrl: (url) => {
         capturedAuthUrl = url;
       },
@@ -72,7 +82,10 @@ export async function orchestrateMcpOAuthConnect(args: {
   const { codePromise } = await provider.startCallbackServer();
 
   // Resolve effective headers: credential store takes precedence, then config
-  const storedHeaders = await getMcpHeaders(serverId);
+  const storedHeaders =
+    credentialTarget.source === "workspace"
+      ? await getMcpHeaders(serverId)
+      : null;
   const effectiveHeaders = storedHeaders ?? transport.headers;
 
   // Build the MCP transport and client
@@ -126,7 +139,7 @@ export async function orchestrateMcpOAuthConnect(args: {
   // finishes cannot have its slot overwritten by stale completion writes
   // from the older attempt.
   const attemptId = crypto.randomUUID();
-  setMcpAuthPending(serverId, capturedAuthUrl, attemptId);
+  setMcpAuthPending(serverId, capturedAuthUrl, attemptId, targetKey);
 
   // Fire-and-forget background tail — completes the token exchange once
   // the user approves in the browser.
@@ -153,7 +166,7 @@ export async function orchestrateMcpOAuthConnect(args: {
         }),
       ]);
       await mcpTransport.finishAuth(code);
-      const applied = setMcpAuthComplete(serverId, attemptId);
+      const applied = setMcpAuthComplete(serverId, attemptId, targetKey);
       if (!applied) {
         log.info(
           { serverId, attemptId },
@@ -182,7 +195,7 @@ export async function orchestrateMcpOAuthConnect(args: {
       }
     } catch (err) {
       const message = err instanceof Error ? err.message : String(err);
-      const applied = setMcpAuthError(serverId, message, attemptId);
+      const applied = setMcpAuthError(serverId, message, attemptId, targetKey);
       if (!applied) {
         log.info(
           { serverId, attemptId, error: message },