keystone-cli 0.1.1 → 0.2.0

package/src/runner/llm-executor.test.ts

@@ -3,11 +3,18 @@ import { mkdirSync, writeFileSync } from 'node:fs';
 import { join } from 'node:path';
 import type { ExpressionContext } from '../expression/evaluator';
 import type { LlmStep, Step } from '../parser/schema';
-import { AnthropicAdapter, CopilotAdapter, OpenAIAdapter } from './llm-adapter';
-import { MCPClient } from './mcp-client';
+import { ConfigLoader } from '../utils/config-loader';
+import {
+  AnthropicAdapter,
+  CopilotAdapter,
+  type LLMMessage,
+  type LLMResponse,
+  type LLMTool,
+  OpenAIAdapter,
+} from './llm-adapter';
 import { executeLlmStep } from './llm-executor';
+import { MCPClient, type MCPResponse } from './mcp-client';
 import { MCPManager } from './mcp-manager';
-import { ConfigLoader } from '../utils/config-loader';
 import type { StepResult } from './step-executor';
 
 // Mock adapters
@@ -302,9 +309,7 @@ You are a test agent.`;
     const context: ExpressionContext = { inputs: {}, steps: {} };
     const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
 
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue(
-      {} as unknown as any
-    );
+    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({} as MCPResponse);
     const listSpy = spyOn(MCPClient.prototype, 'listTools').mockResolvedValue([
       { name: 'mcp-tool', inputSchema: {} },
     ]);
@@ -317,7 +322,7 @@ You are a test agent.`;
     const originalAnthropicChatInner = AnthropicAdapter.prototype.chat;
     let toolErrorCaptured = false;
 
-    const mockChat = mock(async (messages: any[]) => {
+    const mockChat = mock(async (messages: LLMMessage[]) => {
       const toolResultMessage = messages.find((m) => m.role === 'tool');
       if (toolResultMessage?.content?.includes('Error: Tool failed')) {
         toolErrorCaptured = true;
@@ -331,7 +336,7 @@ You are a test agent.`;
           ],
         },
       };
-    }) as any;
+    }) as unknown as typeof originalOpenAIChat;
 
     OpenAIAdapter.prototype.chat = mockChat;
     CopilotAdapter.prototype.chat = mockChat;
@@ -377,21 +382,19 @@ You are a test agent.`;
     const context: ExpressionContext = { inputs: {}, steps: {} };
     const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
 
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue(
-      {} as unknown as any
-    );
+    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({} as MCPResponse);
     const listSpy = spyOn(MCPClient.prototype, 'listTools').mockResolvedValue([
       { name: 'global-tool', description: 'A global tool', inputSchema: {} },
     ]);
 
     let toolFound = false;
     const originalOpenAIChatInner = OpenAIAdapter.prototype.chat;
-    const mockChat = mock(async (_messages: any[], options: any) => {
-      if (options.tools?.some((t: any) => t.function.name === 'global-tool')) {
+    const mockChat = mock(async (_messages: LLMMessage[], options: { tools?: LLMTool[] }) => {
+      if (options.tools?.some((t: LLMTool) => t.function.name === 'global-tool')) {
         toolFound = true;
       }
       return { message: { role: 'assistant', content: 'hello' } };
-    }) as any;
+    }) as unknown as typeof originalOpenAIChat;
 
     OpenAIAdapter.prototype.chat = mockChat;
 
@@ -499,15 +502,13 @@ You are a test agent.`;
     const context: ExpressionContext = { inputs: {}, steps: {} };
     const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
 
-    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue(
-      {} as unknown as any
-    );
+    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockResolvedValue({} as MCPResponse);
     const listSpy = spyOn(MCPClient.prototype, 'listTools').mockResolvedValue([]);
 
     const originalOpenAIChatInner = OpenAIAdapter.prototype.chat;
     const mockChat = mock(async () => ({
       message: { role: 'assistant', content: 'hello' },
-    })) as any;
+    })) as unknown as typeof originalOpenAIChat;
     OpenAIAdapter.prototype.chat = mockChat;
 
     const managerSpy = spyOn(manager, 'getGlobalServers');
@@ -534,4 +535,44 @@ You are a test agent.`;
     managerSpy.mockRestore();
     ConfigLoader.clear();
   });
+
+  it('should handle object prompts by stringifying them', async () => {
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: '${{ steps.prev.output }}' as unknown as string,
+      needs: [],
+    };
+    const context: ExpressionContext = {
+      inputs: {},
+      steps: {
+        prev: { output: { key: 'value' }, status: 'success' },
+      },
+    };
+
+    let capturedPrompt = '';
+    const originalOpenAIChatInner = OpenAIAdapter.prototype.chat;
+    const mockChat = mock(async (messages: LLMMessage[]) => {
+      // console.log('MESSAGES:', JSON.stringify(messages, null, 2));
+      capturedPrompt = messages.find((m) => m.role === 'user')?.content || '';
+      return { message: { role: 'assistant', content: 'Response' } };
+    }) as unknown as typeof originalOpenAIChat;
+    OpenAIAdapter.prototype.chat = mockChat;
+    CopilotAdapter.prototype.chat = mockChat;
+    AnthropicAdapter.prototype.chat = mockChat;
+
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>
+    );
+
+    expect(capturedPrompt).toContain('"key": "value"');
+    expect(capturedPrompt).not.toContain('[object Object]');
+
+    OpenAIAdapter.prototype.chat = originalOpenAIChatInner;
+  });
 });

package/src/runner/llm-executor.ts

@@ -30,7 +30,7 @@ export async function executeLlmStep(
 
   const provider = step.provider || agent.provider;
   const model = step.model || agent.model || 'gpt-4o';
-  const prompt = ExpressionEvaluator.evaluate(step.prompt, context) as string;
+  const prompt = ExpressionEvaluator.evaluateString(step.prompt, context);
 
   const fullModelString = provider ? `${provider}:${model}` : model;
   const { adapter, resolvedModel } = getAdapter(fullModelString);

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();
   }
 }