keystone-cli 0.1.0

package/src/runner/llm-executor.ts

@@ -0,0 +1,256 @@
+import { join } from 'node:path';
+import type { ExpressionContext } from '../expression/evaluator';
+import { ExpressionEvaluator } from '../expression/evaluator';
+import { parseAgent, resolveAgentPath } from '../parser/agent-parser';
+import type { AgentTool, LlmStep, Step } from '../parser/schema';
+import { type LLMMessage, getAdapter } from './llm-adapter';
+import { MCPClient } from './mcp-client';
+import type { MCPManager } from './mcp-manager';
+import type { StepResult } from './step-executor';
+import type { Logger } from './workflow-runner';
+
+interface ToolDefinition {
+  name: string;
+  description?: string;
+  parameters: unknown;
+  source: 'agent' | 'step' | 'mcp';
+  execution?: Step;
+  mcpClient?: MCPClient;
+}
+
+export async function executeLlmStep(
+  step: LlmStep,
+  context: ExpressionContext,
+  executeStepFn: (step: Step, context: ExpressionContext) => Promise<StepResult>,
+  logger: Logger = console,
+  mcpManager?: MCPManager
+): Promise<StepResult> {
+  const agentPath = resolveAgentPath(step.agent);
+  const agent = parseAgent(agentPath);
+
+  const provider = step.provider || agent.provider;
+  const model = step.model || agent.model || 'gpt-4o';
+  const prompt = ExpressionEvaluator.evaluate(step.prompt, context) as string;
+
+  const fullModelString = provider ? `${provider}:${model}` : model;
+  const { adapter, resolvedModel } = getAdapter(fullModelString);
+
+  // Inject schema instructions if present
+  let systemPrompt = agent.systemPrompt;
+  if (step.schema) {
+    systemPrompt += `\n\nIMPORTANT: You must output valid JSON that matches the following schema:\n${JSON.stringify(step.schema, null, 2)}`;
+  }
+
+  const messages: LLMMessage[] = [
+    { role: 'system', content: systemPrompt },
+    { role: 'user', content: prompt },
+  ];
+
+  const localMcpClients: MCPClient[] = [];
+  const allTools: ToolDefinition[] = [];
+
+  try {
+    // 1. Add agent tools
+    for (const tool of agent.tools) {
+      allTools.push({
+        name: tool.name,
+        description: tool.description,
+        parameters: tool.parameters || {
+          type: 'object',
+          properties: {},
+          additionalProperties: true,
+        },
+        source: 'agent',
+        execution: tool.execution,
+      });
+    }
+
+    // 2. Add step tools
+    if (step.tools) {
+      for (const tool of step.tools) {
+        allTools.push({
+          name: tool.name,
+          description: tool.description,
+          parameters: tool.parameters || {
+            type: 'object',
+            properties: {},
+            additionalProperties: true,
+          },
+          source: 'step',
+          execution: tool.execution,
+        });
+      }
+    }
+
+    // 3. Add MCP tools
+    const mcpServersToConnect = [...(step.mcpServers || [])];
+    if (step.useGlobalMcp && mcpManager) {
+      const globalServers = mcpManager.getGlobalServers();
+      for (const globalServer of globalServers) {
+        // Only add if not already explicitly listed
+        const alreadyListed = mcpServersToConnect.some((s) =>
+          typeof s === 'string' ? s === globalServer.name : s.name === globalServer.name
+        );
+        if (!alreadyListed) {
+          mcpServersToConnect.push(globalServer);
+        }
+      }
+    }
+
+    if (mcpServersToConnect.length > 0) {
+      for (const server of mcpServersToConnect) {
+        let client: MCPClient | undefined;
+
+        if (mcpManager) {
+          client = await mcpManager.getClient(server, logger);
+        } else {
+          // Fallback if no manager (should not happen in normal workflow run)
+          if (typeof server === 'string') {
+            logger.error(`  ✗ Cannot reference global MCP server '${server}' without MCPManager`);
+            continue;
+          }
+          logger.log(`  🔌 Connecting to MCP server: ${server.name}`);
+          client = new MCPClient(server.command, server.args, server.env);
+          try {
+            await client.initialize();
+            localMcpClients.push(client);
+          } catch (error) {
+            logger.error(
+              `  ✗ Failed to connect to MCP server ${server.name}: ${error instanceof Error ? error.message : String(error)}`
+            );
+            client.stop();
+            client = undefined;
+          }
+        }
+
+        if (client) {
+          const mcpTools = await client.listTools();
+          for (const tool of mcpTools) {
+            allTools.push({
+              name: tool.name,
+              description: tool.description,
+              parameters: tool.inputSchema,
+              source: 'mcp',
+              mcpClient: client,
+            });
+          }
+        }
+      }
+    }
+
+    const llmTools = allTools.map((t) => ({
+      type: 'function',
+      function: {
+        name: t.name,
+        description: t.description,
+        parameters: t.parameters,
+      },
+    }));
+
+    // ReAct Loop
+    let iterations = 0;
+    const maxIterations = step.maxIterations || 10;
+
+    while (iterations < maxIterations) {
+      iterations++;
+
+      const response = await adapter.chat(messages, {
+        model: resolvedModel,
+        tools: llmTools.length > 0 ? llmTools : undefined,
+      });
+
+      const { message } = response;
+      messages.push(message);
+
+      if (message.content && !step.schema) {
+        logger.log(`\n${message.content}`);
+      }
+
+      if (!message.tool_calls || message.tool_calls.length === 0) {
+        let output = message.content;
+
+        // If schema is defined, attempt to parse JSON
+        if (step.schema && typeof output === 'string') {
+          try {
+            // Attempt to extract JSON if wrapped in markdown code blocks or just finding the first {
+            const jsonMatch =
+              output.match(/```(?:json)?\s*([\s\S]*?)\s*```/i) || output.match(/\{[\s\S]*\}/);
+            const jsonStr = jsonMatch ? jsonMatch[1] || jsonMatch[0] : output;
+            output = JSON.parse(jsonStr);
+          } catch (e) {
+            throw new Error(
+              `Failed to parse LLM output as JSON matching schema: ${e instanceof Error ? e.message : String(e)}\nOutput: ${output}`
+            );
+          }
+        }
+
+        return {
+          output,
+          status: 'success',
+        };
+      }
+
+      // Execute tools
+      for (const toolCall of message.tool_calls) {
+        logger.log(`  🛠️  Tool Call: ${toolCall.function.name}`);
+        const toolInfo = allTools.find((t) => t.name === toolCall.function.name);
+
+        if (!toolInfo) {
+          messages.push({
+            role: 'tool',
+            tool_call_id: toolCall.id,
+            name: toolCall.function.name,
+            content: `Error: Tool ${toolCall.function.name} not found`,
+          });
+          continue;
+        }
+
+        const args = JSON.parse(toolCall.function.arguments);
+
+        if (toolInfo.source === 'mcp' && toolInfo.mcpClient) {
+          try {
+            const result = await toolInfo.mcpClient.callTool(toolInfo.name, args);
+            messages.push({
+              role: 'tool',
+              tool_call_id: toolCall.id,
+              name: toolCall.function.name,
+              content: JSON.stringify(result),
+            });
+          } catch (error) {
+            messages.push({
+              role: 'tool',
+              tool_call_id: toolCall.id,
+              name: toolCall.function.name,
+              content: `Error: ${error instanceof Error ? error.message : String(error)}`,
+            });
+          }
+        } else if (toolInfo.execution) {
+          // Execute the tool as a step
+          const toolContext: ExpressionContext = {
+            ...context,
+            args,
+          };
+
+          const result = await executeStepFn(toolInfo.execution, toolContext);
+
+          messages.push({
+            role: 'tool',
+            tool_call_id: toolCall.id,
+            name: toolCall.function.name,
+            content:
+              result.status === 'success'
+                ? JSON.stringify(result.output)
+                : `Error: ${result.error}`,
+          });
+        }
+      }
+    }
+
+    throw new Error('Max ReAct iterations reached');
+  } finally {
+    // Cleanup LOCAL MCP clients only. Shared clients are managed by MCPManager.
+    for (const client of localMcpClients) {
+      client.stop();
+    }
+  }
+}

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

