keystone-cli 0.2.0 → 0.3.1

package/src/runner/mcp-client.ts

@@ -57,12 +57,16 @@ class StdConfigTransport implements MCPTransport {
         const response = JSON.parse(line) as MCPResponse;
         callback(response);
       } catch (e) {
-        // Ignore non-JSON lines
+        // 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();
   }
 }
@@ -70,49 +74,168 @@ class StdConfigTransport implements MCPTransport {
 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;
+  private abortController: AbortController | null = null;
+  private sessionId?: string;
 
   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 });
+  async connect(timeout = 60000): Promise<void> {
+    this.abortController = new AbortController();
 
-      if (!this.eventSource) {
-        reject(new Error('Failed to create EventSource'));
-        return;
-      }
+    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,
+            });
+          }
 
-      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();
-      });
+          if (!response.ok) {
+            clearTimeout(timeoutId);
+            reject(new Error(`SSE connection failed: ${response.status} ${response.statusText}`));
+            return;
+          }
 
-      this.eventSource.addEventListener('message', (event: MessageEvent) => {
-        if (this.onMessageCallback) {
-          try {
-            const response = JSON.parse(event.data) as MCPResponse;
-            this.onMessageCallback(response);
-          } catch (e) {
-            // Ignore
+          // 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;
           }
-        }
-      });
 
-      this.eventSource.onerror = (err) => {
-        const error = err as ErrorEvent;
-        reject(new Error(`SSE connection failed: ${error?.message || 'Unknown error'}`));
-      };
+          // 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);
+        }
+      })();
     });
   }
 
@@ -121,17 +244,87 @@ class SSETransport implements MCPTransport {
       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: {
-        'Content-Type': 'application/json',
-        ...this.headers,
-      },
+      headers,
       body: JSON.stringify(message),
     });
 
     if (!response.ok) {
-      throw new Error(`Failed to send message to MCP server: ${response.statusText}`);
+      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
+          }
+        })();
+      }
     }
   }
 
@@ -140,7 +333,7 @@ class SSETransport implements MCPTransport {
   }
 
   close(): void {
-    this.eventSource?.close();
+    this.abortController?.abort();
   }
 }
 
