keystone-cli 0.5.0 → 0.6.0

package/src/runner/mcp-client.ts

@@ -1,6 +1,9 @@
 import { type ChildProcess, spawn } from 'node:child_process';
+import { lookup } from 'node:dns/promises';
+import { isIP } from 'node:net';
 import { type Interface, createInterface } from 'node:readline';
 import pkg from '../../package.json' with { type: 'json' };
+import { ConsoleLogger, type Logger } from '../utils/logger.ts';
 
 // MCP Protocol version - update when upgrading to newer MCP spec
 export const MCP_PROTOCOL_VERSION = '2024-11-05';
@@ -8,6 +11,142 @@ export const MCP_PROTOCOL_VERSION = '2024-11-05';
 // Maximum buffer size for incoming messages (10MB) to prevent memory exhaustion
 const MAX_BUFFER_SIZE = 10 * 1024 * 1024;
 
+/**
+ * Validate a URL to prevent SSRF attacks
+ * Blocks private IP ranges, localhost, and other internal addresses
+ * @param url The URL to validate
+ * @param options.allowInsecure If true, skips all security checks (use only for development/testing)
+ * @throws Error if the URL is potentially dangerous
+ */
+function isPrivateIpAddress(address: string): boolean {
+  const normalized = address.toLowerCase();
+  const parseMappedIpv4 = (mapped: string): string | null => {
+    const rest = mapped.replace(/^::ffff:/i, '');
+    if (rest.includes('.')) {
+      return rest;
+    }
+    const parts = rest.split(':');
+    if (parts.length !== 2) {
+      return null;
+    }
+    const high = Number.parseInt(parts[0], 16);
+    const low = Number.parseInt(parts[1], 16);
+    if (Number.isNaN(high) || Number.isNaN(low)) {
+      return null;
+    }
+    const a = (high >> 8) & 0xff;
+    const b = high & 0xff;
+    const c = (low >> 8) & 0xff;
+    const d = low & 0xff;
+    return `${a}.${b}.${c}.${d}`;
+  };
+
+  // IPv4 checks
+  const ipv4Match = normalized.match(/^(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$/);
+  if (ipv4Match) {
+    const [, a, b] = ipv4Match.map(Number);
+    return (
+      a === 10 || // 10.0.0.0/8
+      (a === 172 && b >= 16 && b <= 31) || // 172.16.0.0/12
+      (a === 192 && b === 168) || // 192.168.0.0/16
+      (a === 169 && b === 254) || // 169.254.0.0/16 (link-local)
+      a === 127 // 127.0.0.0/8
+    );
+  }
+
+  // IPv6-mapped IPv4 (::ffff:127.0.0.1 or ::ffff:7f00:1)
+  if (normalized.startsWith('::ffff:')) {
+    const mappedIpv4 = parseMappedIpv4(normalized);
+    if (mappedIpv4) {
+      return isPrivateIpAddress(mappedIpv4);
+    }
+  }
+
+  // IPv6 checks (best-effort, without full CIDR parsing)
+  const ipv6 = normalized.replace(/^\[|\]$/g, '');
+  return (
+    ipv6 === '::1' || // Loopback
+    ipv6.startsWith('fe80:') || // Link-local
+    ipv6.startsWith('fc') || // Unique local (fc00::/7)
+    ipv6.startsWith('fd') // Unique local (fc00::/7)
+  );
+}
+
+export async function validateRemoteUrl(
+  url: string,
+  options: { allowInsecure?: boolean } = {}
+): Promise<void> {
+  let parsed: URL;
+  try {
+    parsed = new URL(url);
+  } catch {
+    throw new Error(`Invalid MCP server URL: ${url}`);
+  }
+
+  // Skip all security checks if allowInsecure is set (for development/testing)
+  if (options.allowInsecure) {
+    return;
+  }
+
+  // Require HTTPS in production
+  if (parsed.protocol !== 'https:') {
+    throw new Error(
+      `SSRF Protection: MCP remote URL must use HTTPS. Got: ${parsed.protocol}. Set allowInsecure option to true if you trust this server.`
+    );
+  }
+
+  const hostname = parsed.hostname.toLowerCase();
+
+  // Block localhost variants
+  if (
+    hostname === 'localhost' ||
+    hostname === '127.0.0.1' ||
+    hostname === '::1' ||
+    hostname === '0.0.0.0' ||
+    hostname.endsWith('.localhost')
+  ) {
+    throw new Error(`SSRF Protection: Cannot connect to localhost/loopback address: ${hostname}`);
+  }
+
+  // Block private IP ranges (IPv4/IPv6) for literal IPs
+  if (isIP(hostname)) {
+    if (isPrivateIpAddress(hostname)) {
+      throw new Error(
+        `SSRF Protection: Cannot connect to private/internal IP address: ${hostname}`
+      );
+    }
+  }
+
+  // Block cloud metadata endpoints
+  if (
+    hostname === '169.254.169.254' || // AWS/GCP/Azure metadata
+    hostname === 'metadata.google.internal' ||
+    hostname.endsWith('.internal')
+  ) {
+    throw new Error(`SSRF Protection: Cannot connect to cloud metadata endpoint: ${hostname}`);
+  }
+
+  // Resolve DNS to prevent hostnames that map to private IPs (DNS rebinding)
+  if (!isIP(hostname)) {
+    try {
+      const resolved = await lookup(hostname, { all: true });
+      for (const record of resolved) {
+        if (isPrivateIpAddress(record.address)) {
+          throw new Error(
+            `SSRF Protection: Hostname "${hostname}" resolves to private/internal address: ${record.address}`
+          );
+        }
+      }
+    } catch (error) {
+      throw new Error(
+        `SSRF Protection: Failed to resolve hostname "${hostname}": ${
+          error instanceof Error ? error.message : String(error)
+        }`
+      );
+    }
+  }
+}
+
 interface MCPTool {
   name: string;
   description?: string;
@@ -32,17 +171,20 @@ interface MCPTransport {
   send(message: unknown): Promise<void>;
   onMessage(callback: (message: MCPResponse) => void): void;
   close(): void;
+  setLogger(logger: Logger): void;
 }
 
 class StdConfigTransport implements MCPTransport {
   private process: ChildProcess;
   private rl: Interface;
+  private logger: Logger = new ConsoleLogger();
 
   constructor(command: string, args: string[] = [], env: Record<string, string> = {}) {
     // Filter out sensitive environment variables from the host process
     // unless they are explicitly provided in the 'env' argument
     const safeEnv: Record<string, string> = {};
-    const sensitivePattern = /(?:key|token|secret|password|credential|auth|private)/i;
+    const sensitivePattern =
+      /(?:key|token|secret|password|credential|auth|private|cookie|session|signature)/i;
 
     for (const [key, value] of Object.entries(process.env)) {
       if (value && !sensitivePattern.test(key)) {
@@ -64,6 +206,10 @@ class StdConfigTransport implements MCPTransport {
     });
   }
 
+  setLogger(logger: Logger): void {
+    this.logger = logger;
+  }
+
   async send(message: unknown): Promise<void> {
     this.process.stdin?.write(`${JSON.stringify(message)}\n`);
   }
@@ -72,8 +218,8 @@ class StdConfigTransport implements MCPTransport {
     this.rl.on('line', (line) => {
       // Safety check for extremely long lines that might have bypassed readline's internal limits
       if (line.length > MAX_BUFFER_SIZE) {
-        process.stderr.write(
-          `[MCP Error] Received line exceeding maximum size (${line.length} bytes), ignoring.\n`
+        this.logger.error(
+          `[MCP Error] Received line exceeding maximum size (${line.length} bytes), ignoring.`
         );
         return;
       }
@@ -84,7 +230,7 @@ class StdConfigTransport implements MCPTransport {
       } catch (e) {
         // Log non-JSON lines to stderr so they show up in the terminal
         if (line.trim()) {
-          process.stderr.write(`[MCP Server Output] ${line}\n`);
+          this.logger.log(`[MCP Server Output] ${line}`);
         }
       }
     });
@@ -104,12 +250,17 @@ class SSETransport implements MCPTransport {
   private abortController: AbortController | null = null;
   private sessionId?: string;
   private activeReaders: Set<ReadableStreamDefaultReader<Uint8Array>> = new Set();
+  private logger: Logger = new ConsoleLogger();
 
   constructor(url: string, headers: Record<string, string> = {}) {
     this.url = url;
     this.headers = headers;
   }
 
+  setLogger(logger: Logger): void {
+    this.logger = logger;
+  }
+
   async connect(timeout = 60000): Promise<void> {
     this.abortController = new AbortController();
 
@@ -179,6 +330,18 @@ class SSETransport implements MCPTransport {
             const dispatchEvent = () => {
               if (currentEvent.data) {
                 if (currentEvent.event === 'endpoint') {
+                  // Validate endpoint to prevent SSRF - only allow relative paths
+                  const endpointValue = currentEvent.data;
+                  if (
+                    endpointValue &&
+                    (endpointValue.startsWith('http://') ||
+                      endpointValue.startsWith('https://') ||
+                      endpointValue.startsWith('//'))
+                  ) {
+                    throw new Error(
+                      `SSE endpoint must be a relative path, got absolute URL: ${endpointValue.substring(0, 50)}`
+                    );
+                  }
                   this.endpoint = currentEvent.data;
                   if (this.endpoint) {
                     this.endpoint = new URL(this.endpoint, this.url).href;
@@ -414,15 +577,25 @@ class SSETransport implements MCPTransport {
 export class MCPClient {
   private transport: MCPTransport;
   private messageId = 0;
-  private pendingRequests = new Map<number, (response: MCPResponse) => void>();
+  private pendingRequests = new Map<
+    number,
+    {
+      resolve: (response: MCPResponse) => void;
+      reject: (error: Error) => void;
+      timeoutId: ReturnType<typeof setTimeout>;
+    }
+  >();
   private timeout: number;
+  private logger: Logger;
 
   constructor(
     transportOrCommand: MCPTransport | string,
     timeoutOrArgs: number | string[] = [],
     env: Record<string, string> = {},
-    timeout = 60000
+    timeout = 60000,
+    logger: Logger = new ConsoleLogger()
   ) {
+    this.logger = logger;
     if (typeof transportOrCommand === 'string') {
       this.transport = new StdConfigTransport(transportOrCommand, timeoutOrArgs as string[], env);
       this.timeout = timeout;
@@ -430,13 +603,15 @@ export class MCPClient {
       this.transport = transportOrCommand;
       this.timeout = timeoutOrArgs as number;
     }
+    this.transport.setLogger(this.logger);
 
     this.transport.onMessage((response) => {
       if (response.id !== undefined && this.pendingRequests.has(response.id)) {
-        const resolve = this.pendingRequests.get(response.id);
-        if (resolve) {
+        const pending = this.pendingRequests.get(response.id);
+        if (pending) {
           this.pendingRequests.delete(response.id);
-          resolve(response);
+          clearTimeout(pending.timeoutId);
+          pending.resolve(response);
         }
       }
     });
@@ -446,20 +621,27 @@ export class MCPClient {
     command: string,
     args: string[] = [],
     env: Record<string, string> = {},
-    timeout = 60000
+    timeout = 60000,
+    logger: Logger = new ConsoleLogger()
   ): Promise<MCPClient> {
     const transport = new StdConfigTransport(command, args, env);
-    return new MCPClient(transport, timeout);
+    return new MCPClient(transport, timeout, {}, 0, logger);
   }
 
   static async createRemote(
     url: string,
     headers: Record<string, string> = {},
-    timeout = 60000
+    timeout = 60000,
+    options: { allowInsecure?: boolean; logger?: Logger } = {}
   ): Promise<MCPClient> {
+    // Validate URL to prevent SSRF attacks
+    await validateRemoteUrl(url, options);
+
+    const logger = options.logger || new ConsoleLogger();
     const transport = new SSETransport(url, headers);
+    transport.setLogger(logger);
     await transport.connect(timeout);
-    return new MCPClient(transport, timeout);
+    return new MCPClient(transport, timeout, {}, 0, logger);
   }
 
   private async request(
@@ -482,10 +664,7 @@ export class MCPClient {
         }
       }, this.timeout);
 
-      this.pendingRequests.set(id, (response) => {
-        clearTimeout(timeoutId);
-        resolve(response);
-      });
+      this.pendingRequests.set(id, { resolve, reject, timeoutId });
 
       this.transport.send(message).catch((err) => {
         clearTimeout(timeoutId);
@@ -524,8 +703,9 @@ export class MCPClient {
 
   stop() {
     // Reject all pending requests to prevent hanging callers
-    for (const [id, resolve] of this.pendingRequests) {
-      resolve({ id, error: { code: -1, message: 'MCP client stopped' } });
+    for (const [, pending] of this.pendingRequests) {
+      clearTimeout(pending.timeoutId);
+      pending.reject(new Error('MCP client stopped'));
     }
     this.pendingRequests.clear();
     this.transport.close();

package/src/runner/mcp-manager.ts

@@ -1,6 +1,6 @@
 import { ConfigLoader } from '../utils/config-loader';
+import { ConsoleLogger, type Logger } from '../utils/logger.ts';
 import { MCPClient } from './mcp-client';
-import type { Logger } from './workflow-runner';
 
 export interface MCPServerConfig {
   name: string;
@@ -20,9 +20,16 @@ export class MCPManager {
   private clients: Map<string, MCPClient> = new Map();
   private connectionPromises: Map<string, Promise<MCPClient | undefined>> = new Map();
   private sharedServers: Map<string, MCPServerConfig> = new Map();
+  private logger: Logger;
 
-  constructor() {
+  constructor(logger: Logger = new ConsoleLogger()) {
+    this.logger = logger;
     this.loadGlobalConfig();
+
+    // Ensure cleanup on process exit
+    process.on('exit', () => {
+      this.stopAll();
+    });
   }
 
   private loadGlobalConfig() {
@@ -39,14 +46,15 @@ export class MCPManager {
 
   async getClient(
     serverRef: string | MCPServerConfig,
-    logger: Logger = console
+    logger?: Logger
   ): Promise<MCPClient | undefined> {
+    const activeLogger = logger || this.logger;
     let config: MCPServerConfig;
 
     if (typeof serverRef === 'string') {
       const shared = this.sharedServers.get(serverRef);
       if (!shared) {
-        logger.error(`  ✗ Global MCP server not found: ${serverRef}`);
+        activeLogger.error(`  ✗ Global MCP server not found: ${serverRef}`);
         return undefined;
       }
       config = shared;
@@ -68,7 +76,7 @@ export class MCPManager {
 
     // Start a new connection and cache the promise
     const connectionPromise = (async () => {
-      logger.log(`  🔌 Connecting to MCP server: ${config.name} (${config.type || 'local'})`);
+      activeLogger.log(`  🔌 Connecting to MCP server: ${config.name} (${config.type || 'local'})`);
 
       let client: MCPClient;
       try {
@@ -91,7 +99,9 @@ export class MCPManager {
             headers.Authorization = `Bearer ${token}`;
           }
 
-          client = await MCPClient.createRemote(config.url, headers, config.timeout);
+          client = await MCPClient.createRemote(config.url, headers, config.timeout, {
+            logger: activeLogger,
+          });
         } else {
           if (!config.command) throw new Error('Local MCP server missing command');
 
@@ -118,7 +128,8 @@ export class MCPManager {
             config.command,
             config.args || [],
             env,
-            config.timeout
+            config.timeout,
+            activeLogger
           );
         }
 
@@ -126,7 +137,7 @@ export class MCPManager {
         this.clients.set(key, client);
         return client;
       } catch (error) {
-        logger.error(
+        activeLogger.error(
           `  ✗ Failed to connect to MCP server ${config.name}: ${error instanceof Error ? error.message : String(error)}`
         );
         return undefined;

package/src/runner/mcp-server.test.ts

@@ -166,7 +166,7 @@ describe('MCPServer', () => {
     expect(JSON.parse(response?.result?.content?.[0]?.text || '{}').status).toBe('success');
 
     // Verify DB was updated
-    const steps = db.getStepsByRun(runId);
+    const steps = await db.getStepsByRun(runId);
     expect(steps[0].status).toBe('success');
     expect(steps[0].output).toBeDefined();
     if (steps[0].output) {
@@ -301,10 +301,13 @@ describe('MCPServer', () => {
     expect(status.hint).toContain('still running');
   });
 
-  it('should call get_run_status tool for completed workflow', async () => {
-    const runId = 'completed-test-run';
+  it('should call get_run_status tool for success workflow', async () => {
+    const runId = 'success-test-run';
     await db.createRun(runId, 'test-wf', {});
-    await db.updateRunStatus(runId, 'completed', { output: 'done' });
+    await db.updateRunStatus(runId, 'success', { output: 'done' });
+
+    // Wait for the async run to finish and update DB
+    await new Promise((resolve) => setTimeout(resolve, 200));
 
     const response = await handleMessage({
       jsonrpc: '2.0',
@@ -314,7 +317,7 @@ describe('MCPServer', () => {
     });
 
     const status = JSON.parse(response?.result?.content?.[0]?.text || '{}');
-    expect(status.status).toBe('completed');
+    expect(status.status).toBe('success');
     expect(status.outputs).toEqual({ output: 'done' });
     expect(status.hint).toBeUndefined();
   });

package/src/runner/mcp-server.ts

@@ -3,6 +3,7 @@ import type { Readable, Writable } from 'node:stream';
 import pkg from '../../package.json' with { type: 'json' };
 import { WorkflowDb } from '../db/workflow-db';
 import { WorkflowParser } from '../parser/workflow-parser';
+import { ConsoleLogger, type Logger } from '../utils/logger';
 import { generateMermaidGraph } from '../utils/mermaid';
 import { WorkflowRegistry } from '../utils/workflow-registry';
 import { WorkflowSuspendedError } from './step-executor';
@@ -19,11 +20,18 @@ export class MCPServer {
   private db: WorkflowDb;
   private input: Readable;
   private output: Writable;
-
-  constructor(db?: WorkflowDb, input: Readable = process.stdin, output: Writable = process.stdout) {
+  private logger: Logger;
+
+  constructor(
+    db?: WorkflowDb,
+    input: Readable = process.stdin,
+    output: Writable = process.stdout,
+    logger: Logger = new ConsoleLogger()
+  ) {
     this.db = db || new WorkflowDb();
     this.input = input;
     this.output = output;
+    this.logger = logger;
   }
 
   async start() {
@@ -43,7 +51,7 @@ export class MCPServer {
             this.output.write(`${JSON.stringify(response)}\n`);
           }
         } catch (error) {
-          console.error('Error handling MCP message:', error);
+          this.logger.error(`Error handling MCP message: ${error}`);
         }
       });
 
@@ -54,7 +62,7 @@ export class MCPServer {
 
       // Handle stream errors
       this.input.on('error', (err: Error) => {
-        console.error('stdin error:', err);
+        this.logger.error(`stdin error: ${err}`);
       });
     });
   }
@@ -223,6 +231,8 @@ export class MCPServer {
               log: (msg: string) => logs.push(msg),
               error: (msg: string) => logs.push(`ERROR: ${msg}`),
               warn: (msg: string) => logs.push(`WARN: ${msg}`),
+              info: (msg: string) => logs.push(`INFO: ${msg}`),
+              debug: (msg: string) => logs.push(`DEBUG: ${msg}`),
             };
 
             const runner = new WorkflowRunner(workflow, {
@@ -307,13 +317,13 @@ export class MCPServer {
           // --- Tool: get_run_logs ---
           if (toolParams.name === 'get_run_logs') {
             const { run_id } = toolParams.arguments as { run_id: string };
-            const run = this.db.getRun(run_id);
+            const run = await this.db.getRun(run_id);
 
             if (!run) {
               throw new Error(`Run ID ${run_id} not found`);
             }
 
-            const steps = this.db.getStepsByRun(run_id);
+            const steps = await this.db.getStepsByRun(run_id);
             const summary = {
               workflow: run.workflow_name,
               status: run.status,
@@ -360,7 +370,7 @@ export class MCPServer {
           // --- Tool: answer_human_input ---
           if (toolParams.name === 'answer_human_input') {
             const { run_id, input } = toolParams.arguments as { run_id: string; input: string };
-            const run = this.db.getRun(run_id);
+            const run = await this.db.getRun(run_id);
             if (!run) {
               throw new Error(`Run ID ${run_id} not found`);
             }
@@ -370,7 +380,7 @@ export class MCPServer {
             }
 
             // Find the pending or suspended step
-            const steps = this.db.getStepsByRun(run_id);
+            const steps = await this.db.getStepsByRun(run_id);
             const pendingStep = steps.find(
               (s) => s.status === 'pending' || s.status === 'suspended'
             );
@@ -403,6 +413,8 @@ export class MCPServer {
               log: (msg: string) => logs.push(msg),
               error: (msg: string) => logs.push(`ERROR: ${msg}`),
               warn: (msg: string) => logs.push(`WARN: ${msg}`),
+              info: (msg: string) => logs.push(`INFO: ${msg}`),
+              debug: (msg: string) => logs.push(`DEBUG: ${msg}`),
             };
 
             const runner = new WorkflowRunner(workflow, {
@@ -497,6 +509,8 @@ export class MCPServer {
               log: () => {},
               error: () => {},
               warn: () => {},
+              info: () => {},
+              debug: () => {},
             };
 
             const runner = new WorkflowRunner(workflow, {
@@ -507,18 +521,18 @@ export class MCPServer {
 
             const runId = runner.getRunId();
 
-            // Start the workflow asynchronously - don't await
+            // Start the workflow asynchronously
             runner.run().then(
-              (outputs) => {
-                // Update DB with success on completion (RunStatus uses 'completed')
-                this.db.updateRunStatus(runId, 'completed', outputs);
+              async (outputs) => {
+                // Update DB with success on completion
+                await this.db.updateRunStatus(runId, 'success', outputs);
               },
-              (error) => {
+              async (error) => {
                 // Update DB with failure
                 if (error instanceof WorkflowSuspendedError) {
-                  this.db.updateRunStatus(runId, 'paused');
+                  await this.db.updateRunStatus(runId, 'paused');
                 } else {
-                  this.db.updateRunStatus(
+                  await this.db.updateRunStatus(
                     runId,
                     'failed',
                     undefined,
@@ -554,7 +568,7 @@ export class MCPServer {
           // --- Tool: get_run_status ---
           if (toolParams.name === 'get_run_status') {
             const { run_id } = toolParams.arguments as { run_id: string };
-            const run = this.db.getRun(run_id);
+            const run = await this.db.getRun(run_id);
 
             if (!run) {
               throw new Error(`Run ID ${run_id} not found`);
@@ -567,7 +581,7 @@ export class MCPServer {
             };
 
             // Include outputs if completed successfully
-            if (run.status === 'completed' && run.outputs) {
+            if (run.status === 'success' && run.outputs) {
               response.outputs = JSON.parse(run.outputs);
             }