@@ -0,0 +1,122 @@
+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';
+
+interface MockProcess extends EventEmitter {
+  stdout: Readable;
+  stdin: Writable;
+  kill: ReturnType<typeof mock>;
+}
+
+describe('MCPClient', () => {
+  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(() => {});
+
+    mockProcess = Object.assign(emitter, { stdout, stdin, kill }) as unknown as MockProcess;
+
+    spawnSpy = spyOn(child_process, 'spawn').mockReturnValue(
+      mockProcess as unknown as child_process.ChildProcess
+    );
+  });
+
+  afterEach(() => {
+    spawnSpy.mockRestore();
+  });
+
+  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));
+  });
+
+  it('should list tools', async () => {
+    const client = new MCPClient('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');
+  });
+
+  it('should call tool', async () => {
+    const client = new MCPClient('node');
+    const callPromise = client.callTool('my-tool', { arg: 1 });
+
+    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 handle tool call error', async () => {
+    const client = new MCPClient('node');
+    const callPromise = client.callTool('my-tool', {});
+
+    mockProcess.stdout.push(
+      `${JSON.stringify({
+        jsonrpc: '2.0',
+        id: 0,
+        error: { code: -32000, message: 'Tool error' },
+      })}\n`
+    );
+
+    await expect(callPromise).rejects.toThrow(/MCP tool call failed/);
+  });
+
+  it('should timeout on request', async () => {
+    // Set a very short timeout for testing
+    const client = new MCPClient('node', [], {}, 10);
+    const requestPromise = client.initialize();
+
+    // Don't push any response to stdout
+    await expect(requestPromise).rejects.toThrow(/MCP request timeout/);
+  });
+
+  it('should stop the process', () => {
+    const client = new MCPClient('node');
+    client.stop();
+    expect(mockProcess.kill).toHaveBeenCalled();
+  });
+});

