keystone-cli 0.1.0 → 0.2.0

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

@@ -1,8 +1,8 @@
 import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
-import { MCPClient } from './mcp-client';
 import * as child_process from 'node:child_process';
 import { EventEmitter } from 'node:events';
 import { Readable, Writable } from 'node:stream';
+import { MCPClient } from './mcp-client';
 
 interface MockProcess extends EventEmitter {
   stdout: Readable;
@@ -11,112 +11,190 @@ interface MockProcess extends EventEmitter {
 }
 
 describe('MCPClient', () => {
-  let mockProcess: MockProcess;
-  let spawnSpy: ReturnType<typeof spyOn>;
+  describe('Local Transport', () => {
+    let mockProcess: MockProcess;
+    let spawnSpy: ReturnType<typeof spyOn>;
+
+    beforeEach(() => {
+      const emitter = new EventEmitter();
+      const stdout = new Readable({
+        read() {},
+      });
+      const stdin = new Writable({
+        write(_chunk, _encoding, callback) {
+          callback();
+        },
+      });
+      const kill = mock(() => {});
 
-  beforeEach(() => {
-    const emitter = new EventEmitter();
-    const stdout = new Readable({
-      read() {},
+      mockProcess = Object.assign(emitter, { stdout, stdin, kill }) as unknown as MockProcess;
+
+      spawnSpy = spyOn(child_process, 'spawn').mockReturnValue(
+        mockProcess as unknown as child_process.ChildProcess
+      );
     });
-    const stdin = new Writable({
-      write(_chunk, _encoding, callback) {
-        callback();
-      },
+
+    afterEach(() => {
+      spawnSpy.mockRestore();
     });
-    const kill = mock(() => {});
 
-    mockProcess = Object.assign(emitter, { stdout, stdin, kill }) as unknown as MockProcess;
+    it('should initialize correctly', async () => {
+      const client = await MCPClient.createLocal('node', ['server.js']);
+      const initPromise = client.initialize();
+
+      // Simulate server response
+      mockProcess.stdout.push(
+        `${JSON.stringify({
+          jsonrpc: '2.0',
+          id: 0,
+          result: { protocolVersion: '2024-11-05' },
+        })}\n`
+      );
+
+      const response = await initPromise;
+      expect(response.result?.protocolVersion).toBe('2024-11-05');
+      expect(spawnSpy).toHaveBeenCalledWith('node', ['server.js'], expect.any(Object));
+    });
 
-    spawnSpy = spyOn(child_process, 'spawn').mockReturnValue(
-      mockProcess as unknown as child_process.ChildProcess
-    );
-  });
+    it('should list tools', async () => {
+      const client = await MCPClient.createLocal('node');
+      const listPromise = client.listTools();
+
+      mockProcess.stdout.push(
+        `${JSON.stringify({
+          jsonrpc: '2.0',
+          id: 0,
+          result: {
+            tools: [{ name: 'test-tool', inputSchema: {} }],
+          },
+        })}\n`
+      );
+
+      const tools = await listPromise;
+      expect(tools).toHaveLength(1);
+      expect(tools[0].name).toBe('test-tool');
+    });
 
-  afterEach(() => {
-    spawnSpy.mockRestore();
-  });
+    it('should call tool', async () => {
+      const client = await MCPClient.createLocal('node');
+      const callPromise = client.callTool('my-tool', { arg: 1 });
 
-  it('should initialize correctly', async () => {
-    const client = new MCPClient('node', ['server.js']);
-    const initPromise = client.initialize();
-
-    // Simulate server response
-    mockProcess.stdout.push(
-      `${JSON.stringify({
-        jsonrpc: '2.0',
-        id: 0,
-        result: { protocolVersion: '2024-11-05' },
-      })}\n`
-    );
-
-    const response = await initPromise;
-    expect(response.result?.protocolVersion).toBe('2024-11-05');
-    expect(spawnSpy).toHaveBeenCalledWith('node', ['server.js'], expect.any(Object));
-  });
+      mockProcess.stdout.push(
+        `${JSON.stringify({
+          jsonrpc: '2.0',
+          id: 0,
+          result: { content: [{ type: 'text', text: 'success' }] },
+        })}\n`
+      );
+
+      const result = await callPromise;
+      expect((result as { content: Array<{ text: string }> }).content[0].text).toBe('success');
+    });
 
-  it('should list tools', async () => {
-    const client = new MCPClient('node');
-    const listPromise = client.listTools();
+    it('should handle tool call error', async () => {
+      const client = await MCPClient.createLocal('node');
+      const callPromise = client.callTool('my-tool', {});
 
-    mockProcess.stdout.push(
-      `${JSON.stringify({
-        jsonrpc: '2.0',
-        id: 0,
-        result: {
-          tools: [{ name: 'test-tool', inputSchema: {} }],
-        },
-      })}\n`
-    );
+      mockProcess.stdout.push(
+        `${JSON.stringify({
+          jsonrpc: '2.0',
+          id: 0,
+          error: { code: -32000, message: 'Tool error' },
+        })}\n`
+      );
 
-    const tools = await listPromise;
-    expect(tools).toHaveLength(1);
-    expect(tools[0].name).toBe('test-tool');
-  });
+      await expect(callPromise).rejects.toThrow(/MCP tool call failed/);
+    });
 
-  it('should call tool', async () => {
-    const client = new MCPClient('node');
-    const callPromise = client.callTool('my-tool', { arg: 1 });
+    it('should timeout on request', async () => {
+      // Set a very short timeout for testing
+      const client = await MCPClient.createLocal('node', [], {}, 10);
+      const requestPromise = client.initialize();
 
-    mockProcess.stdout.push(
-      `${JSON.stringify({
-        jsonrpc: '2.0',
-        id: 0,
-        result: { content: [{ type: 'text', text: 'success' }] },
-      })}\n`
-    );
+      // Don't push any response to stdout
+      await expect(requestPromise).rejects.toThrow(/MCP request timeout/);
+    });
 
-    const result = await callPromise;
-    expect((result as { content: Array<{ text: string }> }).content[0].text).toBe('success');
+    it('should stop the process', async () => {
+      const client = await MCPClient.createLocal('node');
+      client.stop();
+      expect(mockProcess.kill).toHaveBeenCalled();
+    });
   });
 
-  it('should handle tool call error', async () => {
-    const client = new MCPClient('node');
-    const callPromise = client.callTool('my-tool', {});
+  describe('SSE Transport', () => {
+    it('should connect and receive endpoint', async () => {
+      const mockEventSource = new EventEmitter();
+      // @ts-ignore
+      mockEventSource.addEventListener = mockEventSource.on;
+      // @ts-ignore
+      mockEventSource.close = mock(() => {});
+      // @ts-ignore
+      global.EventSource = mock(() => mockEventSource);
+
+      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse');
+
+      // Simulate endpoint event
+      mockEventSource.emit('endpoint', { data: '/endpoint' });
+
+      const client = await clientPromise;
+      expect(client).toBeDefined();
+
+      // Test sending a message
+      const fetchMock = spyOn(global, 'fetch').mockImplementation(() =>
+        Promise.resolve(new Response(JSON.stringify({ ok: true })))
+      );
+
+      const initPromise = client.initialize();
+
+      // Simulate message event (response from server)
+      mockEventSource.emit('message', {
+        data: JSON.stringify({
+          jsonrpc: '2.0',
+          id: 0,
+          result: { protocolVersion: '2024-11-05' },
+        }),
+      });
+
+      const response = await initPromise;
+      expect(response.result?.protocolVersion).toBe('2024-11-05');
+      expect(fetchMock).toHaveBeenCalledWith('http://localhost:8080/endpoint', expect.any(Object));
+
+      client.stop();
+      // @ts-ignore
+      expect(mockEventSource.close).toHaveBeenCalled();
+
+      fetchMock.mockRestore();
+      // @ts-ignore
+      global.EventSource = undefined;
+    });
 
-    mockProcess.stdout.push(
-      `${JSON.stringify({
-        jsonrpc: '2.0',
-        id: 0,
-        error: { code: -32000, message: 'Tool error' },
-      })}\n`
-    );
+    it('should handle SSE connection failure', async () => {
+      const mockEventSource = {
+        addEventListener: mock((_event: string, callback: (arg: unknown) => void) => {
+          if (_event === 'error') {
+            // Store the callback to trigger it later
+            mockEventSource.onerror = callback;
+          }
+        }),
+        onerror: null as ((arg: unknown) => void) | null,
+        close: mock(() => {}),
+      };
 
-    await expect(callPromise).rejects.toThrow(/MCP tool call failed/);
-  });
+      // @ts-ignore
+      global.EventSource = mock(() => mockEventSource);
 
-  it('should timeout on request', async () => {
-    // Set a very short timeout for testing
-    const client = new MCPClient('node', [], {}, 10);
-    const requestPromise = client.initialize();
+      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse');
 
-    // Don't push any response to stdout
-    await expect(requestPromise).rejects.toThrow(/MCP request timeout/);
-  });
+      // Trigger the onerror callback
+      if (mockEventSource.onerror) {
+        mockEventSource.onerror({ message: 'Connection failed' });
+      }
 
-  it('should stop the process', () => {
-    const client = new MCPClient('node');
-    client.stop();
-    expect(mockProcess.kill).toHaveBeenCalled();
+      await expect(clientPromise).rejects.toThrow(/SSE connection failed/);
+
+      // @ts-ignore
+      global.EventSource = undefined;
+    });
   });
 });

package/src/runner/mcp-client.ts

@@ -1,5 +1,6 @@
 import { type ChildProcess, spawn } from 'node:child_process';
 import { type Interface, createInterface } from 'node:readline';
+import pkg from '../../package.json' with { type: 'json' };
 
 interface MCPTool {
   name: string;
@@ -7,7 +8,7 @@ interface MCPTool {
   inputSchema: unknown;
 }
 
-interface MCPResponse {
+export interface MCPResponse {
   id?: number;
   result?: {
     tools?: MCPTool[];
@@ -21,20 +22,17 @@ interface MCPResponse {
   };
 }
 
-export class MCPClient {
+interface MCPTransport {
+  send(message: unknown): Promise<void>;
+  onMessage(callback: (message: MCPResponse) => void): void;
+  close(): void;
+}
+
+class StdConfigTransport implements MCPTransport {
   private process: ChildProcess;
   private rl: Interface;
-  private messageId = 0;
-  private pendingRequests = new Map<number, (response: MCPResponse) => void>();
-  private timeout: number;
 
-  constructor(
-    command: string,
-    args: string[] = [],
-    env: Record<string, string> = {},
-    timeout = 30000
-  ) {
-    this.timeout = timeout;
+  constructor(command: string, args: string[] = [], env: Record<string, string> = {}) {
     this.process = spawn(command, args, {
       env: { ...process.env, ...env },
       stdio: ['pipe', 'pipe', 'inherit'],
@@ -47,23 +45,156 @@ export class MCPClient {
     this.rl = createInterface({
       input: this.process.stdout,
     });
+  }
+
+  async send(message: unknown): Promise<void> {
+    this.process.stdin?.write(`${JSON.stringify(message)}\n`);
+  }
 
+  onMessage(callback: (message: MCPResponse) => void): void {
     this.rl.on('line', (line) => {
       try {
         const response = JSON.parse(line) as MCPResponse;
-        if (response.id !== undefined && this.pendingRequests.has(response.id)) {
-          const resolve = this.pendingRequests.get(response.id);
-          if (resolve) {
-            this.pendingRequests.delete(response.id);
-            resolve(response);
-          }
-        }
+        callback(response);
       } catch (e) {
         // Ignore non-JSON lines
       }
     });
   }
 
+  close(): void {
+    this.process.kill();
+  }
+}
+
+class SSETransport implements MCPTransport {
+  private url: string;
+  private headers: Record<string, string>;
+  private eventSource: EventSource | null = null;
+  private endpoint?: string;
+  private onMessageCallback?: (message: MCPResponse) => void;
+
+  constructor(url: string, headers: Record<string, string> = {}) {
+    this.url = url;
+    this.headers = headers;
+  }
+
+  async connect(): Promise<void> {
+    return new Promise((resolve, reject) => {
+      // @ts-ignore - Bun supports EventSource
+      this.eventSource = new EventSource(this.url, { headers: this.headers });
+
+      if (!this.eventSource) {
+        reject(new Error('Failed to create EventSource'));
+        return;
+      }
+
+      this.eventSource.addEventListener('endpoint', (event: MessageEvent) => {
+        this.endpoint = event.data;
+        if (this.endpoint?.startsWith('/')) {
+          const urlObj = new URL(this.url);
+          this.endpoint = `${urlObj.origin}${this.endpoint}`;
+        }
+        resolve();
+      });
+
+      this.eventSource.addEventListener('message', (event: MessageEvent) => {
+        if (this.onMessageCallback) {
+          try {
+            const response = JSON.parse(event.data) as MCPResponse;
+            this.onMessageCallback(response);
+          } catch (e) {
+            // Ignore
+          }
+        }
+      });
+
+      this.eventSource.onerror = (err) => {
+        const error = err as ErrorEvent;
+        reject(new Error(`SSE connection failed: ${error?.message || 'Unknown error'}`));
+      };
+    });
+  }
+
+  async send(message: unknown): Promise<void> {
+    if (!this.endpoint) {
+      throw new Error('SSE transport not connected or endpoint not received');
+    }
+
+    const response = await fetch(this.endpoint, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        ...this.headers,
+      },
+      body: JSON.stringify(message),
+    });
+
+    if (!response.ok) {
+      throw new Error(`Failed to send message to MCP server: ${response.statusText}`);
+    }
+  }
+
+  onMessage(callback: (message: MCPResponse) => void): void {
+    this.onMessageCallback = callback;
+  }
+
+  close(): void {
+    this.eventSource?.close();
+  }
+}
+
+export class MCPClient {
+  private transport: MCPTransport;
+  private messageId = 0;
+  private pendingRequests = new Map<number, (response: MCPResponse) => void>();
+  private timeout: number;
+
+  constructor(
+    transportOrCommand: MCPTransport | string,
+    timeoutOrArgs: number | string[] = [],
+    env: Record<string, string> = {},
+    timeout = 30000
+  ) {
+    if (typeof transportOrCommand === 'string') {
+      this.transport = new StdConfigTransport(transportOrCommand, timeoutOrArgs as string[], env);
+      this.timeout = timeout;
+    } else {
+      this.transport = transportOrCommand;
+      this.timeout = (timeoutOrArgs as number) || 30000;
+    }
+
+    this.transport.onMessage((response) => {
+      if (response.id !== undefined && this.pendingRequests.has(response.id)) {
+        const resolve = this.pendingRequests.get(response.id);
+        if (resolve) {
+          this.pendingRequests.delete(response.id);
+          resolve(response);
+        }
+      }
+    });
+  }
+
+  static async createLocal(
+    command: string,
+    args: string[] = [],
+    env: Record<string, string> = {},
+    timeout = 30000
+  ): Promise<MCPClient> {
+    const transport = new StdConfigTransport(command, args, env);
+    return new MCPClient(transport, timeout);
+  }
+
+  static async createRemote(
+    url: string,
+    headers: Record<string, string> = {},
+    timeout = 30000
+  ): Promise<MCPClient> {
+    const transport = new SSETransport(url, headers);
+    await transport.connect();
+    return new MCPClient(transport, timeout);
+  }
+
   private async request(
     method: string,
     params: Record<string, unknown> = {}
@@ -78,7 +209,10 @@ export class MCPClient {
 
     return new Promise((resolve, reject) => {
       this.pendingRequests.set(id, resolve);
-      this.process.stdin?.write(`${JSON.stringify(message)}\n`);
+      this.transport.send(message).catch((err) => {
+        this.pendingRequests.delete(id);
+        reject(err);
+      });
 
       // Add a timeout
       setTimeout(() => {
@@ -96,7 +230,7 @@ export class MCPClient {
       capabilities: {},
       clientInfo: {
         name: 'keystone-cli',
-        version: '0.1.0',
+        version: pkg.version,
       },
     });
   }
@@ -118,6 +252,6 @@ export class MCPClient {
   }
 
   stop() {
-    this.process.kill();
+    this.transport.close();
   }
 }

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

@@ -1,10 +1,12 @@
 import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test';
-import { MCPManager } from './mcp-manager';
-import { MCPClient } from './mcp-client';
-import { ConfigLoader } from '../utils/config-loader';
 import * as child_process from 'node:child_process';
-import { Readable, Writable } from 'node:stream';
 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 { Config } from '../parser/config-schema';
 
 describe('MCPManager', () => {
   let spawnSpy: ReturnType<typeof spyOn>;
@@ -15,7 +17,7 @@ describe('MCPManager', () => {
     const mockProcess = Object.assign(new EventEmitter(), {
       stdout: new Readable({ read() {} }),
       stdin: new Writable({
-        write(_chunk, _encoding, cb) {
+        write(_chunk, _encoding, cb: (error?: Error | null) => void) {
           cb();
         },
       }),
@@ -35,6 +37,7 @@ describe('MCPManager', () => {
     ConfigLoader.setConfig({
       mcp_servers: {
         'test-server': {
+          type: 'local',
           command: 'node',
           args: ['test.js'],
           env: { FOO: 'bar' },
@@ -43,7 +46,9 @@ describe('MCPManager', () => {
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-    });
+      storage: { retention_days: 30 },
+      workflows_directory: 'workflows',
+    } as unknown as Config);
 
     const manager = new MCPManager();
     const servers = manager.getGlobalServers();
@@ -56,13 +61,16 @@ describe('MCPManager', () => {
     ConfigLoader.setConfig({
       mcp_servers: {
         'test-server': {
+          type: 'local',
           command: 'node',
         },
       },
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-    });
+      storage: { retention_days: 30 },
+      workflows_directory: 'workflows',
+    } as unknown as Config);
 
     const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({
       result: { protocolVersion: '1.0' },
@@ -99,6 +107,7 @@ describe('MCPManager', () => {
     const manager = new MCPManager();
     const client = await manager.getClient({
       name: 'adhoc',
+      type: 'local',
       command: 'node',
     });
 
@@ -114,30 +123,79 @@ describe('MCPManager', () => {
     expect(client).toBeUndefined();
   });
 
+  it('should handle concurrent connection requests without double-spawning', async () => {
+    ConfigLoader.setConfig({
+      mcp_servers: {
+        'concurrent-server': {
+          type: 'local',
+          command: 'node',
+        },
+      },
+      providers: {},
+      model_mappings: {},
+      default_provider: 'openai',
+      storage: { retention_days: 30 },
+      workflows_directory: 'workflows',
+    } as unknown as Config);
+
+    // Mock initialize to take some time
+    let initCalls = 0;
+    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockImplementation(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();
+
+    // Fire off multiple requests concurrently
+    const p1 = manager.getClient('concurrent-server');
+    const p2 = manager.getClient('concurrent-server');
+    const p3 = manager.getClient('concurrent-server');
+
+    const [c1, c2, c3] = await Promise.all([p1, p2, p3]);
+
+    expect(c1).toBeDefined();
+    expect(c1).toBe(c2);
+    expect(c1).toBe(c3);
+    expect(initCalls).toBe(1); // Crucial: only one initialization
+
+    initSpy.mockRestore();
+  });
+
   it('should handle connection failure', async () => {
     ConfigLoader.setConfig({
       mcp_servers: {
         'fail-server': {
+          type: 'local',
           command: 'fail',
         },
       },
       providers: {},
       model_mappings: {},
       default_provider: 'openai',
-    });
-
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockRejectedValue(
-      new Error('Connection failed')
+      storage: { retention_days: 30 },
+      workflows_directory: 'workflows',
+    } 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 stopSpy = spyOn(MCPClient.prototype, 'stop').mockReturnValue(undefined);
 
     const manager = new MCPManager();
     const client = await manager.getClient('fail-server');
 
     expect(client).toBeUndefined();
-    expect(stopSpy).toHaveBeenCalled();
 
-    initSpy.mockRestore();
-    stopSpy.mockRestore();
+    createLocalSpy.mockRestore();
   });
 });