keystone-cli 0.8.0 → 1.0.1

package/src/runner/llm-executor.ts

@@ -3,6 +3,7 @@ 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 { LIMITS } from '../utils/constants';
 import { extractJson } from '../utils/json-parser';
 import { ConsoleLogger, type Logger } from '../utils/logger.ts';
 import { RedactionBuffer, Redactor } from '../utils/redactor';
@@ -12,11 +13,117 @@ import type { MCPManager, MCPServerConfig } from './mcp-manager';
 import { STANDARD_TOOLS, validateStandardToolSecurity } from './standard-tools';
 import type { StepResult } from './step-executor';
 
+/**
+ * Truncate message history to prevent unbounded memory growth.
+ * Preserves system messages and keeps the most recent messages.
+ */
+function estimateMessageBytes(message: LLMMessage): number {
+  let size = 0;
+  if (typeof message.content === 'string') {
+    size += Buffer.byteLength(message.content, 'utf8');
+  }
+  if (message.tool_calls) {
+    size += Buffer.byteLength(JSON.stringify(message.tool_calls), 'utf8');
+  }
+  if (message.reasoning) {
+    size += Buffer.byteLength(JSON.stringify(message.reasoning), 'utf8');
+  }
+  if (message.name) {
+    size += Buffer.byteLength(message.name, 'utf8');
+  }
+  return size;
+}
+
+function truncateStringByBytes(value: string, maxBytes: number): string {
+  if (maxBytes <= 0) return '';
+  if (Buffer.byteLength(value, 'utf8') <= maxBytes) return value;
+
+  let low = 0;
+  let high = value.length;
+  while (low < high) {
+    const mid = Math.ceil((low + high) / 2);
+    const slice = value.slice(0, mid);
+    if (Buffer.byteLength(slice, 'utf8') <= maxBytes) {
+      low = mid;
+    } else {
+      high = mid - 1;
+    }
+  }
+  return value.slice(0, low);
+}
+
+function truncateToolOutput(content: string, maxBytes: number): string {
+  const contentBytes = Buffer.byteLength(content, 'utf8');
+  if (contentBytes <= maxBytes) return content;
+
+  const suffix = '... [truncated output]';
+  const suffixBytes = Buffer.byteLength(suffix, 'utf8');
+  const truncated = truncateStringByBytes(content, Math.max(0, maxBytes - suffixBytes));
+  return `${truncated}${suffix}`;
+}
+
+function safeJsonStringify(value: unknown): string {
+  try {
+    return JSON.stringify(value);
+  } catch {
+    const seen = new WeakSet<object>();
+    try {
+      return JSON.stringify(value, (_key, val) => {
+        if (typeof val === 'bigint') return val.toString();
+        if (typeof val === 'object' && val !== null) {
+          if (seen.has(val)) return '[Circular]';
+          seen.add(val);
+        }
+        return val;
+      });
+    } catch {
+      return String(value);
+    }
+  }
+}
+
+function truncateMessages(
+  messages: LLMMessage[],
+  maxHistory: number,
+  maxBytes: number
+): LLMMessage[] {
+  if (messages.length === 0) return messages;
+
+  // Keep all system messages
+  const systemMessages = messages.filter((m) => m.role === 'system');
+  const nonSystem = messages.filter((m) => m.role !== 'system');
+
+  // Keep most recent non-system messages, accounting for system messages
+  const nonSystemLimit = Math.max(0, maxHistory - systemMessages.length);
+  let keep = nonSystem.slice(-nonSystemLimit);
+
+  // Enforce total byte budget with a most-recent tail
+  if (maxBytes > 0) {
+    const systemBytes = systemMessages.reduce((total, msg) => total + estimateMessageBytes(msg), 0);
+    let remaining = maxBytes - systemBytes;
+    if (remaining <= 0) {
+      return systemMessages;
+    }
+
+    const tail: LLMMessage[] = [];
+    for (let i = keep.length - 1; i >= 0; i--) {
+      const msg = keep[i];
+      const msgBytes = estimateMessageBytes(msg);
+      if (msgBytes > remaining) break;
+      tail.push(msg);
+      remaining -= msgBytes;
+    }
+    keep = tail.reverse();
+  }
+
+  return [...systemMessages, ...keep];
+}
+
 interface ToolDefinition {
   name: string;
   description?: string;
   parameters: unknown;
-  source: 'agent' | 'step' | 'mcp' | 'standard';
+  source: 'agent' | 'step' | 'mcp' | 'standard' | 'handoff';
   execution?: Step;
   mcpClient?: MCPClient;
 }