package/src/runner/mcp-client.ts

@@ -0,0 +1,123 @@
+import { type ChildProcess, spawn } from 'node:child_process';
+import { type Interface, createInterface } from 'node:readline';
+
+interface MCPTool {
+  name: string;
+  description?: string;
+  inputSchema: unknown;
+}
+
+interface MCPResponse {
+  id?: number;
+  result?: {
+    tools?: MCPTool[];
+    content?: Array<{ type: string; text: string }>;
+    [key: string]: unknown;
+  };
+  error?: {
+    code: number;
+    message: string;
+    data?: unknown;
+  };
+}
+
+export class MCPClient {
+  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;
+    this.process = spawn(command, args, {
+      env: { ...process.env, ...env },
+      stdio: ['pipe', 'pipe', 'inherit'],
+    });
+
+    if (!this.process.stdout || !this.process.stdin) {
+      throw new Error('Failed to start MCP server: stdio not available');
+    }
+
+    this.rl = createInterface({
+      input: this.process.stdout,
+    });
+
+    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);
+          }
+        }
+      } catch (e) {
+        // Ignore non-JSON lines
+      }
+    });
+  }
+
+  private async request(
+    method: string,
+    params: Record<string, unknown> = {}
+  ): Promise<MCPResponse> {
+    const id = this.messageId++;
+    const message = {
+      jsonrpc: '2.0',
+      id,
+      method,
+      params,
+    };
+
+    return new Promise((resolve, reject) => {
+      this.pendingRequests.set(id, resolve);
+      this.process.stdin?.write(`${JSON.stringify(message)}\n`);
+
+      // Add a timeout
+      setTimeout(() => {
+        if (this.pendingRequests.has(id)) {
+          this.pendingRequests.delete(id);
+          reject(new Error(`MCP request timeout: ${method}`));
+        }
+      }, this.timeout);
+    });
+  }
+
+  async initialize() {
+    return this.request('initialize', {
+      protocolVersion: '2024-11-05',
+      capabilities: {},
+      clientInfo: {
+        name: 'keystone-cli',
+        version: '0.1.0',
+      },
+    });
+  }
+
+  async listTools(): Promise<MCPTool[]> {
+    const response = await this.request('tools/list');
+    return response.result?.tools || [];
+  }
+
+  async callTool(name: string, args: Record<string, unknown>): Promise<unknown> {
+    const response = await this.request('tools/call', {
+      name,
+      arguments: args,
+    });
+    if (response.error) {
+      throw new Error(`MCP tool call failed: ${JSON.stringify(response.error)}`);
+    }
+    return response.result;
+  }
+
+  stop() {
+    this.process.kill();
+  }
+}

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

