keystone-cli 0.8.0 → 1.0.1

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

@@ -3,13 +3,31 @@ import * as child_process from 'node:child_process';
 import { EventEmitter } from 'node:events';
 import { Readable, Writable } from 'node:stream';
 import { ConfigLoader } from '../utils/config-loader';
-import { MCPClient, type MCPResponse } from './mcp-client';
-import { MCPManager } from './mcp-manager';
+import type { MCPClient } from './mcp-client';
+import { type MCPClientFactory, MCPManager } from './mcp-manager';
 
 import type { Config } from '../parser/config-schema';
 
 describe('MCPManager', () => {
   let spawnSpy: ReturnType<typeof spyOn>;
+  const createMockClient = (initializeImpl?: () => Promise<unknown>) => {
+    const initialize = mock(
+      initializeImpl ||
+        (async () => ({
+          result: { protocolVersion: '1.0' },
+          jsonrpc: '2.0',
+          id: 0,
+        }))
+    );
+    const stop = mock(() => undefined);
+    const client = { initialize, stop } as unknown as MCPClient;
+    return { client, initialize, stop };
+  };
+
+  const createMockFactory = (client: MCPClient): MCPClientFactory => ({
+    createLocal: mock(async () => client),
+    createRemote: mock(async () => client),
+  });
 
   beforeEach(() => {
     ConfigLoader.clear();
@@ -46,8 +64,9 @@ describe('MCPManager', () => {
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-      storage: { retention_days: 30 },
-      workflows_directory: 'workflows',
+      storage: { retention_days: 30, redact_secrets_at_rest: true },
+      engines: { allowlist: {}, denylist: [] },
+      concurrency: { default: 10, pools: { llm: 2, shell: 5, http: 10, engine: 2 } },
     } as unknown as Config);
 
     const manager = new MCPManager();
@@ -68,43 +87,32 @@ describe('MCPManager', () => {
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-      storage: { retention_days: 30 },
-      workflows_directory: 'workflows',
+      storage: { retention_days: 30, redact_secrets_at_rest: true },
+      engines: { allowlist: {}, denylist: [] },
+      concurrency: { default: 10, pools: { llm: 2, shell: 5, http: 10, engine: 2 } },
     } as unknown as Config);
 
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({
-      result: { protocolVersion: '1.0' },
-      jsonrpc: '2.0',
-      id: 0,
-    });
-    const stopSpy = spyOn(MCPClient.prototype, 'stop').mockReturnValue(undefined);
-
-    const manager = new MCPManager();
+    const { client: mockClient, initialize, stop } = createMockClient();
+    const clientFactory = createMockFactory(mockClient);
+    const manager = new MCPManager(undefined, clientFactory);
     const client = await manager.getClient('test-server');
 
     expect(client).toBeDefined();
-    expect(initSpy).toHaveBeenCalled();
+    expect(initialize).toHaveBeenCalled();
 
     // Should reuse client
     const client2 = await manager.getClient('test-server');
     expect(client2).toBe(client);
-    expect(initSpy).toHaveBeenCalledTimes(1);
+    expect(initialize).toHaveBeenCalledTimes(1);
 
     await manager.stopAll();
-    expect(stopSpy).toHaveBeenCalled();
-
-    initSpy.mockRestore();
-    stopSpy.mockRestore();
+    expect(stop).toHaveBeenCalled();
   });
 
   it('should get client for ad-hoc server config', async () => {
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({
-      result: { protocolVersion: '1.0' },
-      jsonrpc: '2.0',
-      id: 0,
-    });
-
-    const manager = new MCPManager();
+    const { client: mockClient, initialize } = createMockClient();
+    const clientFactory = createMockFactory(mockClient);
+    const manager = new MCPManager(undefined, clientFactory);
     const client = await manager.getClient({
       name: 'adhoc',
       type: 'local',
@@ -112,9 +120,7 @@ describe('MCPManager', () => {
     });
 
     expect(client).toBeDefined();
-    expect(initSpy).toHaveBeenCalled();
-
-    initSpy.mockRestore();
+    expect(initialize).toHaveBeenCalled();
   });
 
   it('should return undefined if global server not found', async () => {
@@ -134,23 +140,24 @@ describe('MCPManager', () => {
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-      storage: { retention_days: 30 },
-      workflows_directory: 'workflows',
+      storage: { retention_days: 30, redact_secrets_at_rest: true },
+      engines: { allowlist: {}, denylist: [] },
+      concurrency: { default: 10, pools: { llm: 2, shell: 5, http: 10, engine: 2 } },
     } as unknown as Config);
 
     // Mock initialize to take some time
     let initCalls = 0;
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockImplementation(async () => {
+    const { client: mockClient, initialize } = createMockClient(async () => {
       initCalls++;
       await new Promise((resolve) => setTimeout(resolve, 50));
       return {
         result: { protocolVersion: '1.0' },
         jsonrpc: '2.0',
         id: 0,
-      } as MCPResponse;
+      };
     });
-
-    const manager = new MCPManager();
+    const clientFactory = createMockFactory(mockClient);
+    const manager = new MCPManager(undefined, clientFactory);
 
     // Fire off multiple requests concurrently
     const p1 = manager.getClient('concurrent-server');
@@ -163,8 +170,7 @@ describe('MCPManager', () => {
     expect(c1).toBe(c2);
     expect(c1).toBe(c3);
     expect(initCalls).toBe(1); // Crucial: only one initialization
-
-    initSpy.mockRestore();
+    expect(initialize).toHaveBeenCalledTimes(1);
   });
 
   it('should handle connection failure', async () => {
@@ -178,24 +184,18 @@ describe('MCPManager', () => {
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-      storage: { retention_days: 30 },
-      workflows_directory: 'workflows',
+      storage: { retention_days: 30, redact_secrets_at_rest: true },
+      engines: { allowlist: {}, denylist: [] },
+      concurrency: { default: 10, pools: { llm: 2, shell: 5, http: 10, engine: 2 } },
     } as unknown as Config);
 
-    const createLocalSpy = spyOn(MCPClient, 'createLocal').mockImplementation(
-      async (_cmd: string) => {
-        const client = Object.create(MCPClient.prototype);
-        spyOn(client, 'initialize').mockRejectedValue(new Error('Connection failed'));
-        spyOn(client, 'stop').mockReturnValue(undefined);
-        return client;
-      }
-    );
-
-    const manager = new MCPManager();
+    const { client: mockClient } = createMockClient(async () => {
+      throw new Error('Connection failed');
+    });
+    const clientFactory = createMockFactory(mockClient);
+    const manager = new MCPManager(undefined, clientFactory);
     const client = await manager.getClient('fail-server');
 
     expect(client).toBeUndefined();
-
-    createLocalSpy.mockRestore();
   });
 });

package/src/runner/mcp-manager.ts

@@ -1,6 +1,8 @@
 import { ConfigLoader } from '../utils/config-loader';
 import { ConsoleLogger, type Logger } from '../utils/logger.ts';
-import { MCPClient } from './mcp-client';
+import { MCPClient, validateRemoteUrl } from './mcp-client';
+
+export type MCPClientFactory = Pick<typeof MCPClient, 'createLocal' | 'createRemote'>;
 
 export interface MCPServerConfig {
   name: string;
@@ -21,9 +23,11 @@ export class MCPManager {
   private connectionPromises: Map<string, Promise<MCPClient | undefined>> = new Map();
   private sharedServers: Map<string, MCPServerConfig> = new Map();
   private logger: Logger;
+  private clientFactory: MCPClientFactory;
 
-  constructor(logger: Logger = new ConsoleLogger()) {
-    this.logger = logger;
+  constructor(logger?: Logger, clientFactory: MCPClientFactory = MCPClient) {
+    this.logger = logger || new ConsoleLogger();
+    this.clientFactory = clientFactory;
     this.loadGlobalConfig();
 
     // Ensure cleanup on process exit
@@ -83,6 +87,9 @@ export class MCPManager {
         if (config.type === 'remote') {
           if (!config.url) throw new Error('Remote MCP server missing URL');
 
+          // SSRF Protection: Validate URL before connecting
+          await validateRemoteUrl(config.url);
+
           const headers = { ...(config.headers || {}) };
 
           if (config.oauth) {
@@ -99,7 +106,7 @@ export class MCPManager {
             headers.Authorization = `Bearer ${token}`;
           }
 
-          client = await MCPClient.createRemote(config.url, headers, config.timeout, {
+          client = await this.clientFactory.createRemote(config.url, headers, config.timeout, {
             logger: activeLogger,
           });
         } else {
@@ -124,7 +131,7 @@ export class MCPManager {
             env.MCP_TOKEN = token;
           }
 
-          client = await MCPClient.createLocal(
+          client = await this.clientFactory.createLocal(
             config.command,
             config.args || [],
             env,

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

@@ -1,24 +1,38 @@
-import { beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
+import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
 import { WorkflowDb } from '../db/workflow-db';
 import { WorkflowParser } from '../parser/workflow-parser';
+import { ConsoleLogger } from '../utils/logger';
 import { WorkflowRegistry } from '../utils/workflow-registry';
 import { MCPServer } from './mcp-server';
 import { WorkflowSuspendedError } from './step-executor';
-import { WorkflowRunner } from './workflow-runner';
+import type { WorkflowRunner } from './workflow-runner';
 
 describe('MCPServer', () => {
   let db: WorkflowDb;
   let server: MCPServer;
+  const activeSpies: Array<{ mockRestore: () => void }> = [];
+  const trackSpy = <T extends { mockRestore: () => void }>(spy: T): T => {
+    activeSpies.push(spy);
+    return spy;
+  };
 
   beforeEach(() => {
     db = new WorkflowDb(':memory:');
     server = new MCPServer(db);
-    mock.restore();
+  });
+  afterEach(() => {
+    for (const spy of activeSpies) {
+      spy.mockRestore();
+    }
+    activeSpies.length = 0;
   });
 
-  const handleMessage = (msg: unknown) => {
+  const createServerWithRunner = (runner: WorkflowRunner) =>
+    new MCPServer(db, process.stdin, process.stdout, new ConsoleLogger(), () => runner);
+
+  const handleMessage = (msg: unknown, targetServer: MCPServer = server) => {
     // @ts-ignore
-    return server.handleMessage(msg);
+    return targetServer.handleMessage(msg);
   };
 
   it('should handle initialize request', async () => {
@@ -44,7 +58,7 @@ describe('MCPServer', () => {
   });
 
   it('should call list_workflows tool', async () => {
-    spyOn(WorkflowRegistry, 'listWorkflows').mockReturnValue([
+    trackSpy(spyOn(WorkflowRegistry, 'listWorkflows')).mockReturnValue([
       { name: 'test-wf', description: 'Test Workflow' },
     ]);
 
@@ -59,76 +73,90 @@ describe('MCPServer', () => {
   });
 
   it('should call run_workflow tool successfully', async () => {
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('test.yaml');
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('test.yaml');
     // @ts-ignore
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue({
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue({
       name: 'test-wf',
       steps: [],
     });
 
-    // Mock WorkflowRunner
     const mockRun = mock(() => Promise.resolve({ result: 'ok' }));
-    // @ts-ignore
-    spyOn(WorkflowRunner.prototype, 'run').mockImplementation(mockRun);
-
-    const response = await handleMessage({
-      jsonrpc: '2.0',
-      id: 4,
-      method: 'tools/call',
-      params: {
-        name: 'run_workflow',
-        arguments: { workflow_name: 'test-wf', inputs: {} },
+    const runner = { run: mockRun } as unknown as WorkflowRunner;
+    const testServer = createServerWithRunner(runner);
+
+    const response = await handleMessage(
+      {
+        jsonrpc: '2.0',
+        id: 4,
+        method: 'tools/call',
+        params: {
+          name: 'run_workflow',
+          arguments: { workflow_name: 'test-wf', inputs: {} },
+        },
       },
-    });
+      testServer
+    );
 
     expect(JSON.parse(response?.result?.content?.[0]?.text || '{}').status).toBe('success');
   });
 
   it('should handle run_workflow failure', async () => {
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('test.yaml');
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('test.yaml');
     // @ts-ignore
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue({
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue({
       name: 'test-wf',
       steps: [],
     });
 
-    spyOn(WorkflowRunner.prototype, 'run').mockRejectedValue(new Error('workflow failed'));
-
-    const response = await handleMessage({
-      jsonrpc: '2.0',
-      id: 5,
-      method: 'tools/call',
-      params: {
-        name: 'run_workflow',
-        arguments: { workflow_name: 'test-wf' },
+    const runner = {
+      run: mock(() => Promise.reject(new Error('workflow failed'))),
+    } as unknown as WorkflowRunner;
+    const testServer = createServerWithRunner(runner);
+
+    const response = await handleMessage(
+      {
+        jsonrpc: '2.0',
+        id: 5,
+        method: 'tools/call',
+        params: {
+          name: 'run_workflow',
+          arguments: { workflow_name: 'test-wf' },
+        },
       },
-    });
+      testServer
+    );
 
     expect(response?.result?.isError).toBe(true);
     expect(response?.result?.content?.[0]?.text).toContain('Workflow failed');
   });
 
   it('should handle workflow suspension in run_workflow', async () => {
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('test.yaml');
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('test.yaml');
     // @ts-ignore
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue({
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue({
       name: 'test-wf',
       steps: [],
     });
 
     const suspendedError = new WorkflowSuspendedError('Input needed', 'step1', 'text');
-    spyOn(WorkflowRunner.prototype, 'run').mockRejectedValue(suspendedError);
-    spyOn(WorkflowRunner.prototype, 'getRunId').mockReturnValue('run123');
-
-    const response = await handleMessage({
-      jsonrpc: '2.0',
-      id: 6,
-      method: 'tools/call',
-      params: {
-        name: 'run_workflow',
-        arguments: { workflow_name: 'test-wf' },
+    const runner = {
+      run: mock(() => Promise.reject(suspendedError)),
+      getRunId: mock(() => 'run123'),
+    } as unknown as WorkflowRunner;
+    const testServer = createServerWithRunner(runner);
+
+    const response = await handleMessage(
+      {
+        jsonrpc: '2.0',
+        id: 6,
+        method: 'tools/call',
+        params: {
+          name: 'run_workflow',
+          arguments: { workflow_name: 'test-wf' },
+        },
       },
-    });
+      testServer
+    );
 
     const result = JSON.parse(response?.result?.content?.[0]?.text || '{}');
     expect(result.status).toBe('paused');
@@ -142,26 +170,29 @@ describe('MCPServer', () => {
     await db.updateRunStatus(runId, 'paused');
     await db.createStep('step-exec-1', runId, 's1');
 
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('test.yaml');
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('test.yaml');
     // @ts-ignore
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue({
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue({
       name: 'test-wf',
       steps: [{ id: 's1', type: 'human' }],
     });
 
     const mockRun = mock(() => Promise.resolve({ result: 'resumed' }));
-    // @ts-ignore
-    spyOn(WorkflowRunner.prototype, 'run').mockImplementation(mockRun);
-
-    const response = await handleMessage({
-      jsonrpc: '2.0',
-      id: 7,
-      method: 'tools/call',
-      params: {
-        name: 'answer_human_input',
-        arguments: { run_id: runId, input: 'my response' },
+    const runner = { run: mockRun } as unknown as WorkflowRunner;
+    const testServer = createServerWithRunner(runner);
+
+    const response = await handleMessage(
+      {
+        jsonrpc: '2.0',
+        id: 7,
+        method: 'tools/call',
+        params: {
+          name: 'answer_human_input',
+          arguments: { run_id: runId, input: 'my response' },
+        },
       },
-    });
+      testServer
+    );
 
     expect(JSON.parse(response?.result?.content?.[0]?.text || '{}').status).toBe('success');
 
@@ -223,8 +254,8 @@ describe('MCPServer', () => {
     // Create a new server for this test to use the streams
     const testServer = new MCPServer(db, input, outputStream);
 
-    const writeSpy = spyOn(outputStream, 'write').mockImplementation(() => true);
-    const consoleSpy = spyOn(console, 'error').mockImplementation(() => {});
+    const writeSpy = trackSpy(spyOn(outputStream, 'write')).mockImplementation(() => true);
+    const consoleSpy = trackSpy(spyOn(console, 'error')).mockImplementation(() => {});
 
     const startPromise = testServer.start();
 
@@ -245,36 +276,37 @@ describe('MCPServer', () => {
 
     input.end();
     await startPromise;
-
-    writeSpy.mockRestore();
-    consoleSpy.mockRestore();
   });
 
   it('should call start_workflow tool and return immediately', async () => {
-    spyOn(WorkflowRegistry, 'resolvePath').mockReturnValue('test.yaml');
+    trackSpy(spyOn(WorkflowRegistry, 'resolvePath')).mockReturnValue('test.yaml');
     // @ts-ignore
-    spyOn(WorkflowParser, 'loadWorkflow').mockReturnValue({
+    trackSpy(spyOn(WorkflowParser, 'loadWorkflow')).mockReturnValue({
       name: 'test-wf',
       steps: [],
     });
 
-    // Mock WorkflowRunner - simulate a slow workflow
     const mockRun = mock(
       () => new Promise((resolve) => setTimeout(() => resolve({ result: 'ok' }), 100))
     );
-    // @ts-ignore
-    spyOn(WorkflowRunner.prototype, 'run').mockImplementation(mockRun);
-    spyOn(WorkflowRunner.prototype, 'getRunId').mockReturnValue('async-run-123');
-
-    const response = await handleMessage({
-      jsonrpc: '2.0',
-      id: 10,
-      method: 'tools/call',
-      params: {
-        name: 'start_workflow',
-        arguments: { workflow_name: 'test-wf', inputs: {} },
+    const runner = {
+      run: mockRun,
+      getRunId: mock(() => 'async-run-123'),
+    } as unknown as WorkflowRunner;
+    const testServer = createServerWithRunner(runner);
+
+    const response = await handleMessage(
+      {
+        jsonrpc: '2.0',
+        id: 10,
+        method: 'tools/call',
+        params: {
+          name: 'start_workflow',
+          arguments: { workflow_name: 'test-wf', inputs: {} },
+        },
       },
-    });
+      testServer
+    );
 
     const result = JSON.parse(response?.result?.content?.[0]?.text || '{}');
     expect(result.status).toBe('running');
@@ -366,4 +398,11 @@ describe('MCPServer', () => {
 
     expect(response?.error?.message).toContain('not found');
   });
+
+  it('should close database connection when stop is called', () => {
+    const dbCloseSpy = spyOn(db, 'close');
+    server.stop();
+    expect(dbCloseSpy).toHaveBeenCalled();
+    dbCloseSpy.mockRestore();
+  });
 });

package/src/runner/mcp-server.ts

@@ -16,22 +16,31 @@ interface MCPMessage {
   id?: string | number;
 }
 
+type WorkflowRunnerFactory = (
+  workflow: ConstructorParameters<typeof WorkflowRunner>[0],
+  options: ConstructorParameters<typeof WorkflowRunner>[1]
+) => WorkflowRunner;
+
 export class MCPServer {
   private db: WorkflowDb;
   private input: Readable;
   private output: Writable;
   private logger: Logger;
+  private runnerFactory: WorkflowRunnerFactory;
 
   constructor(
     db?: WorkflowDb,
     input: Readable = process.stdin,
     output: Writable = process.stdout,
-    logger: Logger = new ConsoleLogger()
+    logger: Logger = new ConsoleLogger(),
+    runnerFactory: WorkflowRunnerFactory = (workflow, options) =>
+      new WorkflowRunner(workflow, options)
   ) {
     this.db = db || new WorkflowDb();
     this.input = input;
     this.output = output;
     this.logger = logger;
+    this.runnerFactory = runnerFactory;
   }
 
   async start() {
@@ -235,7 +244,7 @@ export class MCPServer {
               debug: (msg: string) => logs.push(`DEBUG: ${msg}`),
             };
 
-            const runner = new WorkflowRunner(workflow, {
+            const runner = this.runnerFactory(workflow, {
               inputs,
               logger,
               preventExit: true,
@@ -417,7 +426,7 @@ export class MCPServer {
               debug: (msg: string) => logs.push(`DEBUG: ${msg}`),
             };
 
-            const runner = new WorkflowRunner(workflow, {
+            const runner = this.runnerFactory(workflow, {
               resumeRunId: run_id,
               resumeInputs: { [pendingStep.step_id]: { __answer: output } },
               logger,
@@ -513,7 +522,7 @@ export class MCPServer {
               debug: () => {},
             };
 
-            const runner = new WorkflowRunner(workflow, {
+            const runner = this.runnerFactory(workflow, {
               inputs: inputs || {},
               logger,
               preventExit: true,