@@ -27,7 +134,9 @@ export async function executeLlmStep(
   executeStepFn: (step: Step, context: ExpressionContext) => Promise<StepResult>,
   logger: Logger = new ConsoleLogger(),
   mcpManager?: MCPManager,
-  workflowDir?: string
+  workflowDir?: string,
+  abortSignal?: AbortSignal,
+  getAdapterFn?: typeof getAdapter
 ): Promise<StepResult> {
   const agentPath = resolveAgentPath(step.agent, workflowDir);
   const agent = parseAgent(agentPath);
@@ -37,19 +146,30 @@ export async function executeLlmStep(
   const prompt = ExpressionEvaluator.evaluateString(step.prompt, context);
 
   const fullModelString = provider ? `${provider}:${model}` : model;
-  const { adapter, resolvedModel } = getAdapter(fullModelString);
+  const { adapter, resolvedModel } = (getAdapterFn || 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)}`;
+  if (step.outputSchema) {
+    systemPrompt += `\n\nIMPORTANT: You must output valid JSON that matches the following schema:\n${JSON.stringify(step.outputSchema, null, 2)}`;
   }
 
-  const messages: LLMMessage[] = [];
+  let messages: LLMMessage[] = [];
+  const maxToolOutputBytes = LIMITS.MAX_TOOL_OUTPUT_BYTES;
 
   // Resume from state if provided
-  if (context.output && typeof context.output === 'object' && 'messages' in context.output) {
-    messages.push(...(context.output.messages as LLMMessage[]));
+  const stepState =
+    context.steps && typeof context.steps === 'object'
+      ? (context.steps as Record<string, { output?: unknown }>)[step.id]
+      : undefined;
+  const stepOutput = stepState?.output;
+  const resumeOutput =
+    stepOutput && typeof stepOutput === 'object' && 'messages' in stepOutput
+      ? stepOutput
+      : context.output;
+
+  if (resumeOutput && typeof resumeOutput === 'object' && 'messages' in resumeOutput) {
+    messages.push(...(resumeOutput.messages as LLMMessage[]));
 
     // If we have an answer in inputs, add it as a tool result for the last tool call
     const stepInputs = context.inputs?.[step.id] as Record<string, unknown> | undefined;
@@ -62,7 +182,7 @@ export async function executeLlmStep(
           role: 'tool',
           tool_call_id: askCall.id,
           name: 'ask',
-          content: String(answer),
+          content: truncateToolOutput(String(answer), maxToolOutputBytes),
         });
       }
     }
@@ -72,11 +192,22 @@ export async function executeLlmStep(
 
   const localMcpClients: MCPClient[] = [];
   const allTools: ToolDefinition[] = [];
+  const toolRegistry = new Map<string, string>();
+  const registerTool = (tool: ToolDefinition) => {
+    const existing = toolRegistry.get(tool.name);
+    if (existing) {
+      throw new Error(
+        `Duplicate tool name "${tool.name}" from ${tool.source}; already defined by ${existing}. Rename one of them.`
+      );
+    }
+    toolRegistry.set(tool.name, tool.source);
+    allTools.push(tool);
+  };
 
   try {
     // 1. Add agent tools
     for (const tool of agent.tools) {
-      allTools.push({
+      registerTool({
         name: tool.name,
         description: tool.description,
         parameters: tool.parameters || {
@@ -92,7 +223,7 @@ export async function executeLlmStep(
     // 2. Add step tools
     if (step.tools) {
       for (const tool of step.tools) {
-        allTools.push({
+        registerTool({
           name: tool.name,
           description: tool.description,
           parameters: tool.parameters || {
@@ -109,7 +240,7 @@ export async function executeLlmStep(
     // 3. Add Standard tools
     if (step.useStandardTools) {
       for (const tool of STANDARD_TOOLS) {
-        allTools.push({
+        registerTool({
           name: tool.name,
           description: tool.description,
           parameters: tool.parameters || {
@@ -123,7 +254,39 @@ export async function executeLlmStep(
       }
     }
 
-    // 4. Add MCP tools
+    // 4. Add Engine handoff tool
+    if (step.handoff) {
+      const toolName = step.handoff.name || 'handoff';
+      const description =
+        step.handoff.description || `Delegate to engine ${step.handoff.engine.command}`;
+      const parameters = step.handoff.inputSchema || {
+        type: 'object',
+        properties: {},
+        additionalProperties: true,
+      };
+
+      const handoffStep: Step = {
+        id: `${step.id}-handoff`,
+        type: 'engine',
+        command: step.handoff.engine.command,
+        args: step.handoff.engine.args,
+        env: step.handoff.engine.env,
+        cwd: step.handoff.engine.cwd,
+        timeout: step.handoff.engine.timeout,
+        outputSchema: step.handoff.engine.outputSchema,
+        input: step.handoff.engine.input ?? '${{ args }}',
+      };
+
+      registerTool({
+        name: toolName,
+        description,
+        parameters,
+        source: 'handoff',
+        execution: handoffStep,
+      });
+    }
+
+    // 5. Add MCP tools
     const mcpServersToConnect: (string | MCPServerConfig)[] = [...(step.mcpServers || [])];
     if (step.useGlobalMcp && mcpManager) {
       const globalServers = mcpManager.getGlobalServers();
@@ -169,7 +332,7 @@ export async function executeLlmStep(
             if (client) {
               const mcpTools = await client.listTools();
               for (const tool of mcpTools) {
-                allTools.push({
+                registerTool({
                   name: tool.name,
                   description: tool.description,
                   parameters: tool.inputSchema,
@@ -200,6 +363,11 @@ export async function executeLlmStep(
     }));
 
     if (step.allowClarification) {
+      if (toolRegistry.has('ask')) {
+        throw new Error(
+          'Tool name "ask" is reserved for clarification. Rename your tool or disable allowClarification.'
+        );
+      }
       llmTools.push({
         type: 'function' as const,
         function: {
@@ -230,23 +398,37 @@ export async function executeLlmStep(
     };
 
     // Create redactor once outside the loop for performance (regex compilation)
-    const redactor = new Redactor(context.secrets || {});
+    const redactor = new Redactor(context.secrets || {}, {
+      forcedSecrets: context.secretValues || [],
+    });
     const redactionBuffer = new RedactionBuffer(redactor);
+    const maxHistory = step.maxMessageHistory || LIMITS.MAX_MESSAGE_HISTORY;
+    const maxConversationBytes = LIMITS.MAX_CONVERSATION_BYTES;
+    const formatToolContent = (content: string): string =>
+      truncateToolOutput(content, maxToolOutputBytes);
 
     while (iterations < maxIterations) {
       iterations++;
+      if (abortSignal?.aborted) {
+        throw new Error('Step canceled');
+      }
+
+      // Truncate message history to prevent unbounded growth
+      messages = truncateMessages(messages, maxHistory, maxConversationBytes);
+      const truncatedMessages = messages;
 
-      const response = await adapter.chat(messages, {
+      const response = await adapter.chat(truncatedMessages, {
         model: resolvedModel,
         tools: llmTools.length > 0 ? llmTools : undefined,
         onStream: (chunk) => {
-          if (!step.schema) {
+          if (!step.outputSchema) {
             process.stdout.write(redactionBuffer.process(chunk));
           }
         },
+        signal: abortSignal,
       });
 
-      if (!step.schema) {
+      if (!step.outputSchema) {
         process.stdout.write(redactionBuffer.flush());
       }
 
@@ -263,7 +445,7 @@ export async function executeLlmStep(
         let output = message.content;
 
         // If schema is defined, attempt to parse JSON
-        if (step.schema && typeof output === 'string') {
+        if (step.outputSchema && typeof output === 'string') {
           try {
             output = extractJson(output) as typeof output;
           } catch (e) {
@@ -287,6 +469,9 @@ export async function executeLlmStep(
 
       // Execute tools
       for (const toolCall of message.tool_calls) {
+        if (abortSignal?.aborted) {
+          throw new Error('Step canceled');
+        }
         const argsStr = toolCall.function.arguments;
         let displayArgs = '';
         try {
@@ -315,7 +500,9 @@ export async function executeLlmStep(
                 role: 'tool',
                 tool_call_id: toolCall.id,
                 name: 'ask',
-                content: `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`,
+                content: formatToolContent(
+                  `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`
+                ),
               });
               continue;
             }
@@ -337,11 +524,12 @@ export async function executeLlmStep(
                 role: 'tool',
                 tool_call_id: toolCall.id,
                 name: 'ask',
-                content: String(result.output),
+                content: formatToolContent(String(result.output)),
               });
               continue;
             }
             // In non-TTY, we suspend
+            messages = truncateMessages(messages, maxHistory, maxConversationBytes);
             return {
               status: 'suspended',
               output: {
@@ -356,7 +544,7 @@ export async function executeLlmStep(
             role: 'tool',
             tool_call_id: toolCall.id,
             name: toolCall.function.name,
-            content: `Error: Tool ${toolCall.function.name} not found`,
+            content: formatToolContent(`Error: Tool ${toolCall.function.name} not found`),
           });
           continue;
         }
@@ -369,7 +557,9 @@ export async function executeLlmStep(
             role: 'tool',
             tool_call_id: toolCall.id,
             name: toolCall.function.name,
-            content: `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`,
+            content: formatToolContent(
+              `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`
+            ),
           });
           continue;
         }
@@ -381,14 +571,16 @@ export async function executeLlmStep(
               role: 'tool',
               tool_call_id: toolCall.id,
               name: toolCall.function.name,
-              content: JSON.stringify(result),
+              content: formatToolContent(safeJsonStringify(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)}`,
+              content: formatToolContent(
+                `Error: ${error instanceof Error ? error.message : String(error)}`
+              ),
             });
           }
         } else if (toolInfo.execution) {
@@ -404,7 +596,9 @@ export async function executeLlmStep(
                 role: 'tool',
                 tool_call_id: toolCall.id,
                 name: toolCall.function.name,
-                content: `Security Error: ${error instanceof Error ? error.message : String(error)}`,
+                content: formatToolContent(
+                  `Security Error: ${error instanceof Error ? error.message : String(error)}`
+                ),
               });
               continue;
             }