@@ -154,14 +347,14 @@ export class MCPClient {
     transportOrCommand: MCPTransport | string,
     timeoutOrArgs: number | string[] = [],
     env: Record<string, string> = {},
-    timeout = 30000
+    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) || 30000;
+      this.timeout = (timeoutOrArgs as number) || 60000;
     }
 
     this.transport.onMessage((response) => {
@@ -179,7 +372,7 @@ export class MCPClient {
     command: string,
     args: string[] = [],
     env: Record<string, string> = {},
-    timeout = 30000
+    timeout = 60000
   ): Promise<MCPClient> {
     const transport = new StdConfigTransport(command, args, env);
     return new MCPClient(transport, timeout);
@@ -188,10 +381,10 @@ export class MCPClient {
   static async createRemote(
     url: string,
     headers: Record<string, string> = {},
-    timeout = 30000
+    timeout = 60000
   ): Promise<MCPClient> {
     const transport = new SSETransport(url, headers);
-    await transport.connect();
+    await transport.connect(timeout);
     return new MCPClient(transport, timeout);
   }
 

package/src/runner/mcp-manager.ts

@@ -10,6 +10,9 @@ export interface MCPServerConfig {
   env?: Record<string, string>;
   url?: string;
   headers?: Record<string, string>;
+  oauth?: {
+    scope?: string;
+  };
 }
 
 export class MCPManager {
@@ -70,10 +73,47 @@ export class MCPManager {
       try {
         if (config.type === 'remote') {
           if (!config.url) throw new Error('Remote MCP server missing URL');
-          client = await MCPClient.createRemote(config.url, config.headers || {});
+
+          const headers = { ...(config.headers || {}) };
+
+          if (config.oauth) {
+            const { AuthManager } = await import('../utils/auth-manager');
+            const auth = AuthManager.load();
+            const token = auth.mcp_tokens?.[config.name]?.access_token;
+
+            if (!token) {
+              throw new Error(
+                `MCP server ${config.name} requires OAuth. Please run "keystone mcp login ${config.name}" first.`
+              );
+            }
+
+            headers.Authorization = `Bearer ${token}`;
+          }
+
+          client = await MCPClient.createRemote(config.url, headers);
         } else {
           if (!config.command) throw new Error('Local MCP server missing command');
-          client = await MCPClient.createLocal(config.command, config.args || [], config.env || {});
+
+          const env = { ...(config.env || {}) };
+
+          if (config.oauth) {
+            const { AuthManager } = await import('../utils/auth-manager');
+            const auth = AuthManager.load();
+            const token = auth.mcp_tokens?.[config.name]?.access_token;
+
+            if (!token) {
+              throw new Error(
+                `MCP server ${config.name} requires OAuth. Please run "keystone mcp login ${config.name}" first.`
+              );
+            }
+
+            // Pass token to the local proxy via environment variables
+            // Most proxies expect AUTHORIZATION or MCP_TOKEN
+            env.AUTHORIZATION = `Bearer ${token}`;
+            env.MCP_TOKEN = token;
+          }
+
+          client = await MCPClient.createLocal(config.command, config.args || [], env);
         }
 
         await client.initialize();

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

@@ -28,7 +28,7 @@ describe('MCPServer', () => {
       method: 'initialize',
     });
 
-    expect(response.result.serverInfo.name).toBe('keystone-mcp');
+    expect(response?.result?.serverInfo?.name).toBe('keystone-mcp');
   });
 
   it('should list tools', async () => {
@@ -38,9 +38,9 @@ describe('MCPServer', () => {
       method: 'tools/list',
     });
 
-    expect(response.result.tools).toHaveLength(5);
+    expect(response?.result?.tools).toHaveLength(5);
     // @ts-ignore
-    expect(response.result.tools.map((t) => t.name)).toContain('run_workflow');
+    expect(response?.result?.tools?.map((t) => t.name)).toContain('run_workflow');
   });
 
   it('should call list_workflows tool', async () => {
@@ -55,7 +55,7 @@ describe('MCPServer', () => {
       params: { name: 'list_workflows', arguments: {} },
     });
 
-    expect(response.result.content[0].text).toContain('test-wf');
+    expect(response?.result?.content?.[0]?.text).toContain('test-wf');
   });
 
   it('should call run_workflow tool successfully', async () => {
@@ -104,8 +104,8 @@ describe('MCPServer', () => {
       },
     });
 
-    expect(response.result.isError).toBe(true);
-    expect(response.result.content[0].text).toContain('Workflow failed');
+    expect(response?.result?.isError).toBe(true);
+    expect(response?.result?.content?.[0]?.text).toContain('Workflow failed');
   });
 
   it('should handle workflow suspension in run_workflow', async () => {
@@ -130,7 +130,7 @@ describe('MCPServer', () => {
       },
     });
 
-    const result = JSON.parse(response.result.content[0].text);
+    const result = JSON.parse(response?.result?.content?.[0]?.text);
     expect(result.status).toBe('paused');
     expect(result.run_id).toBe('run123');
     expect(result.message).toBe('Input needed');
@@ -187,7 +187,7 @@ describe('MCPServer', () => {
       params: { name: 'get_run_logs', arguments: { run_id: runId } },
     });
 
-    const summary = JSON.parse(response.result.content[0].text);
+    const summary = JSON.parse(response?.result?.content?.[0]?.text);
     expect(summary.workflow).toBe('test-wf');
     expect(summary.steps).toHaveLength(1);
     expect(summary.steps[0].step).toBe('s1');
@@ -202,7 +202,7 @@ describe('MCPServer', () => {
       params: { name: 'unknown_tool', arguments: {} },
     });
 
-    expect(response.error.message).toContain('Unknown tool');
+    expect(response?.error?.message).toContain('Unknown tool');
   });
 
   it('should handle unknown method', async () => {
@@ -212,14 +212,21 @@ describe('MCPServer', () => {
       method: 'unknown_method',
     });
 
-    expect(response.error.message).toContain('Method not found');
+    expect(response?.error?.message).toContain('Method not found');
   });
 
   it('should start and handle messages from stdin', async () => {
-    const writeSpy = spyOn(process.stdout, 'write').mockImplementation(() => true);
+    const { PassThrough } = await import('node:stream');
+    const input = new PassThrough();
+    const outputStream = new PassThrough();
+
+    // 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 startPromise = server.start();
+    const startPromise = testServer.start();
 
     // Simulate stdin data
     const message = {
@@ -227,16 +234,16 @@ describe('MCPServer', () => {
       id: 9,
       method: 'initialize',
     };
-    process.stdin.emit('data', Buffer.from(`${JSON.stringify(message)}\n`));
+    input.write(`${JSON.stringify(message)}\n`);
 
     // Wait for async processing
-    await new Promise((resolve) => setTimeout(resolve, 50));
+    await new Promise((resolve) => setTimeout(resolve, 100));
 
     expect(writeSpy).toHaveBeenCalled();
     const output = JSON.parse(writeSpy.mock.calls[0][0] as string);
     expect(output.id).toBe(9);
 
-    process.stdin.emit('close');
+    input.end();
     await startPromise;
 
     writeSpy.mockRestore();

package/src/runner/mcp-server.ts

@@ -1,4 +1,5 @@
 import * as readline from 'node:readline';
+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';
@@ -16,14 +17,18 @@ interface MCPMessage {
 
 export class MCPServer {
   private db: WorkflowDb;
+  private input: Readable;
+  private output: Writable;
 
-  constructor(db?: WorkflowDb) {
+  constructor(db?: WorkflowDb, input: Readable = process.stdin, output: Writable = process.stdout) {
     this.db = db || new WorkflowDb();
+    this.input = input;
+    this.output = output;
   }
 
   async start() {
     const rl = readline.createInterface({
-      input: process.stdin,
+      input: this.input,
       terminal: false,
     });
 
@@ -35,7 +40,7 @@ export class MCPServer {
           const message = JSON.parse(line) as MCPMessage;
           const response = await this.handleMessage(message);
           if (response) {
-            process.stdout.write(`${JSON.stringify(response)}\n`);
+            this.output.write(`${JSON.stringify(response)}\n`);
           }
         } catch (error) {
           console.error('Error handling MCP message:', error);
@@ -46,6 +51,11 @@ export class MCPServer {
         this.stop();
         resolve();
       });
+
+      // Handle stream errors
+      this.input.on('error', (err: Error) => {
+        console.error('stdin error:', err);
+      });
     });
   }
 
@@ -333,7 +343,14 @@ export class MCPServer {
             }
 
             // Fulfill the step in the DB
-            const output = input === 'confirm' ? true : input;
+            let output: unknown = input;
+            const lowerInput = input.trim().toLowerCase();
+            if (lowerInput === 'confirm' || lowerInput === 'y' || lowerInput === 'yes' || lowerInput === '') {
+              output = true;
+            } else if (lowerInput === 'n' || lowerInput === 'no') {
+              output = false;
+            }
+
             await this.db.completeStep(pendingStep.id, 'success', output);
 
             // Resume the workflow