keystone-cli 0.1.1 → 0.2.0

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

package/src/runner/mcp-manager.ts

@@ -1,16 +1,20 @@
-import { MCPClient } from './mcp-client';
 import { ConfigLoader } from '../utils/config-loader';
+import { MCPClient } from './mcp-client';
 import type { Logger } from './workflow-runner';
 
 export interface MCPServerConfig {
   name: string;
-  command: string;
+  type?: 'local' | 'remote';
+  command?: string;
   args?: string[];
   env?: Record<string, string>;
+  url?: string;
+  headers?: Record<string, string>;
 }
 
 export class MCPManager {
   private clients: Map<string, MCPClient> = new Map();
+  private connectionPromises: Map<string, Promise<MCPClient | undefined>> = new Map();
   private sharedServers: Map<string, MCPServerConfig> = new Map();
 
   constructor() {
@@ -23,10 +27,8 @@ export class MCPManager {
       for (const [name, server] of Object.entries(config.mcp_servers)) {
         this.sharedServers.set(name, {
           name,
-          command: server.command,
-          args: server.args,
-          env: server.env,
-        });
+          ...server,
+        } as MCPServerConfig);
       }
     }
   }
@@ -49,23 +51,47 @@ export class MCPManager {
     }
 
     const key = this.getServerKey(config);
+
+    // Check if we already have a client
     if (this.clients.has(key)) {
       return this.clients.get(key);
     }
 