@@ -422,13 +616,16 @@ export async function executeLlmStep(
             role: 'tool',
             tool_call_id: toolCall.id,
             name: toolCall.function.name,
-            content:
+            content: formatToolContent(
               result.status === 'success'
-                ? JSON.stringify(result.output)
-                : `Error: ${result.error}`,
+                ? safeJsonStringify(result.output)
+                : `Error: ${result.error}`
+            ),
           });
         }
       }
+
+      messages = truncateMessages(messages, maxHistory, maxConversationBytes);
     }
 
     throw new Error('Max ReAct iterations reached');

package/src/runner/mcp-client.ts

@@ -80,7 +80,7 @@ export async function validateRemoteUrl(
   try {
     parsed = new URL(url);
   } catch {
-    throw new Error(`Invalid MCP server URL: ${url}`);
+    throw new Error(`Invalid URL: ${url}`);
   }
 
   // Skip all security checks if allowInsecure is set (for development/testing)
@@ -91,7 +91,7 @@ export async function validateRemoteUrl(
   // Require HTTPS in production
   if (parsed.protocol !== 'https:') {
     throw new Error(
-      `SSRF Protection: MCP remote URL must use HTTPS. Got: ${parsed.protocol}. Set allowInsecure option to true if you trust this server.`
+      `SSRF Protection: URL must use HTTPS. Got: ${parsed.protocol}. Set allowInsecure option to true if you trust this server.`
     );
   }
 
@@ -587,6 +587,8 @@ export class MCPClient {
   >();
   private timeout: number;
   private logger: Logger;
+  private _isHealthy = true;
+  private lastHealthCheck = Date.now();
 
   constructor(
     transportOrCommand: MCPTransport | string,
@@ -617,6 +619,21 @@ export class MCPClient {
     });
   }
 
+  /**
+   * Returns whether the client is considered healthy.
+   * Updated after each health check or failed request.
+   */
+  get isHealthy(): boolean {
+    return this._isHealthy;
+  }
+
+  /**
+   * Returns the timestamp of the last health check.
+   */
+  get lastHealthCheckTime(): number {
+    return this.lastHealthCheck;
+  }
+
   static async createLocal(
     command: string,
     args: string[] = [],
@@ -660,6 +677,7 @@ export class MCPClient {
       const timeoutId = setTimeout(() => {
         if (this.pendingRequests.has(id)) {
           this.pendingRequests.delete(id);
+          this._isHealthy = false;
           reject(new Error(`MCP request timeout: ${method}`));
         }
       }, this.timeout);
@@ -669,13 +687,14 @@ export class MCPClient {
       this.transport.send(message).catch((err) => {
         clearTimeout(timeoutId);
         this.pendingRequests.delete(id);
+        this._isHealthy = false;
         reject(err);
       });
     });
   }
 
   async initialize() {
-    return this.request('initialize', {
+    const response = await this.request('initialize', {
       protocolVersion: MCP_PROTOCOL_VERSION,
       capabilities: {},
       clientInfo: {
@@ -683,6 +702,8 @@ export class MCPClient {
         version: pkg.version,
       },
     });
+    this._isHealthy = true;
+    return response;
   }
 
   async listTools(): Promise<MCPTool[]> {
@@ -701,6 +722,52 @@ export class MCPClient {
     return response.result;
   }
 
+  /**
+   * Perform a health check by sending a ping-like request.
+   * Uses a shorter timeout than normal requests.
+   *
+   * @param timeout Optional timeout in ms (default: 5000)
+   * @returns true if healthy, false otherwise
+   */
+  async healthCheck(timeout = 5000): Promise<boolean> {
+    this.lastHealthCheck = Date.now();
+    const id = this.messageId++;
+    const message = {
+      jsonrpc: '2.0',
+      id,
+      method: 'ping', // Standard MCP ping or falls back gracefully
+    };
+
+    return new Promise((resolve) => {
+      const timeoutId = setTimeout(() => {
+        if (this.pendingRequests.has(id)) {
+          this.pendingRequests.delete(id);
+          this._isHealthy = false;
+          resolve(false);
+        }
+      }, timeout);
+
+      this.pendingRequests.set(id, {
+        resolve: () => {
+          this._isHealthy = true;
+          resolve(true);
+        },
+        reject: () => {
+          this._isHealthy = false;
+          resolve(false);
+        },
+        timeoutId,
+      });
+
+      this.transport.send(message).catch(() => {
+        clearTimeout(timeoutId);
+        this.pendingRequests.delete(id);
+        this._isHealthy = false;
+        resolve(false);
+      });
+    });
+  }
+
   stop() {
     // Reject all pending requests to prevent hanging callers
     for (const [, pending] of this.pendingRequests) {