@vellumai/assistant 0.7.0 → 0.7.2

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

@@ -0,0 +1,304 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+// ── Module mocks (must precede imports) ───────────────────────────────────────
+
+// Track captured callbacks and deferred code promises for each McpOAuthProvider
+// instance so tests can drive the flow.
+let capturedOnAuthorizationUrl: ((url: string) => void) | undefined;
+let deferredCodeResolve: ((code: string) => void) | undefined;
+let deferredCodeReject: ((err: Error) => void) | undefined;
+
+const mockInvalidateCredentials = mock(async () => {});
+const mockStartCallbackServer = mock(async () => {
+  const codePromise = new Promise<string>((resolve, reject) => {
+    deferredCodeResolve = resolve;
+    deferredCodeReject = reject;
+  });
+  return { codePromise };
+});
+const mockStopCallbackServer = mock(() => {});
+
+mock.module("../mcp-oauth-provider.js", () => ({
+  McpOAuthProvider: class {
+    constructor(
+      _serverId: string,
+      _serverUrl: string,
+      _interactive: boolean,
+      _callbackTransport: string,
+      options: { onAuthorizationUrl?: (url: string) => void } = {},
+    ) {
+      capturedOnAuthorizationUrl = options.onAuthorizationUrl;
+    }
+    invalidateCredentials = mockInvalidateCredentials;
+    startCallbackServer = mockStartCallbackServer;
+    stopCallbackServer = mockStopCallbackServer;
+  },
+}));
+
+const mockSetMcpAuthPending = mock(
+  (_serverId: string, _authUrl: string, _attemptId: 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,
+);
+const mockSetMcpAuthError = mock(
+  (_serverId: string, _error: string, _attemptId: string): boolean => true,
+);
+
+mock.module("../mcp-auth-state.js", () => ({
+  setMcpAuthPending: (...args: unknown[]) =>
+    mockSetMcpAuthPending(...(args as [string, string, string])),
+  setMcpAuthComplete: (...args: unknown[]) =>
+    mockSetMcpAuthComplete(...(args as [string, string])),
+  setMcpAuthError: (...args: unknown[]) =>
+    mockSetMcpAuthError(...(args as [string, string, string])),
+}));
+
+const mockReloadMcpServers = mock(async () => ({
+  ok: true,
+  reloaded: 0,
+  servers: [],
+}));
+
+mock.module("../../daemon/mcp-reload-service.js", () => ({
+  reloadMcpServers: () => mockReloadMcpServers(),
+}));
+
+mock.module("../../config/env-registry.js", () => ({
+  getIsContainerized: () => false,
+}));
+
+mock.module("../../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// Create a fake UnauthorizedError class that the orchestrator's instanceof check
+// will recognize (since we're also mocking the auth module).
+class FakeUnauthorizedError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "UnauthorizedError";
+  }
+}
+
+mock.module("@modelcontextprotocol/sdk/client/auth.js", () => ({
+  UnauthorizedError: FakeUnauthorizedError,
+}));
+
+const mockFinishAuth = mock(async (_code: string) => {});
+let mockConnectCallCount = 0;
+
+mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({
+  Client: class {
+    async connect() {
+      mockConnectCallCount++;
+      // Every connect fires onAuthorizationUrl and throws UnauthorizedError —
+      // the orchestrator never calls connect() a second time after finishAuth.
+      if (capturedOnAuthorizationUrl) {
+        capturedOnAuthorizationUrl("https://auth.example.com/oauth");
+      }
+      throw new FakeUnauthorizedError("unauthorized");
+    }
+    async close() {}
+  },
+}));
+
+mock.module("@modelcontextprotocol/sdk/client/sse.js", () => ({
+  SSEClientTransport: class {
+    constructor(_url: URL, _opts: unknown) {}
+    finishAuth = mockFinishAuth;
+  },
+}));
+
+mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
+  StreamableHTTPClientTransport: class {
+    constructor(_url: URL, _opts: unknown) {}
+    finishAuth = mockFinishAuth;
+  },
+}));
+
+// ── Import SUT after mocks ─────────────────────────────────────────────────────
+
+const { orchestrateMcpOAuthConnect } =
+  await import("../mcp-auth-orchestrator.js");
+
+// ── Helpers ────────────────────────────────────────────────────────────────────
+
+function resetMocks() {
+  capturedOnAuthorizationUrl = undefined;
+  deferredCodeResolve = undefined;
+  deferredCodeReject = undefined;
+  mockInvalidateCredentials.mockClear();
+  mockStartCallbackServer.mockClear();
+  mockStopCallbackServer.mockClear();
+  mockSetMcpAuthPending.mockClear();
+  mockSetMcpAuthComplete.mockClear();
+  mockSetMcpAuthError.mockClear();
+  mockReloadMcpServers.mockClear();
+  mockFinishAuth.mockClear();
+  // Reset complete/error mocks to default "applied=true" behavior; tests
+  // that exercise the superseded branch override these in-test.
+  mockSetMcpAuthComplete.mockImplementation(() => true);
+  mockSetMcpAuthError.mockImplementation(() => true);
+}
+
+// ── Tests ──────────────────────────────────────────────────────────────────────
+
+describe("orchestrateMcpOAuthConnect", () => {
+  beforeEach(() => {
+    resetMocks();
+    mockConnectCallCount = 0;
+  });
+
+  afterEach(() => {
+    resetMocks();
+    mockConnectCallCount = 0;
+  });
+
+  test("happy path — returns auth_url and sets state to pending", async () => {
+    const result = await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    expect(result.auth_url).toBe("https://auth.example.com/oauth");
+    expect(mockSetMcpAuthPending.mock.calls[0]).toEqual([
+      "test-server",
+      "https://auth.example.com/oauth",
+      expect.any(String) as unknown as string, // attemptId UUID
+    ]);
+    // Sanity-check the attemptId looks UUID-shaped
+    expect(mockSetMcpAuthPending.mock.calls[0][2]).toMatch(
+      /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/,
+    );
+    expect(result).toBeDefined();
+  });
+
+  test("tail completion — codePromise resolves → state goes to complete", async () => {
+    await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    // Resolve the code to trigger the background tail
+    deferredCodeResolve!("auth-code-123");
+
+    // Wait for fire-and-forget tail to settle
+    await new Promise<void>((resolve) => setTimeout(resolve, 20));
+
+    expect(mockFinishAuth).toHaveBeenCalledWith("auth-code-123");
+    // The orchestrator calls connect() exactly once (the initial attempt that triggers
+    // UnauthorizedError). It does NOT reconnect after finishAuth to avoid the
+    // "already started" error thrown by SSE/StreamableHTTP transports.
+    expect(mockConnectCallCount).toBe(1);
+    expect(mockSetMcpAuthComplete).toHaveBeenCalledWith(
+      "test-server",
+      expect.any(String) as unknown as string,
+    );
+    // Daemon-side reload should be triggered after a successful completion.
+    expect(mockReloadMcpServers).toHaveBeenCalled();
+  });
+
+  test("transport.finishAuth rejects → state goes to error", async () => {
+    mockFinishAuth.mockImplementationOnce(async () => {
+      throw new Error("exchange failed");
+    });
+
+    await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    deferredCodeResolve!("auth-code-456");
+
+    await new Promise<void>((resolve) => setTimeout(resolve, 20));
+
+    expect(mockSetMcpAuthError).toHaveBeenCalledWith(
+      "test-server",
+      "exchange failed",
+      expect.any(String) as unknown as string,
+    );
+    expect(mockSetMcpAuthComplete).not.toHaveBeenCalled();
+  });
+
+  test("codePromise rejects (timeout/user deny) → state goes to error", async () => {
+    await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    deferredCodeReject!(new Error("MCP OAuth callback timed out"));
+
+    await new Promise<void>((resolve) => setTimeout(resolve, 20));
+
+    expect(mockSetMcpAuthError).toHaveBeenCalledWith(
+      "test-server",
+      "MCP OAuth callback timed out",
+      expect.any(String) as unknown as string,
+    );
+    expect(mockSetMcpAuthComplete).not.toHaveBeenCalled();
+  });
+
+  test("re-start for same serverId while previous is pending → new state overwrites with a fresh attemptId", async () => {
+    await orchestrateMcpOAuthConnect({
+      serverId: "srv",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    await orchestrateMcpOAuthConnect({
+      serverId: "srv",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    expect(mockSetMcpAuthPending.mock.calls).toHaveLength(2);
+    expect(mockSetMcpAuthPending.mock.calls[0][0]).toBe("srv");
+    expect(mockSetMcpAuthPending.mock.calls[1][0]).toBe("srv");
+    // Each attempt gets a distinct UUID so superseded tails can be detected.
+    const firstAttemptId = mockSetMcpAuthPending.mock.calls[0][2];
+    const secondAttemptId = mockSetMcpAuthPending.mock.calls[1][2];
+    expect(firstAttemptId).not.toBe(secondAttemptId);
+  });
+
+  test("supersede — when setMcpAuthComplete returns false, daemon-side reload is NOT triggered", async () => {
+    // Simulate a newer attempt having taken the slot: completion writes
+    // for the older attempt should be skipped, and reload should not fire.
+    mockSetMcpAuthComplete.mockImplementation(() => false);
+
+    await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    deferredCodeResolve!("auth-code-superseded");
+    await new Promise<void>((resolve) => setTimeout(resolve, 20));
+
+    // The set was attempted but returned false, so no reload should happen
+    expect(mockSetMcpAuthComplete).toHaveBeenCalled();
+    expect(mockReloadMcpServers).not.toHaveBeenCalled();
+  });
+
+  test("reload failure after completion is logged but does not corrupt success state", async () => {
+    mockReloadMcpServers.mockImplementation(async () => {
+      throw new Error("reload boom");
+    });
+
+    await orchestrateMcpOAuthConnect({
+      serverId: "test-server",
+      transport: { url: "https://example.com", type: "sse" },
+    });
+
+    deferredCodeResolve!("auth-code-789");
+    await new Promise<void>((resolve) => setTimeout(resolve, 20));
+
+    // Completion was applied even though reload threw — the polling CLI
+    // still observes status=complete.
+    expect(mockSetMcpAuthComplete).toHaveBeenCalled();
+    expect(mockSetMcpAuthError).not.toHaveBeenCalled();
+    expect(mockReloadMcpServers).toHaveBeenCalled();
+  });
+});

package/src/mcp/manager.ts

@@ -169,8 +169,3 @@ export function getMcpServerManager(): McpServerManager {
   }
   return instance;
 }
-
-/** Reset singleton for testing. */
-export function __resetMcpManagerForTesting(): void {
-  instance = null;
-}

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

@@ -0,0 +1,213 @@
+/**
+ * Daemon-side orchestrator for MCP OAuth flows.
+ *
+ * Runs the entire OAuth exchange (callback registration, auth URL capture,
+ * code exchange, token persistence) inside the daemon heap so that
+ * registerPendingCallback and consumeCallback always execute in the same
+ * process.  The CLI receives only the authorization URL via IPC and polls
+ * for completion.
+ */
+
+import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+
+import { getIsContainerized } from "../config/env-registry.js";
+import { reloadMcpServers } from "../daemon/mcp-reload-service.js";
+import { getLogger } from "../util/logger.js";
+import {
+  setMcpAuthComplete,
+  setMcpAuthError,
+  setMcpAuthPending,
+} from "./mcp-auth-state.js";
+import {
+  type McpOAuthCallbackTransport,
+  McpOAuthProvider,
+} from "./mcp-oauth-provider.js";
+
+const log = getLogger("mcp-auth-orchestrator");
+
+export interface McpAuthTransportConfig {
+  url: string;
+  type: "sse" | "streamable-http";
+  headers?: Record<string, string>;
+}
+
+export interface OrchestrateMcpOAuthConnectResult {
+  auth_url: string;
+  already_authenticated?: true;
+}
+
+/**
+ * Start a daemon-owned MCP OAuth flow.
+ *
+ * Returns immediately with the authorization URL for the CLI to open in
+ * the browser.  The token exchange runs in the background in the daemon heap
+ * and updates the in-memory auth state map on completion.
+ */
+export async function orchestrateMcpOAuthConnect(args: {
+  serverId: string;
+  transport: McpAuthTransportConfig;
+}): Promise<OrchestrateMcpOAuthConnectResult> {
+  const { serverId, transport } = args;
+
+  // Containerized deployments (platform-managed AND self-hosted Docker) must
+  // use the gateway transport: the browser is outside the daemon container,
+  // so a loopback callback inside the container is unreachable. Bare-metal
+  // daemons can use loopback as before.
+  const callbackTransport: McpOAuthCallbackTransport = getIsContainerized()
+    ? "gateway"
+    : "loopback";
+
+  let capturedAuthUrl: string | undefined;
+  const provider = new McpOAuthProvider(
+    serverId,
+    transport.url,
+    /* interactive */ false,
+    callbackTransport,
+    {
+      onAuthorizationUrl: (url) => {
+        capturedAuthUrl = url;
+      },
+    },
+  );
+
+  // Clear stale credentials so the flow starts fresh
+  await provider.invalidateCredentials("client");
+  await provider.invalidateCredentials("discovery");
+
+  // Register the pending callback in the daemon heap
+  const { codePromise } = await provider.startCallbackServer();
+
+  // Build the MCP transport and client
+  const serverUrl = new URL(transport.url);
+  const TransportClass =
+    transport.type === "sse"
+      ? SSEClientTransport
+      : StreamableHTTPClientTransport;
+  const mcpTransport = new TransportClass(serverUrl, {
+    authProvider: provider,
+    requestInit: transport.headers ? { headers: transport.headers } : undefined,
+  });
+  const client = new Client({ name: "vellum-assistant", version: "1.0.0" });
+
+  try {
+    await client.connect(mcpTransport);
+    // No error — server is already authenticated
+    provider.stopCallbackServer();
+    try {
+      await client.close();
+    } catch {
+      /* ignore */
+    }
+    return { auth_url: "", already_authenticated: true };
+  } catch (err) {
+    if (err instanceof UnauthorizedError) {
+      // Expected — onAuthorizationUrl has fired, capturedAuthUrl is set
+    } else {
+      provider.stopCallbackServer();
+      try {
+        await client.close();
+      } catch {
+        /* ignore */
+      }
+      throw err;
+    }
+  }
+
+  if (!capturedAuthUrl) {
+    provider.stopCallbackServer();
+    try {
+      await client.close();
+    } catch {
+      /* ignore */
+    }
+    throw new Error("No authorization URL captured from OAuth provider");
+  }
+
+  // Per-attempt token. Gates fire-and-forget state writes so that a
+  // re-run of `assistant mcp auth <serverId>` before the previous attempt
+  // finishes cannot have its slot overwritten by stale completion writes
+  // from the older attempt.
+  const attemptId = crypto.randomUUID();
+  setMcpAuthPending(serverId, capturedAuthUrl, attemptId);
+
+  // Fire-and-forget background tail — completes the token exchange once
+  // the user approves in the browser.
+  // Note: we do NOT call client.connect() again here — both SSEClientTransport
+  // and StreamableHTTPClientTransport throw "already started" on a second
+  // connect() call.  The tokens are persisted by saveTokens inside finishAuth,
+  // so the daemon can reconnect on the next MCP reload without re-connecting here.
+  void (async () => {
+    try {
+      // Apply an explicit timeout: the loopback callback path has a built-in
+      // 2-minute timer, but the gateway transport's deferred promise relies on
+      // the caller for time-boxing. Without this race, gateway-mode tails leak
+      // forever if the user never completes the OAuth handshake.
+      const code = await Promise.race([
+        codePromise,
+        new Promise<never>((_, reject) => {
+          const t = setTimeout(
+            () => reject(new Error("OAuth callback timed out")),
+            CALLBACK_TIMEOUT_MS,
+          );
+          // Don't keep the event loop alive solely for this timer
+          if (typeof t.unref === "function") t.unref();
+        }),
+      ]);
+      await mcpTransport.finishAuth(code);
+      const applied = setMcpAuthComplete(serverId, attemptId);
+      if (!applied) {
+        log.info(
+          { serverId, attemptId },
+          "MCP OAuth completion superseded by newer attempt — skipping state write",
+        );
+        return;
+      }
+      log.info({ serverId }, "MCP OAuth flow completed");
+      // Trigger MCP reload from inside the daemon so the CLI doesn't need
+      // to fall back on the deprecated file-based signal mechanism.
+      // Best-effort: reload failures are logged but don't poison the
+      // success status the polling CLI is about to observe.
+      try {
+        await reloadMcpServers();
+      } catch (reloadErr) {
+        log.warn(
+          {
+            serverId,
+            err:
+              reloadErr instanceof Error
+                ? reloadErr.message
+                : String(reloadErr),
+          },
+          "MCP reload after auth completion failed",
+        );
+      }
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      const applied = setMcpAuthError(serverId, message, attemptId);
+      if (!applied) {
+        log.info(
+          { serverId, attemptId, error: message },
+          "MCP OAuth error superseded by newer attempt — skipping state write",
+        );
+      } else {
+        log.warn({ serverId, error: message }, "MCP OAuth flow failed");
+      }
+    } finally {
+      provider.stopCallbackServer();
+      try {
+        await client.close();
+      } catch {
+        /* ignore */
+      }
+    }
+  })();
+
+  return { auth_url: capturedAuthUrl };
+}
+
+// Matches the loopback callback timeout; keep both in lockstep so loopback
+// and gateway transports time-box OAuth identically.
+const CALLBACK_TIMEOUT_MS = 2 * 60 * 1000;

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

@@ -0,0 +1,133 @@
+/**
+ * In-memory MCP OAuth flow status map.
+ *
+ * Tracks the current state of daemon-owned MCP OAuth flows so the CLI can
+ * poll for completion via the IPC route.
+ */
+
+/**
+ * Sibling: `assistant/src/security/oauth-callback-registry.ts`. The two look
+ * similar (in-memory map, supersede semantics, ~5 min TTL) but live at
+ * different layers. The callback registry stores the deferred resolve/reject
+ * pair for a single OAuth code arrival, keyed by OAuth `state`. This file
+ * stores observable status (pending / complete / error) keyed by serverId so
+ * the polling CLI can render progress without holding a long-lived IPC
+ * connection.
+ */
+type McpAuthState =
+  | { status: "pending"; authUrl: string; attemptId: string; expiresAt: number }
+  | {
+      status: "complete";
+      serverId: string;
+      attemptId: string;
+      completedAt: number;
+    }
+  | { status: "error"; error: string; attemptId: string; failedAt: number };
+
+const activeMcpAuthFlows = new Map<string, McpAuthState>();
+
+const PENDING_TTL_MS = 5 * 60 * 1000; // 5 min — matches oauth-callback-registry.ts
+const COMPLETION_GRACE_MS = 60 * 1000; // 60s so the polling CLI gets one final read
+
+/**
+ * Record that an OAuth flow is pending authorization.
+ * Overwrites any prior state for the same serverId (supersede semantics
+ * matching registerPendingCallback). The caller must pass an `attemptId`
+ * (a unique token per attempt) so that fire-and-forget completion writes
+ * can verify they still own the slot before mutating shared state — see
+ * `setMcpAuthComplete` / `setMcpAuthError`.
+ */
+export function setMcpAuthPending(
+  serverId: string,
+  authUrl: string,
+  attemptId: string,
+): void {
+  activeMcpAuthFlows.set(serverId, {
+    status: "pending",
+    authUrl,
+    attemptId,
+    expiresAt: Date.now() + PENDING_TTL_MS,
+  });
+}
+
+/**
+ * Record that an OAuth flow completed successfully. Returns true if the
+ * write was applied; false if the attempt has been superseded by a newer
+ * one (in which case the caller's tail should silently exit without
+ * touching state).
+ */
+export function setMcpAuthComplete(
+  serverId: string,
+  attemptId: string,
+): boolean {
+  const current = activeMcpAuthFlows.get(serverId);
+  if (current && current.attemptId !== attemptId) {
+    return false; // superseded
+  }
+  activeMcpAuthFlows.set(serverId, {
+    status: "complete",
+    serverId,
+    attemptId,
+    completedAt: Date.now(),
+  });
+  return true;
+}
+
+/**
+ * Record that an OAuth flow failed. Returns true if the write was applied;
+ * false if the attempt has been superseded.
+ */
+export function setMcpAuthError(
+  serverId: string,
+  error: string,
+  attemptId: string,
+): boolean {
+  const current = activeMcpAuthFlows.get(serverId);
+  if (current && current.attemptId !== attemptId) {
+    return false; // superseded
+  }
+  activeMcpAuthFlows.set(serverId, {
+    status: "error",
+    error,
+    attemptId,
+    failedAt: Date.now(),
+  });
+  return true;
+}
+
+/**
+ * Get the current state of an OAuth flow, or null if none exists.
+ */
+export function getMcpAuthState(serverId: string): McpAuthState | null {
+  clearExpiredMcpAuthStates();
+  return activeMcpAuthFlows.get(serverId) ?? null;
+}
+
+/**
+ * Remove expired entries from the auth flow map.
+ */
+export function clearExpiredMcpAuthStates(): void {
+  const now = Date.now();
+  for (const [serverId, state] of activeMcpAuthFlows) {
+    if (state.status === "pending" && now > state.expiresAt) {
+      activeMcpAuthFlows.delete(serverId);
+    } else if (
+      state.status === "complete" &&
+      now > state.completedAt + COMPLETION_GRACE_MS
+    ) {
+      activeMcpAuthFlows.delete(serverId);
+    } else if (
+      state.status === "error" &&
+      now > state.failedAt + COMPLETION_GRACE_MS
+    ) {
+      activeMcpAuthFlows.delete(serverId);
+    }
+  }
+}
+
+/**
+ * Test-only helper — clears all auth flow state for test isolation.
+ */
+export function _clearAllMcpAuthStates(): void {
+  activeMcpAuthFlows.clear();
+}