keystone-cli 0.1.1 → 0.3.0

package/src/runner/mcp-manager.ts

@@ -1,16 +1,23 @@
-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>;
+  oauth?: {
+    scope?: 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 +30,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 +54,84 @@ 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');
+
+          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');
+
+          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();
+        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.test.ts

@@ -374,7 +374,7 @@ describe('step-executor', () => {
       const step: WorkflowStep = {
         id: 'w1',
         type: 'workflow',
-        workflow: 'child.yaml',
+        path: 'child.yaml',
       };
       // @ts-ignore
       const executeWorkflowFn = mock(() =>
@@ -392,7 +392,7 @@ describe('step-executor', () => {
       const step: WorkflowStep = {
         id: 'w1',
         type: 'workflow',
-        workflow: 'child.yaml',
+        path: 'child.yaml',
       };
       const result = await executeStep(step, context);
       expect(result.status).toBe('failed');

package/src/runner/step-executor.ts

@@ -42,7 +42,8 @@ export async function executeStep(
   context: ExpressionContext,
   logger: Logger = console,
   executeWorkflowFn?: (step: WorkflowStep, context: ExpressionContext) => Promise<StepResult>,
-  mcpManager?: MCPManager
+  mcpManager?: MCPManager,
+  workflowDir?: string
 ): Promise<StepResult> {
   try {
     let result: StepResult;
@@ -66,9 +67,10 @@ export async function executeStep(
         result = await executeLlmStep(
           step,
           context,
-          (s, c) => executeStep(s, c, logger, executeWorkflowFn, mcpManager),
+          (s, c) => executeStep(s, c, logger, executeWorkflowFn, mcpManager, workflowDir),
           logger,
-          mcpManager
+          mcpManager,
+          workflowDir
         );
         break;
       case 'workflow':
@@ -84,11 +86,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 +162,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 +181,14 @@ 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);
+
+      // Ensure parent directory exists
+      const fs = await import('node:fs/promises');
+      const pathModule = await import('node:path');
+      const dir = pathModule.dirname(path);
+      await fs.mkdir(dir, { recursive: true });
+
       const bytes = await Bun.write(path, content);
       return {
         output: { path, bytes },
@@ -190,10 +200,15 @@ 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
+      // Ensure parent directory exists
       const fs = await import('node:fs/promises');
+      const pathModule = await import('node:path');
+      const dir = pathModule.dirname(path);
+      await fs.mkdir(dir, { recursive: true });
+
+      // Use Node.js fs for efficient append operation
       await fs.appendFile(path, content, 'utf-8');
 
       return {
@@ -215,13 +230,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 +305,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) {
@@ -309,10 +324,10 @@ async function executeHumanStep(
   try {
     if (step.inputType === 'confirm') {
       logger.log(`\n❓ ${message}`);
-      logger.log('Press Enter to continue, or Ctrl+C to cancel...');
-      await rl.question('');
+      const answer = await rl.question('Confirm? (Y/n): ');
+      const isConfirmed = answer.toLowerCase() !== 'n';
       return {
-        output: true,
+        output: isConfirmed,
         status: 'success',
       };
     }

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

@@ -1,4 +1,5 @@
 import { randomUUID } from 'node:crypto';
+import { dirname } from 'node:path';
 import { WorkflowDb } from '../db/workflow-db.ts';
 import type { ExpressionContext } from '../expression/evaluator.ts';
 import { ExpressionEvaluator } from '../expression/evaluator.ts';
@@ -45,6 +46,8 @@ export interface RunOptions {
   resumeRunId?: string;
   logger?: Logger;
   mcpManager?: MCPManager;
+  preventExit?: boolean; // Defaults to false
+  workflowDir?: string;
 }
 
 export interface StepContext {
@@ -76,9 +79,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 +263,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 +273,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);
+    }
   }
 
   /**
@@ -437,7 +458,8 @@ export class WorkflowRunner {
         context,
         this.logger,
         this.executeSubWorkflow.bind(this),
-        this.mcpManager
+        this.mcpManager,
+        this.options.workflowDir
       );
       if (result.status === 'failed') {
         throw new Error(result.error || 'Step failed');
@@ -681,6 +703,7 @@ export class WorkflowRunner {
   ): Promise<StepResult> {
     const workflowPath = WorkflowRegistry.resolvePath(step.path);
     const workflow = WorkflowParser.loadWorkflow(workflowPath);
+    const subWorkflowDir = dirname(workflowPath);
 
     // Evaluate inputs for the sub-workflow
     const inputs: Record<string, unknown> = {};
@@ -697,6 +720,7 @@ export class WorkflowRunner {
       dbPath: this.db.dbPath,
       logger: this.logger,
       mcpManager: this.mcpManager,
+      workflowDir: subWorkflowDir,
     });
 
     try {
@@ -853,8 +877,11 @@ export class WorkflowRunner {
       await this.db.updateRunStatus(this.runId, 'failed', undefined, errorMsg);
       throw error;
     } finally {
+      this.removeSignalHandlers();
       await this.runFinally();
-      await this.mcpManager.stopAll();
+      if (!this.options.mcpManager) {
+        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.