@@ -0,0 +1,143 @@
+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';
+
+describe('MCPManager', () => {
+  let spawnSpy: ReturnType<typeof spyOn>;
+
+  beforeEach(() => {
+    ConfigLoader.clear();
+
+    const mockProcess = Object.assign(new EventEmitter(), {
+      stdout: new Readable({ read() {} }),
+      stdin: new Writable({
+        write(_chunk, _encoding, cb) {
+          cb();
+        },
+      }),
+      kill: mock(() => {}),
+    });
+    spawnSpy = spyOn(child_process, 'spawn').mockReturnValue(
+      mockProcess as unknown as child_process.ChildProcess
+    );
+  });
+
+  afterEach(() => {
+    ConfigLoader.clear();
+    spawnSpy.mockRestore();
+  });
+
+  it('should load global config on initialization', () => {
+    ConfigLoader.setConfig({
+      mcp_servers: {
+        'test-server': {
+          command: 'node',
+          args: ['test.js'],
+          env: { FOO: 'bar' },
+        },
+      },
+      providers: {},
+      model_mappings: {},
+      default_provider: 'openai',
+    });
+
+    const manager = new MCPManager();
+    const servers = manager.getGlobalServers();
+    expect(servers).toHaveLength(1);
+    expect(servers[0].name).toBe('test-server');
+    expect(servers[0].command).toBe('node');
+  });
+
+  it('should get client for global server', async () => {
+    ConfigLoader.setConfig({
+      mcp_servers: {
+        'test-server': {
+          command: 'node',
+        },
+      },
+      providers: {},
+      model_mappings: {},
+      default_provider: 'openai',
+    });
+
+    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 = await manager.getClient('test-server');
+
+    expect(client).toBeDefined();
+    expect(initSpy).toHaveBeenCalled();
+
+    // Should reuse client
+    const client2 = await manager.getClient('test-server');
+    expect(client2).toBe(client);
+    expect(initSpy).toHaveBeenCalledTimes(1);
+
+    await manager.stopAll();
+    expect(stopSpy).toHaveBeenCalled();
+
+    initSpy.mockRestore();
+    stopSpy.mockRestore();
+  });
+
+  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 = await manager.getClient({
+      name: 'adhoc',
+      command: 'node',
+    });
+
+    expect(client).toBeDefined();
+    expect(initSpy).toHaveBeenCalled();
+
+    initSpy.mockRestore();
+  });
+
+  it('should return undefined if global server not found', async () => {
+    const manager = new MCPManager();
+    const client = await manager.getClient('non-existent');
+    expect(client).toBeUndefined();
+  });
+
+  it('should handle connection failure', async () => {
+    ConfigLoader.setConfig({
+      mcp_servers: {
+        'fail-server': {
+          command: 'fail',
+        },
+      },
+      providers: {},
+      model_mappings: {},
+      default_provider: 'openai',
+    });
+
+    const initSpy = spyOn(MCPClient.prototype, 'initialize').mockRejectedValue(
+      new Error('Connection failed')
+    );
+    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();
+  });
+});