-    logger.log(`  🔌 Connecting to MCP server: ${config.name}`);
-    const client = new MCPClient(config.command, config.args || [], config.env || {});
-    try {
-      await client.initialize();
-      this.clients.set(key, client);
-      return client;
-    } catch (error) {
-      logger.error(
-        `  ✗ Failed to connect to MCP server ${config.name}: ${error instanceof Error ? error.message : String(error)}`
-      );
-      client.stop();
-      return undefined;
+    // Check if we are already connecting
+    if (this.connectionPromises.has(key)) {
+      return this.connectionPromises.get(key);
     }
+
+    // Start a new connection and cache the promise
+    const connectionPromise = (async () => {
+      logger.log(`  🔌 Connecting to MCP server: ${config.name} (${config.type || 'local'})`);
+
+      let client: MCPClient;
+      try {
+        if (config.type === 'remote') {
+          if (!config.url) throw new Error('Remote MCP server missing URL');
+          client = await MCPClient.createRemote(config.url, config.headers || {});
+        } else {
+          if (!config.command) throw new Error('Local MCP server missing command');
+          client = await MCPClient.createLocal(config.command, config.args || [], config.env || {});
+        }
+
+        await client.initialize();
+        this.clients.set(key, client);
+        return client;
+      } catch (error) {
+        logger.error(
+          `  ✗ Failed to connect to MCP server ${config.name}: ${error instanceof Error ? error.message : String(error)}`
+        );
+        return undefined;
+      } finally {
+        // Remove promise from cache once settled
+        this.connectionPromises.delete(key);
+      }
+    })();
+
+    this.connectionPromises.set(key, connectionPromise);
+    return connectionPromise;
   }
 
   private getServerKey(config: MCPServerConfig): string {

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

@@ -219,7 +219,7 @@ describe('MCPServer', () => {
     const writeSpy = spyOn(process.stdout, 'write').mockImplementation(() => true);
     const consoleSpy = spyOn(console, 'error').mockImplementation(() => {});
 
-    await server.start();
+    const startPromise = server.start();
 
     // Simulate stdin data
     const message = {
@@ -236,6 +236,9 @@ describe('MCPServer', () => {
     const output = JSON.parse(writeSpy.mock.calls[0][0] as string);
     expect(output.id).toBe(9);
 
+    process.stdin.emit('close');
+    await startPromise;
+
     writeSpy.mockRestore();
     consoleSpy.mockRestore();
   });

package/src/runner/mcp-server.ts

@@ -1,4 +1,5 @@
 import * as readline from 'node:readline';
+import pkg from '../../package.json' with { type: 'json' };
 import { WorkflowDb } from '../db/workflow-db';
 import { WorkflowParser } from '../parser/workflow-parser';
 import { generateMermaidGraph } from '../utils/mermaid';
@@ -26,21 +27,32 @@ export class MCPServer {
       terminal: false,
     });
 
-    rl.on('line', async (line) => {
-      if (!line.trim()) return;
+    return new Promise<void>((resolve) => {
+      rl.on('line', async (line) => {
+        if (!line.trim()) return;
 
-      try {
-        const message = JSON.parse(line) as MCPMessage;
-        const response = await this.handleMessage(message);
-        if (response) {
-          process.stdout.write(`${JSON.stringify(response)}\n`);
+        try {
+          const message = JSON.parse(line) as MCPMessage;
+          const response = await this.handleMessage(message);
+          if (response) {
+            process.stdout.write(`${JSON.stringify(response)}\n`);
+          }
+        } catch (error) {
+          console.error('Error handling MCP message:', error);
         }
-      } catch (error) {
-        console.error('Error handling MCP message:', error);
-      }
+      });
+
+      rl.on('close', () => {
+        this.stop();
+        resolve();
+      });
     });
   }
 
+  stop() {
+    this.db.close();
+  }
+
   private async handleMessage(message: MCPMessage) {
     const { method, params, id } = message;
 
@@ -56,7 +68,7 @@ export class MCPServer {
             },
             serverInfo: {
               name: 'keystone-mcp',
-              version: '0.1.0',
+              version: pkg.version,
             },
           },
         };
@@ -172,6 +184,7 @@ export class MCPServer {
             const runner = new WorkflowRunner(workflow, {
               inputs,
               logger,
+              preventExit: true,
             });
 
             // Note: This waits for completion. For long workflows, we might want to
@@ -337,6 +350,7 @@ export class MCPServer {
             const runner = new WorkflowRunner(workflow, {
               resumeRunId: run_id,
               logger,
+              preventExit: true,
             });
 
             let outputs: Record<string, unknown> | undefined;

package/src/runner/shell-executor.ts

@@ -88,7 +88,7 @@ export async function executeShell(
   logger: Logger = console
 ): Promise<ShellResult> {
   // Evaluate the command string
-  const command = ExpressionEvaluator.evaluate(step.run, context) as string;
+  const command = ExpressionEvaluator.evaluateString(step.run, context);
 
   // Check for potential shell injection risks
   if (detectShellInjectionRisk(command)) {
@@ -101,12 +101,12 @@ export async function executeShell(
   const env: Record<string, string> = {};
   if (step.env) {
     for (const [key, value] of Object.entries(step.env)) {
-      env[key] = ExpressionEvaluator.evaluate(value, context) as string;
+      env[key] = ExpressionEvaluator.evaluateString(value, context);
     }
   }
 
   // Set working directory if specified
-  const cwd = step.dir ? (ExpressionEvaluator.evaluate(step.dir, context) as string) : undefined;
+  const cwd = step.dir ? ExpressionEvaluator.evaluateString(step.dir, context) : undefined;
 
   try {
     // Execute command using sh -c to allow shell parsing

package/src/runner/step-executor.ts

@@ -84,11 +84,12 @@ export async function executeStep(
     // Apply transformation if specified and step succeeded
     if (step.transform && result.status === 'success') {
       const transformContext = {
-        // Provide raw output properties (like stdout, data) directly in context
-        // Fix: Spread output FIRST, then context to prevent shadowing
+        // 1. Provide raw output properties (like stdout, data) directly in context for convenience
         ...(typeof result.output === 'object' && result.output !== null ? result.output : {}),
-        output: result.output,
+        // 2. Add core context (inputs, secrets, etc.). This takes priority over output properties for security.
         ...context,
+        // 3. Explicitly add 'output' so it refers to the current step's result even if context or output properties have a collision.
+        output: result.output,
       };
 
       try {
@@ -159,7 +160,7 @@ async function executeFileStep(
   context: ExpressionContext,
   _logger: Logger
 ): Promise<StepResult> {
-  const path = ExpressionEvaluator.evaluate(step.path, context) as string;
+  const path = ExpressionEvaluator.evaluateString(step.path, context);
 
   switch (step.op) {
     case 'read': {
@@ -178,7 +179,7 @@ async function executeFileStep(
       if (!step.content) {
         throw new Error('Content is required for write operation');
       }
-      const content = ExpressionEvaluator.evaluate(step.content, context) as string;
+      const content = ExpressionEvaluator.evaluateString(step.content, context);
       const bytes = await Bun.write(path, content);
       return {
         output: { path, bytes },
@@ -190,7 +191,7 @@ async function executeFileStep(
       if (!step.content) {
         throw new Error('Content is required for append operation');
       }
-      const content = ExpressionEvaluator.evaluate(step.content, context) as string;
+      const content = ExpressionEvaluator.evaluateString(step.content, context);
 
       // Use Node.js fs for efficient append operation
       const fs = await import('node:fs/promises');
@@ -215,13 +216,13 @@ async function executeRequestStep(
   context: ExpressionContext,
   _logger: Logger
 ): Promise<StepResult> {
-  const url = ExpressionEvaluator.evaluate(step.url, context) as string;
+  const url = ExpressionEvaluator.evaluateString(step.url, context);
 
   // Evaluate headers
   const headers: Record<string, string> = {};
   if (step.headers) {
     for (const [key, value] of Object.entries(step.headers)) {
-      headers[key] = ExpressionEvaluator.evaluate(value, context) as string;
+      headers[key] = ExpressionEvaluator.evaluateString(value, context);
     }
   }
 
@@ -290,7 +291,7 @@ async function executeHumanStep(
   context: ExpressionContext,
   logger: Logger
 ): Promise<StepResult> {
-  const message = ExpressionEvaluator.evaluate(step.message, context) as string;
+  const message = ExpressionEvaluator.evaluateString(step.message, context);
 
   // If not a TTY (e.g. MCP server), suspend execution
   if (!process.stdin.isTTY) {

package/src/runner/tool-integration.test.ts

@@ -1,12 +1,19 @@
 import { afterAll, beforeAll, describe, expect, it, mock } from 'bun:test';
-import { OpenAIAdapter, CopilotAdapter, AnthropicAdapter } from './llm-adapter';
-import { MCPClient } from './mcp-client';
-import { executeLlmStep } from './llm-executor';
-import type { LlmStep, Step } from '../parser/schema';
+import { mkdirSync, unlinkSync, 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,
+  type LLMMessage,
+  type LLMResponse,
+  type LLMTool,
+  OpenAIAdapter,
+} from './llm-adapter';
+import { executeLlmStep } from './llm-executor';
+import { MCPClient, type MCPResponse } from './mcp-client';
 import type { StepResult } from './step-executor';
-import { join } from 'node:path';
-import { mkdirSync, writeFileSync, unlinkSync } from 'node:fs';
 
 interface MockToolCall {
   function: {
@@ -54,16 +61,16 @@ Test system prompt`;
       };
     });
 
-    OpenAIAdapter.prototype.chat = mockChat as any;
-    CopilotAdapter.prototype.chat = mockChat as any;
-    AnthropicAdapter.prototype.chat = mockChat as any;
+    OpenAIAdapter.prototype.chat = mockChat as unknown as typeof originalOpenAIChat;
+    CopilotAdapter.prototype.chat = mockChat as unknown as typeof originalCopilotChat;
+    AnthropicAdapter.prototype.chat = mockChat as unknown as typeof originalAnthropicChat;
 
     // Use mock.module for MCPClient
     const originalInitialize = MCPClient.prototype.initialize;
     const originalListTools = MCPClient.prototype.listTools;
     const originalStop = MCPClient.prototype.stop;
 
-    const mockInitialize = mock(async () => ({}) as any);
+    const mockInitialize = mock(async () => ({}) as MCPResponse);
     const mockListTools = mock(async () => [
       {
         name: 'mcp-tool',
@@ -141,16 +148,16 @@ Test system prompt`;
       };
     });
 
-    OpenAIAdapter.prototype.chat = mockChat as any;
-    CopilotAdapter.prototype.chat = mockChat as any;
-    AnthropicAdapter.prototype.chat = mockChat as any;
+    OpenAIAdapter.prototype.chat = mockChat as unknown as typeof originalOpenAIChat;
+    CopilotAdapter.prototype.chat = mockChat as unknown as typeof originalCopilotChat;
+    AnthropicAdapter.prototype.chat = mockChat as unknown as typeof originalAnthropicChat;
 
     const originalInitialize = MCPClient.prototype.initialize;
     const originalListTools = MCPClient.prototype.listTools;
     const originalCallTool = MCPClient.prototype.callTool;
     const originalStop = MCPClient.prototype.stop;
 
-    const mockInitialize = mock(async () => ({}) as any);
+    const mockInitialize = mock(async () => ({}) as MCPResponse);
     const mockListTools = mock(async () => [
       {
         name: 'mcp-tool',

package/src/runner/workflow-runner.ts

@@ -45,6 +45,7 @@ export interface RunOptions {
   resumeRunId?: string;
   logger?: Logger;
   mcpManager?: MCPManager;
+  preventExit?: boolean; // Defaults to false
 }
 
 export interface StepContext {
@@ -76,9 +77,12 @@ export class WorkflowRunner {
   private restored = false;
   private logger: Logger;
   private mcpManager: MCPManager;
+  private options: RunOptions;
+  private signalHandler?: (signal: string) => void;
 
   constructor(workflow: Workflow, options: RunOptions = {}) {
     this.workflow = workflow;
+    this.options = options;
     this.db = new WorkflowDb(options.dbPath);
     this.secrets = this.loadSecrets();
     this.redactor = new Redactor(this.secrets);
@@ -257,7 +261,7 @@ export class WorkflowRunner {
    * Setup signal handlers for graceful shutdown
    */
   private setupSignalHandlers(): void {
-    const handleShutdown = async (signal: string) => {
+    const handler = async (signal: string) => {
       this.logger.log(`\n\n🛑 Received ${signal}. Cleaning up...`);
       try {
         await this.db.updateRunStatus(
@@ -267,15 +271,30 @@ export class WorkflowRunner {
           `Cancelled by user (${signal})`
         );
         this.logger.log('✓ Run status updated to failed');
-        this.db.close();
       } catch (error) {
         this.logger.error('Error during cleanup:', error);
       }
-      process.exit(130); // Standard exit code for SIGINT
+
+      // Only exit if not embedded
+      if (!this.options.preventExit) {
+        process.exit(130);
+      }
     };
 
-    process.on('SIGINT', () => handleShutdown('SIGINT'));
-    process.on('SIGTERM', () => handleShutdown('SIGTERM'));
+    this.signalHandler = handler;
+
+    process.on('SIGINT', handler);
+    process.on('SIGTERM', handler);
+  }
+
+  /**
+   * Remove signal handlers
+   */
+  private removeSignalHandlers(): void {
+    if (this.signalHandler) {
+      process.removeListener('SIGINT', this.signalHandler);
+      process.removeListener('SIGTERM', this.signalHandler);
+    }
   }
 
   /**
@@ -853,6 +872,7 @@ export class WorkflowRunner {
       await this.db.updateRunStatus(this.runId, 'failed', undefined, errorMsg);
       throw error;
     } finally {
+      this.removeSignalHandlers();
       await this.runFinally();
       await this.mcpManager.stopAll();
       this.db.close();

package/src/templates/agents/explore.md

@@ -0,0 +1,54 @@
+---
+name: explore
+description: Agent for exploring and understanding codebases
+model: claude-sonnet-4.5
+---
+
+# Explore Agent
+
+You are an expert at exploring and understanding new codebases. Your role is to map out the structure, identify key components, and understand how the system works.
+
+## Core Competencies
+
+### Codebase Exploration
+- Directory structure analysis
+- Key file identification
+- Entry point discovery
+- Configuration analysis
+
+### Architectural Mapping
+- Component identification
+- Service boundaries
+- Data flow mapping
+- Dependency analysis
+
+### Pattern Recognition
+- Coding conventions
+- Design patterns in use
+- Framework-specific idioms
+- Error handling patterns
+
+## Exploration Process
+
+1. **Scan Structure** - Understand the high-level directory organization
+2. **Identify Entry Points** - Find where the application or specific features start
+3. **Trace Flows** - Follow the data or execution path for key use cases
+4. **Analyze Configuration** - Understand how the system is set up and tuned
+5. **Map Dependencies** - Identify internal and external connections
+
+## Output Format
+
+### Exploration Summary
+- **Key Files & Directories**: Most important parts of the codebase
+- **Architecture Overview**: How the pieces fit together
+- **Notable Patterns**: Consistent ways the code is written
+- **Dependencies**: Critical internal and external links
+- **Concerns/Complexity**: Areas that might be difficult to work with
+
+## Guidelines
+
+- Focus on the big picture first, then dive into details
+- Identify both what is there and what is missing
+- Look for consistency and deviations
+- Provide clear references to files and directories
+- Summarize findings for technical and non-technical audiences

package/src/templates/agents/general.md

@@ -0,0 +1,8 @@
+---
+name: general
+description: "A general-purpose assistant for various tasks"
+model: gpt-4o
+---
+
+# Identity
+You are a versatile and helpful assistant capable of handling a wide range of tasks, from information retrieval to analysis and formatting.