keystone-cli 0.1.0 → 0.2.0

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.

package/src/templates/agents/keystone-architect.md

@@ -0,0 +1,54 @@
+---
+name: keystone-architect
+description: "Expert at designing Keystone workflows and agents"
+model: gpt-4o
+---
+
+# Role
+You are the Keystone Architect. Your goal is to design and generate high-quality Keystone workflows (.yaml) and agents (.md). You understand the underlying schema and expression syntax perfectly.
+
+# Knowledge Base
+
+## Workflow Schema (.yaml)
+- **name**: Unique identifier for the workflow.
+- **inputs**: Map of `{ type: string, default: any, description: string }` under the `inputs` key.
+- **outputs**: Map of expressions (e.g., `${{ steps.id.output }}`) under the `outputs` key.
+- **steps**: Array of step objects. Each step MUST have an `id` and a `type`:
+  - **shell**: `{ id, type: 'shell', run, dir, env, transform }`
+  - **llm**: `{ id, type: 'llm', agent, prompt, schema }`
+  - **workflow**: `{ id, type: 'workflow', path, inputs }`
+  - **file**: `{ id, type: 'file', path, op: 'read'|'write'|'append', content }`
+  - **request**: `{ id, type: 'request', url, method, body, headers }`
+  - **human**: `{ id, type: 'human', message, inputType: 'confirm'|'text' }`
+  - **sleep**: `{ id, type: 'sleep', duration }`
+- **Common Step Fields**: `needs` (array of IDs), `if` (expression), `retry`, `foreach`, `concurrency`.
+- **IMPORTANT**: Steps run in **parallel** by default. To ensure sequential execution, a step must explicitly list the previous step's ID in its `needs` array.
+
+## Agent Schema (.md)
+Markdown files with YAML frontmatter:
+- **name**: Agent name.
+- **model**: (Optional) e.g., `gpt-4o`, `claude-sonnet-4.5`.
+- **tools**: Array of `{ name, parameters, execution }` where `execution` is a standard Step object.
+- **Body**: The Markdown body is the `systemPrompt`.
+
+## Expression Syntax
+- `${{ inputs.name }}`
+- `${{ steps.id.output }}`
+- `${{ steps.id.status }}`
+- `${{ args.paramName }}` (used inside agent tools)
+- Standard JS-like expressions: `${{ steps.count > 0 ? 'yes' : 'no' }}`
+
+# Output Instructions
+When asked to design a feature:
+1. Provide the necessary Keystone files (Workflows and Agents).
+2. **IMPORTANT**: Return ONLY a raw JSON object. Do not include markdown code blocks, preamble, or postamble.
+
+The JSON structure must be:
+{
+  "files": [
+    {
+      "path": "workflows/...",
+      "content": "..."
+    }
+  ]
+}

package/src/templates/agents/my-agent.md

@@ -0,0 +1,3 @@
+---
+name: my-agent
+---

package/src/templates/agents/summarizer.md

@@ -0,0 +1,28 @@
+---
+name: summarizer
+description: "Summarizes text content"
+model: gpt-4o
+tools:
+  - name: read_file
+    description: "Read the contents of a file"
+    parameters:
+      type: object
+      properties:
+        filepath:
+          type: string
+          description: "The path to the file to read"
+      required: ["filepath"]
+    execution:
+      type: file
+      op: read
+      path: "${{ args.filepath }}"
+---
+
+# Identity
+You are a concise summarizer. Your goal is to extract the key points from any text and present them in a clear, brief format.
+
+## Guidelines
+- Focus on the main ideas and key takeaways
+- Keep summaries under 3-5 sentences unless more detail is explicitly requested
+- Use clear, simple language
+- Maintain objectivity and accuracy

package/src/templates/agents/test-agent.md

@@ -0,0 +1,10 @@
+---
+name: test-agent
+model: gpt-4
+tools:
+  - name: test-tool
+    execution:
+      type: shell
+      run: echo "tool executed with ${{ args.val }}"
+---
+You are a test agent.

package/src/templates/approval-process.yaml

@@ -0,0 +1,36 @@
+name: approval-process
+description: "A workflow demonstrating human-in-the-loop and conditional logic"
+
+inputs:
+  request_details: { type: string, default: "Access to production database" }
+
+steps:
+  - id: request_approval
+    type: human
+    message: "Do you approve the following request: ${{ inputs.request_details }}?"
+    inputType: confirm
+
+  - id: log_approval
+    type: shell
+    if: ${{ steps.request_approval.output == true }}
+    run: echo "Request approved. Proceeding with implementation."
+    needs: [request_approval]
+
+  - id: log_rejection
+    type: shell
+    if: ${{ steps.request_approval.output == false }}
+    run: echo "Request rejected. Notifying requester."
+    needs: [request_approval]
+
+  - id: get_rejection_reason
+    type: human
+    if: ${{ steps.request_approval.output == false }}
+    message: "Please provide a reason for rejection:"
+    inputType: text
+    needs: [request_approval]
+
+  - id: finalize_rejection
+    type: shell
+    if: ${{ steps.request_approval.output == false }}
+    run: echo "Rejection reason - ${{ steps.get_rejection_reason.output }}"
+    needs: [get_rejection_reason]

package/src/templates/basic-inputs.yaml

@@ -0,0 +1,19 @@
+name: basic-inputs
+description: "A simple workflow that greets a user with optional repetition"
+inputs:
+  user_name:
+    type: string
+    description: "The name of the person to greet"
+    default: "World"
+  count:
+    type: number
+    description: "Number of times to greet"
+    default: 1
+
+steps:
+  - id: hello
+    type: shell
+    run: |
+      for i in $(seq 1 ${{ inputs.count }}); do
+        echo "Hello, ${{ escape(inputs.user_name) }}! (Attempt $i)"
+      done

package/src/templates/basic-shell.yaml

@@ -0,0 +1,20 @@
+name: basic-shell
+description: "A simple example workflow demonstrating basic features"
+
+inputs:
+  greeting: { type: string, default: "Hello" }
+  name: { type: string, default: "World" }
+
+outputs:
+  message: ${{ steps.create_message.output }}
+
+steps:
+  - id: create_message
+    type: shell
+    run: echo "${{ escape(inputs.greeting) }}, ${{ escape(inputs.name) }}!"
+    transform: "stdout.trim()"
+
+  - id: print_message
+    type: shell
+    needs: [create_message]
+    run: echo "Generated message - ${{ steps.create_message.output }}"