keystone-cli 0.1.1 → 0.3.0

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,349 @@ 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) {
+        // Log non-JSON lines to stderr so they show up in the terminal
+        if (line.trim()) {
+          process.stderr.write(`[MCP Server Output] ${line}\n`);
+        }
+      }
+    });
+  }
+
+  close(): void {
+    this.rl.close();
+    this.process.kill();
+  }
+}
+
+class SSETransport implements MCPTransport {
+  private url: string;
+  private headers: Record<string, string>;
+  private endpoint?: string;
+  private onMessageCallback?: (message: MCPResponse) => void;
+  private abortController: AbortController | null = null;
+  private sessionId?: string;
+
+  constructor(url: string, headers: Record<string, string> = {}) {
+    this.url = url;
+    this.headers = headers;
+  }
+
+  async connect(timeout = 60000): Promise<void> {
+    this.abortController = new AbortController();
+
+    return new Promise((resolve, reject) => {
+      const timeoutId = setTimeout(() => {
+        this.close();
+        reject(new Error(`SSE connection timeout: ${this.url}`));
+      }, timeout);
+
+      (async () => {
+        try {
+          let response = await fetch(this.url, {
+            headers: {
+              Accept: 'application/json, text/event-stream',
+              ...this.headers,
+            },
+            signal: this.abortController?.signal,
+          });
+
+          if (response.status === 405) {
+            // Some MCP servers (like GitHub) require POST to start a session
+            response = await fetch(this.url, {
+              method: 'POST',
+              headers: {
+                Accept: 'application/json, text/event-stream',
+                'Content-Type': 'application/json',
+                ...this.headers,
+              },
+              body: JSON.stringify({
+                jsonrpc: '2.0',
+                id: 'ping',
+                method: 'ping',
+              }),
+              signal: this.abortController?.signal,
+            });
+          }
+
+          if (!response.ok) {
+            clearTimeout(timeoutId);
+            reject(new Error(`SSE connection failed: ${response.status} ${response.statusText}`));
+            return;
           }
+
+          // Check for session ID in headers
+          this.sessionId =
+            response.headers.get('mcp-session-id') ||
+            response.headers.get('Mcp-Session-Id') ||
+            undefined;
+
+          const reader = response.body?.getReader();
+          if (!reader) {
+            clearTimeout(timeoutId);
+            reject(new Error('Failed to get response body reader'));
+            return;
+          }
+
+          // Process the stream in the background
+          (async () => {
+            let buffer = '';
+            const decoder = new TextDecoder();
+            let currentEvent: { event?: string; data?: string } = {};
+            let isResolved = false;
+
+            const dispatchEvent = () => {
+              if (currentEvent.data) {
+                if (currentEvent.event === 'endpoint') {
+                  this.endpoint = currentEvent.data;
+                  if (this.endpoint) {
+                    this.endpoint = new URL(this.endpoint, this.url).href;
+                  }
+                  if (!isResolved) {
+                    isResolved = true;
+                    clearTimeout(timeoutId);
+                    resolve();
+                  }
+                } else if (!currentEvent.event || currentEvent.event === 'message') {
+                  // If we get a message before an endpoint, assume the URL itself is the endpoint
+                  // (Common in some MCP over SSE implementations like GitHub's)
+                  if (!this.endpoint) {
+                    this.endpoint = this.url;
+                    if (!isResolved) {
+                      isResolved = true;
+                      clearTimeout(timeoutId);
+                      resolve();
+                    }
+                  }
+
+                  if (this.onMessageCallback && currentEvent.data) {
+                    try {
+                      const message = JSON.parse(currentEvent.data) as MCPResponse;
+                      this.onMessageCallback(message);
+                    } catch (e) {
+                      // Ignore parse errors
+                    }
+                  }
+                }
+              }
+              currentEvent = {};
+            };
+
+            try {
+              while (true) {
+                const { done, value } = await reader.read();
+                if (done) {
+                  // Dispatch any remaining data
+                  dispatchEvent();
+                  break;
+                }
+
+                buffer += decoder.decode(value, { stream: true });
+                const lines = buffer.split(/\r\n|\r|\n/);
+                buffer = lines.pop() || '';
+
+                for (const line of lines) {
+                  if (line.trim() === '') {
+                    dispatchEvent();
+                    continue;
+                  }
+
+                  if (line.startsWith('event:')) {
+                    currentEvent.event = line.substring(6).trim();
+                  } else if (line.startsWith('data:')) {
+                    const data = line.substring(5).trim();
+                    currentEvent.data = currentEvent.data ? `${currentEvent.data}\n${data}` : data;
+                  }
+                }
+              }
+
+              if (!isResolved) {
+                // If the stream ended before we resolved, but we have a session ID, we can try to resolve
+                if (this.sessionId && !this.endpoint) {
+                  this.endpoint = this.url;
+                  isResolved = true;
+                  clearTimeout(timeoutId);
+                  resolve();
+                } else {
+                  clearTimeout(timeoutId);
+                  reject(new Error('SSE stream ended before connection established'));
+                }
+              }
+            } catch (err) {
+              if ((err as Error).name !== 'AbortError' && !isResolved) {
+                clearTimeout(timeoutId);
+                reject(err);
+              }
+            }
+          })();
+        } catch (err) {
+          clearTimeout(timeoutId);
+          reject(err);
+        }
+      })();
+    });
+  }
+
+  async send(message: unknown): Promise<void> {
+    if (!this.endpoint) {
+      throw new Error('SSE transport not connected or endpoint not received');
+    }
+
+    const headers = {
+      'Content-Type': 'application/json',
+      ...this.headers,
+    };
+
+    if (this.sessionId) {
+      headers['mcp-session-id'] = this.sessionId;
+    }
+
+    const response = await fetch(this.endpoint, {
+      method: 'POST',
+      headers,
+      body: JSON.stringify(message),
+    });
+
+    if (!response.ok) {
+      const text = await response.text();
+      throw new Error(
+        `Failed to send message to MCP server: ${response.status} ${response.statusText}${
+          text ? ` - ${text}` : ''
+        }`
+      );
+    }
+
+    // Some MCP servers (like GitHub) send the response directly in the POST response as SSE
+    const contentType = response.headers.get('content-type');
+    if (contentType?.includes('text/event-stream')) {
+      const reader = response.body?.getReader();
+      if (reader) {
+        (async () => {
+          let buffer = '';
+          const decoder = new TextDecoder();
+          let currentEvent: { event?: string; data?: string } = {};
+
+          const dispatchEvent = () => {
+            if (
+              this.onMessageCallback &&
+              currentEvent.data &&
+              (!currentEvent.event || currentEvent.event === 'message')
+            ) {
+              try {
+                const message = JSON.parse(currentEvent.data) as MCPResponse;
+                this.onMessageCallback(message);
+              } catch (e) {
+                // Ignore parse errors
+              }
+            }
+            currentEvent = {};
+          };
+
+          try {
+            while (true) {
+              const { done, value } = await reader.read();
+              if (done) {
+                dispatchEvent();
+                break;
+              }
+
+              buffer += decoder.decode(value, { stream: true });
+              const lines = buffer.split(/\r\n|\r|\n/);
+              buffer = lines.pop() || '';
+
+              for (const line of lines) {
+                if (line.trim() === '') {
+                  dispatchEvent();
+                  continue;
+                }
+
+                if (line.startsWith('event:')) {
+                  currentEvent.event = line.substring(6).trim();
+                } else if (line.startsWith('data:')) {
+                  const data = line.substring(5).trim();
+                  currentEvent.data = currentEvent.data ? `${currentEvent.data}\n${data}` : data;
+                }
+              }
+            }
+          } catch (e) {
+            // Ignore stream errors
+          }
+        })();
+      }
+    }
+  }
+
+  onMessage(callback: (message: MCPResponse) => void): void {
+    this.onMessageCallback = callback;
+  }
+
+  close(): void {
+    this.abortController?.abort();
+  }
+}
+
+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 = 60000
+  ) {
+    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) || 60000;
+    }
+
+    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);
         }
-      } catch (e) {
-        // Ignore non-JSON lines
       }
     });
   }
 
+  static async createLocal(
+    command: string,
+    args: string[] = [],
+    env: Record<string, string> = {},
+    timeout = 60000
+  ): Promise<MCPClient> {
+    const transport = new StdConfigTransport(command, args, env);
+    return new MCPClient(transport, timeout);
+  }
+
+  static async createRemote(
+    url: string,
+    headers: Record<string, string> = {},
+    timeout = 60000
+  ): Promise<MCPClient> {
+    const transport = new SSETransport(url, headers);
+    await transport.connect(timeout);
+    return new MCPClient(transport, timeout);
+  }
+
   private async request(
     method: string,
     params: Record<string, unknown> = {}
@@ -78,7 +402,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 +423,7 @@ export class MCPClient {
       capabilities: {},
       clientInfo: {
         name: 'keystone-cli',
-        version: '0.1.0',
+        version: pkg.version,
       },
     });
   }
@@ -118,6 +445,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();
   });
 });