palaryn 0.1.0 → 0.3.2

package/src/mcp/bridge.ts

@@ -0,0 +1,703 @@
+import { Readable, Writable } from 'stream';
+import { randomUUID } from 'crypto';
+import { Gateway } from '../server/gateway';
+import { ToolCall, ToolCallArgs, Actor, Source, ToolInfo } from '../types/tool-call';
+import { ToolResult } from '../types/tool-result';
+import { GatewayConfig } from '../types/config';
+import { DEFAULT_CONFIG } from '../config/defaults';
+import {
+  MCPToolHandler,
+  MCPToolDefinition,
+  HTTP_TOOLS,
+  FILE_TOOLS,
+  SQL_TOOLS,
+  SHELL_TOOLS,
+  validateUrlArg,
+} from './tool-definitions';
+
+// ---------------------------------------------------------------------------
+// JSON-RPC 2.0 types for the MCP protocol
+// ---------------------------------------------------------------------------
+
+interface JSONRPCRequest {
+  jsonrpc: '2.0';
+  id: string | number;
+  method: string;
+  params?: Record<string, unknown>;
+}
+
+interface JSONRPCNotification {
+  jsonrpc: '2.0';
+  method: string;
+  params?: Record<string, unknown>;
+}
+
+interface JSONRPCResponse {
+  jsonrpc: '2.0';
+  id: string | number | null;
+  result?: unknown;
+  error?: {
+    code: number;
+    message: string;
+    data?: unknown;
+  };
+}
+
+type JSONRPCMessage = JSONRPCRequest | JSONRPCNotification;
+
+// ---------------------------------------------------------------------------
+// MCP protocol types (subset needed for tools)
+// ---------------------------------------------------------------------------
+
+/** MCP text content block */
+interface MCPTextContent {
+  type: 'text';
+  text: string;
+}
+
+/** MCP tool call result */
+interface MCPCallToolResult {
+  content: MCPTextContent[];
+  isError?: boolean;
+}
+
+// ---------------------------------------------------------------------------
+// MCP Bridge configuration
+// ---------------------------------------------------------------------------
+
+/**
+ * Configuration for the MCP bridge defaults.
+ * These values are used when MCP tool calls do not supply them explicitly.
+ */
+export interface MCPBridgeConfig {
+  /** Default workspace ID for tool calls */
+  workspace_id?: string;
+  /** Default actor for tool calls */
+  actor?: Actor;
+  /** Default source platform identifier */
+  source?: Source;
+  /** Default task ID (if not provided, a new UUID is generated per call) */
+  task_id?: string;
+}
+
+const DEFAULT_BRIDGE_CONFIG: Required<MCPBridgeConfig> = {
+  workspace_id: 'ws_mcp_default',
+  actor: { type: 'agent', id: 'mcp-agent', display: 'MCP Agent' },
+  source: { platform: 'mcp' },
+  task_id: '',
+};
+
+// ---------------------------------------------------------------------------
+// MCP protocol constants
+// ---------------------------------------------------------------------------
+
+const LATEST_PROTOCOL_VERSION = '2025-03-26';
+const JSONRPC_VERSION = '2.0';
+
+const SERVER_INFO = {
+  name: 'palaryn-mcp-bridge',
+  version: '1.0.0',
+};
+
+// JSON-RPC error codes
+const PARSE_ERROR = -32700;
+const INVALID_REQUEST = -32600;
+const METHOD_NOT_FOUND = -32601;
+const INVALID_PARAMS = -32602;
+const INTERNAL_ERROR = -32603;
+
+// ---------------------------------------------------------------------------
+// StdioTransport - line-delimited JSON-RPC over stdin/stdout
+// ---------------------------------------------------------------------------
+
+/**
+ * Reads line-delimited JSON-RPC messages from a readable stream
+ * and writes JSON-RPC responses to a writable stream.
+ * Follows the MCP stdio transport specification.
+ */
+class StdioTransport {
+  static readonly MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10MB
+  static readonly BUFFER_INACTIVITY_TIMEOUT_MS = 30_000; // 30s
+
+  private input: Readable;
+  private output: Writable;
+  private buffer: string = '';
+  private onMessage: ((msg: JSONRPCMessage) => void) | null = null;
+  private bufferTimeoutHandle: ReturnType<typeof setTimeout> | null = null;
+
+  constructor(input?: Readable, output?: Writable) {
+    this.input = input || process.stdin;
+    this.output = output || process.stdout;
+  }
+
+  /** Start listening for messages on the input stream. */
+  start(handler: (msg: JSONRPCMessage) => void): void {
+    this.onMessage = handler;
+    this.input.setEncoding('utf8');
+    this.input.on('data', (chunk: string) => this.onData(chunk));
+  }
+
+  /** Send a JSON-RPC response to the output stream. */
+  send(response: JSONRPCResponse): void {
+    const json = JSON.stringify(response);
+    this.output.write(json + '\n');
+  }
+
+  /** Stop listening and clean up. */
+  close(): void {
+    this.input.removeAllListeners('data');
+    this.onMessage = null;
+    this.clearBufferTimeout();
+  }
+
+  private clearBufferTimeout(): void {
+    if (this.bufferTimeoutHandle) {
+      clearTimeout(this.bufferTimeoutHandle);
+      this.bufferTimeoutHandle = null;
+    }
+  }
+
+  /** Reset the inactivity timer for partial buffer data. */
+  private resetBufferTimeout(): void {
+    this.clearBufferTimeout();
+    if (this.buffer.length > 0) {
+      this.bufferTimeoutHandle = setTimeout(() => {
+        if (this.buffer.length > 0) {
+          console.error(`[mcp-bridge] Buffer inactivity timeout (${StdioTransport.BUFFER_INACTIVITY_TIMEOUT_MS}ms) with ${this.buffer.length} bytes pending, clearing`);
+          this.buffer = '';
+          const errorResponse: JSONRPCResponse = {
+            jsonrpc: JSONRPC_VERSION,
+            id: null,
+            error: { code: PARSE_ERROR, message: 'Buffer inactivity timeout: incomplete message cleared' },
+          };
+          this.send(errorResponse);
+        }
+      }, StdioTransport.BUFFER_INACTIVITY_TIMEOUT_MS);
+    }
+  }
+
+  private onData(chunk: string): void {
+    this.buffer += chunk;
+
+    // Prevent unbounded buffer growth
+    if (this.buffer.length > StdioTransport.MAX_BUFFER_SIZE) {
+      console.error('[mcp-bridge] Buffer exceeded maximum size (10MB), clearing');
+      this.buffer = '';
+      this.clearBufferTimeout();
+      const errorResponse: JSONRPCResponse = {
+        jsonrpc: JSONRPC_VERSION,
+        id: null,
+        error: { code: PARSE_ERROR, message: 'Message too large: exceeded 10MB buffer limit' },
+      };
+      this.send(errorResponse);
+      return;
+    }
+
+    // Process complete lines
+    const lines = this.buffer.split('\n');
+    // Keep the last (possibly incomplete) line in the buffer
+    this.buffer = lines.pop() || '';
+
+    for (const line of lines) {
+      const trimmed = line.trim();
+      if (!trimmed) continue;
+
+      try {
+        const parsed = JSON.parse(trimmed);
+        if (this.onMessage) {
+          this.onMessage(parsed as JSONRPCMessage);
+        }
+      } catch {
+        // Send parse error for invalid JSON
+        const errorResponse: JSONRPCResponse = {
+          jsonrpc: JSONRPC_VERSION,
+          id: null,
+          error: { code: PARSE_ERROR, message: 'Parse error: invalid JSON' },
+        };
+        this.send(errorResponse);
+      }
+    }
+
+    // Reset inactivity timer if there's still partial data in the buffer
+    this.resetBufferTimeout();
+  }
+}
+
+// ---------------------------------------------------------------------------
+// MCPBridge - the main bridge class
+// ---------------------------------------------------------------------------
+
+/**
+ * MCPBridge wraps a Palaryn Gateway instance as an MCP server, exposing
+ * the gateway's tool execution capabilities through the Model Context Protocol.
+ *
+ * Communication uses JSON-RPC 2.0 over stdio (line-delimited).
+ *
+ * Tools are managed via a registry pattern. By default, HTTP tools are
+ * registered (http_request, http_get, http_post). Additional tools
+ * (file, sql, shell) can be registered dynamically.
+ *
+ * Each tool constructs a proper ToolCall, runs it through the full gateway
+ * pipeline (policy, DLP, budget, rate limiting, execution), and returns
+ * the ToolResult as the MCP response.
+ *
+ * Supported MCP methods:
+ * - `initialize` - Protocol handshake (returns server info and capabilities)
+ * - `notifications/initialized` - Client acknowledgment (no-op notification)
+ * - `tools/list` - List available tools with their JSON schemas
+ * - `tools/call` - Execute a tool through the gateway
+ * - `ping` - Health check
+ */
+export class MCPBridge {
+  private gateway: Gateway;
+  private bridgeConfig: Required<MCPBridgeConfig>;
+  private transport: StdioTransport | null = null;
+  private initialized: boolean = false;
+  private toolRegistry: Map<string, MCPToolHandler> = new Map();
+
+  constructor(gateway: Gateway, bridgeConfig?: MCPBridgeConfig) {
+    this.gateway = gateway;
+    this.bridgeConfig = {
+      ...DEFAULT_BRIDGE_CONFIG,
+      ...bridgeConfig,
+    };
+
+    // Register HTTP tools by default
+    this.registerTools(HTTP_TOOLS);
+  }
+
+  /**
+   * Register a single tool handler in the registry.
+   */
+  registerTool(handler: MCPToolHandler): void {
+    this.toolRegistry.set(handler.definition.name, handler);
+  }
+
+  /**
+   * Register multiple tool handlers in the registry.
+   */
+  registerTools(handlers: MCPToolHandler[]): void {
+    for (const handler of handlers) {
+      this.registerTool(handler);
+    }
+  }
+
+  /**
+   * Connect via stdio transport (reads from stdin, writes to stdout).
+   * This is the standard way to run an MCP server for CLI-based clients.
+   * Optionally accepts custom input/output streams for testing.
+   */
+  async connectStdio(input?: Readable, output?: Writable): Promise<void> {
+    this.transport = new StdioTransport(input, output);
+    this.transport.start((msg) => this.handleMessage(msg));
+  }
+
+  /**
+   * Close the transport and shut down the gateway.
+   */
+  async close(): Promise<void> {
+    if (this.transport) {
+      this.transport.close();
+      this.transport = null;
+    }
+    this.gateway.shutdown();
+  }
+
+  /**
+   * Returns the underlying Gateway instance.
+   */
+  getGateway(): Gateway {
+    return this.gateway;
+  }
+
+  /**
+   * Whether the MCP handshake has been completed.
+   */
+  isInitialized(): boolean {
+    return this.initialized;
+  }
+
+  // ---------------------------------------------------------------------------
+  // Message dispatch
+  // ---------------------------------------------------------------------------
+
+  /** Handle an incoming JSON-RPC message (request or notification). */
+  private handleMessage(msg: JSONRPCMessage): void {
+    // Notifications have no 'id' field
+    if (!('id' in msg) || msg.id === undefined) {
+      this.handleNotification(msg as JSONRPCNotification);
+      return;
+    }
+
+    const request = msg as JSONRPCRequest;
+    this.handleRequest(request).catch((err) => {
+      this.sendError(request.id, INTERNAL_ERROR, `Internal error: ${err instanceof Error ? err.message : String(err)}`);
+    });
+  }
+
+  /** Handle a JSON-RPC notification (no response expected). */
+  private handleNotification(msg: JSONRPCNotification): void {
+    // The only notification we handle is `notifications/initialized`
+    if (msg.method === 'notifications/initialized') {
+      // Client acknowledges initialization is complete; no action needed.
+      return;
+    }
+    // Unknown notifications are silently ignored per MCP spec.
+  }
+
+  /** Handle a JSON-RPC request and send a response. */
+  private async handleRequest(request: JSONRPCRequest): Promise<void> {
+    switch (request.method) {
+      case 'initialize':
+        this.handleInitialize(request);
+        break;
+      case 'tools/list':
+        this.handleToolsList(request);
+        break;
+      case 'tools/call':
+        await this.handleToolsCall(request);
+        break;
+      case 'ping':
+        this.sendResult(request.id, {});
+        break;
+      default:
+        this.sendError(request.id, METHOD_NOT_FOUND, `Method not found: ${request.method}`);
+        break;
+    }
+  }
+
+  // ---------------------------------------------------------------------------
+  // MCP method handlers
+  // ---------------------------------------------------------------------------
+
+  /** Handle `initialize` - protocol handshake. */
+  private handleInitialize(request: JSONRPCRequest): void {
+    this.initialized = true;
+
+    this.sendResult(request.id, {
+      protocolVersion: LATEST_PROTOCOL_VERSION,
+      capabilities: {
+        tools: {
+          listChanged: false,
+        },
+      },
+      serverInfo: SERVER_INFO,
+    });
+  }
+
+  /** Handle `tools/list` - return all registered tool definitions. */
+  private handleToolsList(request: JSONRPCRequest): void {
+    const tools = Array.from(this.toolRegistry.values()).map(h => h.definition);
+    this.sendResult(request.id, { tools });
+  }
+
+  /** Handle `tools/call` - execute a tool through the gateway. */
+  private async handleToolsCall(request: JSONRPCRequest): Promise<void> {
+    const params = request.params;
+    if (!params || typeof params.name !== 'string') {
+      this.sendError(request.id, INVALID_PARAMS, 'Missing required parameter: name');
+      return;
+    }
+
+    const toolName = params.name as string;
+    const toolArgs = (params.arguments || {}) as Record<string, unknown>;
+
+    const handler = this.toolRegistry.get(toolName);
+    if (!handler) {
+      this.sendError(request.id, INVALID_PARAMS, `Unknown tool: ${toolName}`);
+      return;
+    }
+
+    let result: MCPCallToolResult;
+    try {
+      // Validate required args for HTTP tools (url is required)
+      const urlError = validateUrlArg(toolArgs);
+      if (urlError && (toolName === 'http_request' || toolName === 'http_get' || toolName === 'http_post')) {
+        result = this.toolError(urlError);
+      } else {
+        const toolCall = handler.buildToolCall(toolArgs, this.bridgeConfig);
+        result = await this.executeAndFormat(toolCall);
+      }
+    } catch (err) {
+      result = {
+        content: [
+          {
+            type: 'text',
+            text: JSON.stringify({
+              error: err instanceof Error ? err.message : String(err),
+              status: 'error',
+            }, null, 2),
+          },
+        ],
+        isError: true,
+      };
+    }
+
+    this.sendResult(request.id, result);
+  }
+
+  // ---------------------------------------------------------------------------
+  // Internal helpers
+  // ---------------------------------------------------------------------------
+
+  /**
+   * Build a ToolCall from MCP tool arguments, applying bridge defaults.
+   * Kept for backward compatibility with tests that access it via reflection.
+   */
+  private buildToolCall(params: {
+    toolName: string;
+    capability: ToolInfo['capability'];
+    args: ToolCallArgs;
+    constraints?: { timeout_ms?: number; max_cost_usd?: number };
+    context?: { purpose?: string; labels?: string[] };
+  }): ToolCall {
+    const toolCall: ToolCall = {
+      tool_call_id: randomUUID(),
+      task_id: this.bridgeConfig.task_id || randomUUID(),
+      workspace_id: this.bridgeConfig.workspace_id,
+      actor: { ...this.bridgeConfig.actor },
+      source: { ...this.bridgeConfig.source },
+      tool: {
+        name: params.toolName,
+        version: '1.0.0',
+        capability: params.capability,
+      },
+      args: params.args,
+      timestamp: new Date().toISOString(),
+    };
+
+    // Add constraints if any are specified
+    if (params.constraints?.timeout_ms != null || params.constraints?.max_cost_usd != null) {
+      toolCall.constraints = {};
+      if (params.constraints.timeout_ms != null) {
+        toolCall.constraints.timeout_ms = params.constraints.timeout_ms;
+      }
+      if (params.constraints.max_cost_usd != null) {
+        toolCall.constraints.max_cost_usd = params.constraints.max_cost_usd;
+      }
+    }
+
+    // Add context if any is specified
+    if (params.context?.purpose || params.context?.labels) {
+      toolCall.context = {};
+      if (params.context.purpose) {
+        toolCall.context.purpose = params.context.purpose;
+      }
+      if (params.context.labels) {
+        toolCall.context.labels = params.context.labels;
+      }
+    }
+
+    return toolCall;
+  }
+
+  /**
+   * Execute a ToolCall through the gateway and format as MCP result.
+   */
+  private async executeAndFormat(toolCall: ToolCall): Promise<MCPCallToolResult> {
+    try {
+      const result: ToolResult = await this.gateway.execute(toolCall);
+      return this.formatResult(result);
+    } catch (err) {
+      const errorMessage = err instanceof Error ? err.message : String(err);
+      return {
+        content: [
+          {
+            type: 'text',
+            text: JSON.stringify(
+              {
+                error: errorMessage,
+                tool_call_id: toolCall.tool_call_id,
+                status: 'error',
+              },
+              null,
+              2,
+            ),
+          },
+        ],
+        isError: true,
+      };
+    }
+  }
+
+  /**
+   * Convert a gateway ToolResult into an MCP CallToolResult.
+   *
+   * The result includes:
+   * - The tool output body (or error message) as the primary text content
+   * - Gateway metadata (status, policy decision, DLP report, budget, timing)
+   *   as a second text content block for transparency
+   */
+  private formatResult(result: ToolResult): MCPCallToolResult {
+    const isError = result.status === 'error' || result.status === 'blocked';
+
+    // Primary content: the actual output or error
+    let primaryText: string;
+    if (result.error) {
+      primaryText = result.error;
+    } else if (result.output?.body !== undefined) {
+      primaryText =
+        typeof result.output.body === 'string'
+          ? result.output.body
+          : JSON.stringify(result.output.body, null, 2);
+    } else {
+      primaryText = `Request completed with status: ${result.status}`;
+    }
+
+    // Metadata block for gateway transparency
+    const metadata = {
+      tool_call_id: result.tool_call_id,
+      task_id: result.task_id,
+      status: result.status,
+      policy: result.policy,
+      dlp: {
+        detected: result.dlp.detected,
+        severity: result.dlp.severity,
+        redaction_count: result.dlp.redactions.length,
+      },
+      budget: result.budget,
+      timing: result.timing,
+      http_status: result.output?.http_status,
+      exit_code: result.output?.exit_code,
+      rows_affected: result.output?.rows_affected,
+      paths: result.output?.paths,
+      stderr: result.output?.stderr,
+    };
+
+    return {
+      content: [
+        {
+          type: 'text',
+          text: primaryText,
+        },
+        {
+          type: 'text',
+          text: `--- Gateway Metadata ---\n${JSON.stringify(metadata, null, 2)}`,
+        },
+      ],
+      isError,
+    };
+  }
+
+  /**
+   * Map an HTTP method to a ToolInfo capability level.
+   */
+  private methodToCapability(method: string): ToolInfo['capability'] {
+    switch (method.toUpperCase()) {
+      case 'GET':
+      case 'HEAD':
+      case 'OPTIONS':
+        return 'read';
+      case 'POST':
+      case 'PUT':
+      case 'PATCH':
+        return 'write';
+      case 'DELETE':
+        return 'delete';
+      default:
+        return 'write';
+    }
+  }
+
+  /**
+   * Attempt to parse a body string as JSON, falling back to the raw string.
+   */
+  private parseBody(body: string): unknown {
+    try {
+      return JSON.parse(body);
+    } catch {
+      return body;
+    }
+  }
+
+  /**
+   * Create an MCP tool error result.
+   */
+  private toolError(message: string): MCPCallToolResult {
+    return {
+      content: [{ type: 'text', text: message }],
+      isError: true,
+    };
+  }
+
+  /**
+   * Send a successful JSON-RPC response.
+   */
+  private sendResult(id: string | number, result: unknown): void {
+    if (this.transport) {
+      this.transport.send({
+        jsonrpc: JSONRPC_VERSION,
+        id,
+        result,
+      });
+    }
+  }
+
+  /**
+   * Send a JSON-RPC error response.
+   */
+  private sendError(id: string | number, code: number, message: string, data?: unknown): void {
+    if (this.transport) {
+      this.transport.send({
+        jsonrpc: JSONRPC_VERSION,
+        id,
+        error: { code, message, data },
+      });
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Convenience entry point
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a Gateway instance with the given config (or defaults),
+ * wrap it in an MCPBridge, and connect via stdio transport.
+ *
+ * This is the main entry point for running Palaryn as an MCP server.
+ *
+ * @param gatewayConfig - Full gateway configuration (defaults to DEFAULT_CONFIG with auth disabled)
+ * @param bridgeConfig - MCP bridge defaults for workspace, actor, source
+ * @returns The connected MCPBridge instance
+ */
+export async function startMCPBridge(
+  gatewayConfig?: Partial<GatewayConfig>,
+  bridgeConfig?: MCPBridgeConfig,
+): Promise<MCPBridge> {
+  // Merge with defaults, disabling auth for MCP (auth is handled by the MCP client/transport)
+  const config: GatewayConfig = {
+    ...DEFAULT_CONFIG,
+    ...gatewayConfig,
+    auth: {
+      ...DEFAULT_CONFIG.auth,
+      enabled: false,
+      ...gatewayConfig?.auth,
+    },
+    // Disable console audit output when running as MCP server
+    // to avoid polluting stdout (which is the MCP transport channel)
+    audit: {
+      ...DEFAULT_CONFIG.audit,
+      console_output: false,
+      ...gatewayConfig?.audit,
+    },
+  };
+
+  const gateway = new Gateway(config);
+  const bridge = new MCPBridge(gateway, bridgeConfig);
+
+  // Register non-HTTP tools based on executor config
+  if (config.executor.filesystem?.enabled) {
+    bridge.registerTools(FILE_TOOLS);
+  }
+  if (config.executor.sql?.enabled) {
+    bridge.registerTools(SQL_TOOLS);
+  }
+  if (config.executor.shell?.enabled) {
+    bridge.registerTools(SHELL_TOOLS);
+  }
+
+  await bridge.connectStdio();
+
+  return bridge;
+}