987 lines · 32.1 KB · TypeScript
RawDownload
import { Duration, Effect, Exit, Scope, ScopedCache } from "effect"; import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";import type { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js"; import {  definePlugin,  SecretId,  SetSecretInput,  SourceDetectionResult,  type PluginCtx,  type StorageFailure,} from "@executor/sdk"; import {  makeMcpStore,  mcpSchema,  type McpBindingStore,  type McpStoredSource,} from "./binding-store";import {  createMcpConnector,  type ConnectorInput,  type McpConnection,} from "./connection";import { discoverTools } from "./discover";import {  McpConnectionError,  McpOAuthError,  McpToolDiscoveryError,} from "./errors";import { invokeMcpTool } from "./invoke";import {  deriveMcpNamespace,  type McpToolManifestEntry,} from "./manifest";import { exchangeMcpOAuthCode, startMcpOAuthAuthorization } from "./oauth";import { McpToolBinding, type McpConnectionAuth, type McpStoredSourceData } from "./types"; import {  SECRET_REF_PREFIX,  type ConfigFileSink,  type McpAuthConfig,  type McpRemoteSourceConfig as McpRemoteConfigEntry,  type McpStdioSourceConfig as McpStdioConfigEntry,  type SourceConfig,} from "@executor/config"; // ---------------------------------------------------------------------------// Plugin config — discriminated union on transport// --------------------------------------------------------------------------- export interface McpRemoteSourceConfig {  readonly transport: "remote";  readonly name: string;  readonly endpoint: string;  readonly remoteTransport?: "streamable-http" | "sse" | "auto";  readonly queryParams?: Record<string, string>;  readonly headers?: Record<string, string>;  readonly namespace?: string;  readonly auth?: McpConnectionAuth;} export interface McpStdioSourceConfig {  readonly transport: "stdio";  readonly name: string;  readonly command: string;  readonly args?: string[];  readonly env?: Record<string, string>;  readonly cwd?: string;  readonly namespace?: string;} export type McpSourceConfig = McpRemoteSourceConfig | McpStdioSourceConfig; // ---------------------------------------------------------------------------// Extension types// --------------------------------------------------------------------------- export interface McpOAuthStartInput {  readonly endpoint: string;  readonly redirectUrl: string;  readonly queryParams?: Record<string, string> | null;} export interface McpOAuthStartResponse {  readonly sessionId: string;  readonly authorizationUrl: string;} export interface McpOAuthCompleteInput {  readonly state: string;  readonly code?: string;  readonly error?: string;} export interface McpOAuthCompleteResponse {  readonly accessTokenSecretId: string;  readonly refreshTokenSecretId: string | null;  readonly tokenType: string;  readonly expiresAt: number | null;  readonly scope: string | null;} export interface McpProbeResult {  readonly connected: boolean;  readonly requiresOAuth: boolean;  readonly name: string;  readonly namespace: string;  readonly toolCount: number | null;  readonly serverName: string | null;} export interface McpUpdateSourceInput {  readonly name?: string;  readonly endpoint?: string;  readonly headers?: Record<string, string>;  readonly queryParams?: Record<string, string>;  readonly auth?: McpConnectionAuth;} // ---------------------------------------------------------------------------// Helpers// --------------------------------------------------------------------------- const toStoredSourceData = (config: McpSourceConfig): McpStoredSourceData => {  if (config.transport === "stdio") {    return {      transport: "stdio",      command: config.command,      args: config.args,      env: config.env,      cwd: config.cwd,    };  }  return {    transport: "remote",    endpoint: config.endpoint,    remoteTransport: config.remoteTransport ?? "auto",    queryParams: config.queryParams,    headers: config.headers,    auth: config.auth ?? { kind: "none" },  };}; const normalizeNamespace = (config: McpSourceConfig): string =>  config.namespace ??  deriveMcpNamespace({    name: config.name,    endpoint: config.transport === "remote" ? config.endpoint : undefined,    command: config.transport === "stdio" ? config.command : undefined,  }); const toBinding = (entry: McpToolManifestEntry): McpToolBinding =>  new McpToolBinding({    toolId: entry.toolId,    toolName: entry.toolName,    description: entry.description,    inputSchema: entry.inputSchema,    outputSchema: entry.outputSchema,  }); const makeOAuthProvider = (  accessToken: string,  tokenType: string,  refreshToken?: string,): OAuthClientProvider => ({  get redirectUrl() {    return "http://localhost/oauth/callback";  },  get clientMetadata() {    return {      redirect_uris: ["http://localhost/oauth/callback"],      grant_types: ["authorization_code", "refresh_token"] as string[],      response_types: ["code"] as string[],      token_endpoint_auth_method: "none" as const,      client_name: "Executor",    };  },  clientInformation: () => undefined,  saveClientInformation: () => {},  tokens: async (): Promise<OAuthTokens> => ({    access_token: accessToken,    token_type: tokenType,    ...(refreshToken ? { refresh_token: refreshToken } : {}),  }),  saveTokens: async () => {},  redirectToAuthorization: async () => {    throw new Error("MCP OAuth re-authorization required");  },  saveCodeVerifier: () => {},  codeVerifier: () => {    throw new Error("No active PKCE verifier");  },  saveDiscoveryState: () => {},  discoveryState: () => undefined,}); const remoteConnectionError = (message: string) =>  new McpConnectionError({ transport: "remote", message }); const mcpOAuthError = (message: string) => new McpOAuthError({ message }); const mcpDiscoveryError = (message: string) =>  new McpToolDiscoveryError({ stage: "list_tools", message }); // ---------------------------------------------------------------------------// Shared connector resolution — reads secrets, builds stdio/remote input// --------------------------------------------------------------------------- const resolveConnectorInput = (  sd: McpStoredSourceData,  ctx: PluginCtx<McpBindingStore>,  allowStdio: boolean,): Effect.Effect<ConnectorInput, McpConnectionError | StorageFailure> => {  if (sd.transport === "stdio") {    if (!allowStdio) {      return Effect.fail(        new McpConnectionError({          transport: "stdio",          message:            "MCP stdio transport is disabled. Enable it by passing `dangerouslyAllowStdioMCP: true` to mcpPlugin() — only safe for trusted local contexts.",        }),      );    }    return Effect.succeed({      transport: "stdio" as const,      command: sd.command,      args: sd.args,      env: sd.env,      cwd: sd.cwd,    });  }   return Effect.gen(function* () {    const headers: Record<string, string> = { ...sd.headers };    let authProvider: OAuthClientProvider | undefined;     const auth = sd.auth;    if (auth.kind === "header") {      const val = yield* ctx.secrets.get(auth.secretId);      if (val === null) {        return yield* Effect.fail(          remoteConnectionError(`Failed to resolve secret "${auth.secretId}"`),        );      }      headers[auth.headerName] = auth.prefix ? `${auth.prefix}${val}` : val;    } else if (auth.kind === "oauth2") {      const accessToken = yield* ctx.secrets.get(auth.accessTokenSecretId);      if (accessToken === null) {        return yield* Effect.fail(          remoteConnectionError("Failed to resolve OAuth access token"),        );      }      const refreshToken = auth.refreshTokenSecretId        ? (yield* ctx.secrets.get(auth.refreshTokenSecretId)) ?? undefined        : undefined;      authProvider = makeOAuthProvider(        accessToken,        auth.tokenType ?? "Bearer",        refreshToken,      );    }     return {      transport: "remote" as const,      endpoint: sd.endpoint,      remoteTransport: sd.remoteTransport,      queryParams: sd.queryParams,      headers: Object.keys(headers).length > 0 ? headers : undefined,      authProvider,    };  });}; // ---------------------------------------------------------------------------// Connection cache — kept as plugin-module state so both invokeTool and// the close hook see the same ScopedCache instance. The ScopedCache's// lookup key is the stringified stored source data identity.// --------------------------------------------------------------------------- interface McpRuntime {  readonly connectionCache: ScopedCache.ScopedCache<    string,    McpConnection,    McpConnectionError  >;  readonly pendingConnectors: Map<    string,    Effect.Effect<McpConnection, McpConnectionError>  >;  readonly cacheScope: Scope.CloseableScope;} const makeRuntime = (): Effect.Effect<McpRuntime, never> =>  Effect.gen(function* () {    const cacheScope = yield* Scope.make();    const pendingConnectors = new Map<      string,      Effect.Effect<McpConnection, McpConnectionError>    >();    const connectionCache = yield* ScopedCache.make({      lookup: (key: string) =>        Effect.acquireRelease(          Effect.suspend(() => {            const connector = pendingConnectors.get(key);            if (!connector) {              return Effect.fail(                new McpConnectionError({                  transport: "auto",                  message: `No pending connector for key: ${key}`,                }),              );            }            return connector;          }),          (connection) =>            Effect.promise(() => connection.close().catch(() => {})),        ),      capacity: 64,      timeToLive: Duration.minutes(5),    }).pipe(Scope.extend(cacheScope));     return { connectionCache, pendingConnectors, cacheScope };  }); // ---------------------------------------------------------------------------// Plugin factory// --------------------------------------------------------------------------- export interface McpPluginOptions {  /**   * Allow configuring stdio-transport MCP sources. Off by default.   *   * Stdio sources spawn a local subprocess that inherits the parent   * `process.env`. Only enable for trusted single-user contexts.   */  readonly dangerouslyAllowStdioMCP?: boolean;  /** If provided, source add/remove is mirrored to executor.jsonc   *  (best-effort — file errors are logged, not raised). */  readonly configFile?: ConfigFileSink;} const secretRef = (id: string): string => `${SECRET_REF_PREFIX}${id}`; const authToConfig = (auth: McpConnectionAuth | undefined): McpAuthConfig | undefined => {  if (!auth) return undefined;  if (auth.kind === "none") return { kind: "none" };  if (auth.kind === "header") {    return {      kind: "header",      headerName: auth.headerName,      secret: secretRef(auth.secretId),      prefix: auth.prefix,    };  }  return {    kind: "oauth2",    accessTokenSecret: secretRef(auth.accessTokenSecretId),    refreshTokenSecret: auth.refreshTokenSecretId      ? secretRef(auth.refreshTokenSecretId)      : null,    tokenType: auth.tokenType,  };}; const toMcpConfigEntry = (  namespace: string,  sourceName: string,  config: McpSourceConfig,): SourceConfig => {  if (config.transport === "stdio") {    const entry: McpStdioConfigEntry = {      kind: "mcp",      transport: "stdio",      name: sourceName,      command: config.command,      args: config.args,      env: config.env,      cwd: config.cwd,      namespace,    };    return entry;  }  const entry: McpRemoteConfigEntry = {    kind: "mcp",    transport: "remote",    name: sourceName,    endpoint: config.endpoint,    remoteTransport: config.remoteTransport,    queryParams: config.queryParams,    headers: config.headers,    namespace,    auth: authToConfig(config.auth),  };  return entry;}; export const mcpPlugin = definePlugin(  (options?: McpPluginOptions) => {    const allowStdio = options?.dangerouslyAllowStdioMCP ?? false;    // Per-plugin-instance runtime holder. Captured by closures in    // `extension`, `invokeTool`, and `close`, so all three see the same    // connection cache across a single createExecutor lifecycle.    const runtimeRef: { current: McpRuntime | null } = { current: null };     const ensureRuntime = (): Effect.Effect<McpRuntime, never> =>      runtimeRef.current        ? Effect.succeed(runtimeRef.current)        : makeRuntime().pipe(            Effect.tap((rt) =>              Effect.sync(() => {                runtimeRef.current = rt;              }),            ),          );     return {      id: "mcp" as const,      schema: mcpSchema,      storage: (deps): McpBindingStore => makeMcpStore(deps),       extension: (ctx) => {        const probeEndpoint = (endpoint: string) =>          Effect.gen(function* () {            const trimmed = endpoint.trim();            if (!trimmed) {              return yield* Effect.fail(                remoteConnectionError("Endpoint URL is required"),              );            }             const name = yield* Effect.try(              () => new URL(trimmed).hostname,            ).pipe(Effect.orElseSucceed(() => "mcp"));            const namespace = deriveMcpNamespace({ endpoint: trimmed });             const connector = createMcpConnector({              transport: "remote",              endpoint: trimmed,            });             const result = yield* discoverTools(connector).pipe(              Effect.map((m) => ({ ok: true as const, manifest: m })),              Effect.catchAll(() =>                Effect.succeed({ ok: false as const, manifest: null }),              ),            );             if (result.ok && result.manifest) {              return {                connected: true,                requiresOAuth: false,                name: result.manifest.server?.name ?? name,                namespace,                toolCount: result.manifest.tools.length,                serverName: result.manifest.server?.name ?? null,              } satisfies McpProbeResult;            }             const hasOAuth = yield* startMcpOAuthAuthorization({              endpoint: trimmed,              redirectUrl: "http://127.0.0.1/executor/discovery/oauth/probe",              state: "probe",            }).pipe(              Effect.map(() => true),              Effect.catchAll(() => Effect.succeed(false)),            );             if (hasOAuth) {              return {                connected: false,                requiresOAuth: true,                name,                namespace,                toolCount: null,                serverName: null,              } satisfies McpProbeResult;            }             return yield* Effect.fail(              remoteConnectionError(                "Could not connect to MCP endpoint and no OAuth was detected",              ),            );          });         const configFile = options?.configFile;         const addSource = (config: McpSourceConfig) =>          Effect.gen(function* () {            const namespace = normalizeNamespace(config);            const sd = toStoredSourceData(config);             // Resolve auth (may fail if stdio gate is off)            const ci = yield* resolveConnectorInput(sd, ctx, allowStdio);             const connector = createMcpConnector(ci);            // Try discovery. If it fails (auth, network, bad spec), we still            // want the source to land in the catalog so users see it in

Lines 1–500 of 987.