keystone-cli 1.0.2 → 1.1.0

package/src/runner/llm-executor.ts

@@ -1,638 +0,0 @@
-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 { LIMITS } from '../utils/constants';
-import { extractJson } from '../utils/json-parser';
-import { ConsoleLogger, type Logger } from '../utils/logger.ts';
-import { RedactionBuffer, Redactor } from '../utils/redactor';
-import { type LLMMessage, getAdapter } from './llm-adapter';
-import { MCPClient } from './mcp-client';
-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' | 'handoff';
-  execution?: Step;
-  mcpClient?: MCPClient;
-}
-
-export async function executeLlmStep(
-  step: LlmStep,
-  context: ExpressionContext,
-  executeStepFn: (step: Step, context: ExpressionContext) => Promise<StepResult>,
-  logger: Logger = new ConsoleLogger(),
-  mcpManager?: MCPManager,
-  workflowDir?: string,
-  abortSignal?: AbortSignal,
-  getAdapterFn?: typeof getAdapter
-): Promise<StepResult> {
-  const agentPath = resolveAgentPath(step.agent, workflowDir);
-  const agent = parseAgent(agentPath);
-
-  const provider = step.provider || agent.provider;
-  const model = step.model || agent.model || 'gpt-4o';
-  const prompt = ExpressionEvaluator.evaluateString(step.prompt, context);
-
-  const fullModelString = provider ? `${provider}:${model}` : model;
-  const { adapter, resolvedModel } = (getAdapterFn || getAdapter)(fullModelString);
-
-  // Inject schema instructions if present
-  let systemPrompt = agent.systemPrompt;
-  if (step.outputSchema) {
-    systemPrompt += `\n\nIMPORTANT: You must output valid JSON that matches the following schema:\n${JSON.stringify(step.outputSchema, null, 2)}`;
-  }
-
-  let messages: LLMMessage[] = [];
-  const maxToolOutputBytes = LIMITS.MAX_TOOL_OUTPUT_BYTES;
-
-  // Resume from state if provided
-  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;
-    if (stepInputs && typeof stepInputs === 'object' && '__answer' in stepInputs) {
-      const answer = stepInputs.__answer;
-      const lastMessage = messages[messages.length - 1];
-      const askCall = lastMessage?.tool_calls?.find((tc) => tc.function.name === 'ask');
-      if (askCall) {
-        messages.push({
-          role: 'tool',
-          tool_call_id: askCall.id,
-          name: 'ask',
-          content: truncateToolOutput(String(answer), maxToolOutputBytes),
-        });
-      }
-    }
-  } else {
-    messages.push({ role: 'system', content: systemPrompt }, { role: 'user', content: prompt });
-  }
-
-  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) {
-      registerTool({
-        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) {
-        registerTool({
-          name: tool.name,
-          description: tool.description,
-          parameters: tool.parameters || {
-            type: 'object',
-            properties: {},
-            additionalProperties: true,
-          },
-          source: 'step',
-          execution: tool.execution,
-        });
-      }
-    }
-
-    // 3. Add Standard tools
-    if (step.useStandardTools) {
-      for (const tool of STANDARD_TOOLS) {
-        registerTool({
-          name: tool.name,
-          description: tool.description,
-          parameters: tool.parameters || {
-            type: 'object',
-            properties: {},
-            additionalProperties: true,
-          },
-          source: 'standard',
-          execution: tool.execution,
-        });
-      }
-    }
-
-    // 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();
-      for (const globalServer of globalServers) {
-        // Only add if not already explicitly listed
-        const alreadyListed = mcpServersToConnect.some((s) => {
-          const name = typeof s === 'string' ? s : s.name;
-          return name === globalServer.name;
-        });
-        if (!alreadyListed) {
-          mcpServersToConnect.push(globalServer);
-        }
-      }
-    }
-
-    if (mcpServersToConnect.length > 0) {
-      await Promise.all(
-        mcpServersToConnect.map(async (server) => {
-          let client: MCPClient | undefined;
-          const serverName = typeof server === 'string' ? server : server.name;
-
-          try {
-            if (mcpManager) {
-              client = await mcpManager.getClient(server as string | MCPServerConfig, 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`
-                );
-                return;
-              }
-              logger.log(`  🔌 Connecting to MCP server: ${server.name}`);
-              client = await MCPClient.createLocal(
-                (server as MCPServerConfig).command || 'node',
-                (server as MCPServerConfig).args || [],
-                (server as MCPServerConfig).env || {}
-              );
-              await client.initialize();
-              localMcpClients.push(client);
-            }
-
-            if (client) {
-              const mcpTools = await client.listTools();
-              for (const tool of mcpTools) {
-                registerTool({
-                  name: tool.name,
-                  description: tool.description,
-                  parameters: tool.inputSchema,
-                  source: 'mcp',
-                  mcpClient: client,
-                });
-              }
-            }
-          } catch (error) {
-            logger.error(
-              `  ✗ Failed to list tools from MCP server ${serverName}: ${error instanceof Error ? error.message : String(error)}`
-            );
-            if (!mcpManager && client) {
-              client.stop();
-            }
-          }
-        })
-      );
-    }
-
-    const llmTools = allTools.map((t) => ({
-      type: 'function' as const,
-      function: {
-        name: t.name,
-        description: t.description,
-        parameters: t.parameters as Record<string, unknown>,
-      },
-    }));
-
-    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: {
-          name: 'ask',
-          description:
-            'Ask the user a clarifying question if the initial request is ambiguous or missing information.',
-          parameters: {
-            type: 'object',
-            properties: {
-              question: {
-                type: 'string',
-                description: 'The question to ask the user',
-              },
-            },
-            required: ['question'],
-          } as Record<string, unknown>,
-        },
-      });
-    }
-
-    // ReAct Loop
-    let iterations = 0;
-    const maxIterations = step.maxIterations || 10;
-    const totalUsage = {
-      prompt_tokens: 0,
-      completion_tokens: 0,
-      total_tokens: 0,
-    };
-
-    // Create redactor once outside the loop for performance (regex compilation)
-    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(truncatedMessages, {
-        model: resolvedModel,
-        tools: llmTools.length > 0 ? llmTools : undefined,
-        onStream: (chunk) => {
-          if (!step.outputSchema) {
-            process.stdout.write(redactionBuffer.process(chunk));
-          }
-        },
-        signal: abortSignal,
-      });
-
-      if (!step.outputSchema) {
-        process.stdout.write(redactionBuffer.flush());
-      }
-
-      if (response.usage) {
-        totalUsage.prompt_tokens += response.usage.prompt_tokens;
-        totalUsage.completion_tokens += response.usage.completion_tokens;
-        totalUsage.total_tokens += response.usage.total_tokens;
-      }
-
-      const { message } = response;
-      messages.push(message);
-
-      if (!message.tool_calls || message.tool_calls.length === 0) {
-        let output = message.content;
-
-        // If schema is defined, attempt to parse JSON
-        if (step.outputSchema && typeof output === 'string') {
-          try {
-            output = extractJson(output) as typeof output;
-          } catch (e) {
-            const errorMessage = `Failed to parse LLM output as JSON matching schema: ${e instanceof Error ? e.message : String(e)}`;
-            logger.error(`  ⚠️  ${errorMessage}. Retrying...`);
-
-            messages.push({
-              role: 'user',
-              content: `Error: ${errorMessage}\n\nPlease correct your output to be valid JSON matching the schema.`,
-            });
-            continue;
-          }
-        }
-
-        return {
-          output,
-          status: 'success',
-          usage: totalUsage,
-        };
-      }
-
-      // 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 {
-          const parsedArgs = JSON.parse(argsStr);
-          const keys = Object.keys(parsedArgs);
-          if (keys.length > 0) {
-            const formatted = JSON.stringify(parsedArgs);
-            displayArgs = formatted.length > 100 ? `${formatted.substring(0, 100)}...` : formatted;
-          }
-        } catch (e) {
-          displayArgs = argsStr.length > 100 ? `${argsStr.substring(0, 100)}...` : argsStr;
-        }
-
-        logger.log(
-          `  🛠️  Tool Call: ${toolCall.function.name}${displayArgs ? ` ${displayArgs}` : ''}`
-        );
-        const toolInfo = allTools.find((t) => t.name === toolCall.function.name);
-
-        if (!toolInfo) {
-          if (toolCall.function.name === 'ask' && step.allowClarification) {
-            let args: { question: string };
-            try {
-              args = JSON.parse(toolCall.function.arguments);
-            } catch (e) {
-              messages.push({
-                role: 'tool',
-                tool_call_id: toolCall.id,
-                name: 'ask',
-                content: formatToolContent(
-                  `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`
-                ),
-              });
-              continue;
-            }
-
-            if (process.stdin.isTTY) {
-              // In TTY, we can use a human step to get the answer immediately
-              logger.log(`\n🤔 Question from ${agent.name}: ${args.question}`);
-              const result = await executeStepFn(
-                {
-                  id: `${step.id}-clarify`,
-                  type: 'human',
-                  message: args.question,
-                  inputType: 'text',
-                } as Step,
-                context
-              );
-
-              messages.push({
-                role: 'tool',
-                tool_call_id: toolCall.id,
-                name: 'ask',
-                content: formatToolContent(String(result.output)),
-              });
-              continue;
-            }
-            // In non-TTY, we suspend
-            messages = truncateMessages(messages, maxHistory, maxConversationBytes);
-            return {
-              status: 'suspended',
-              output: {
-                messages,
-                question: args.question,
-              },
-              usage: totalUsage,
-            };
-          }
-
-          messages.push({
-            role: 'tool',
-            tool_call_id: toolCall.id,
-            name: toolCall.function.name,
-            content: formatToolContent(`Error: Tool ${toolCall.function.name} not found`),
-          });
-          continue;
-        }
-
-        let args: Record<string, unknown>;
-        try {
-          args = JSON.parse(toolCall.function.arguments);
-        } catch (e) {
-          messages.push({
-            role: 'tool',
-            tool_call_id: toolCall.id,
-            name: toolCall.function.name,
-            content: formatToolContent(
-              `Error: Invalid JSON in arguments: ${e instanceof Error ? e.message : String(e)}`
-            ),
-          });
-          continue;
-        }
-
-        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: formatToolContent(safeJsonStringify(result)),
-            });
-          } catch (error) {
-            messages.push({
-              role: 'tool',
-              tool_call_id: toolCall.id,
-              name: toolCall.function.name,
-              content: formatToolContent(
-                `Error: ${error instanceof Error ? error.message : String(error)}`
-              ),
-            });
-          }
-        } else if (toolInfo.execution) {
-          // Security validation for standard tools
-          if (toolInfo.source === 'standard') {
-            try {
-              validateStandardToolSecurity(toolInfo.name, args, {
-                allowOutsideCwd: step.allowOutsideCwd,
-                allowInsecure: step.allowInsecure,
-              });
-            } catch (error) {
-              messages.push({
-                role: 'tool',
-                tool_call_id: toolCall.id,
-                name: toolCall.function.name,
-                content: formatToolContent(
-                  `Security Error: ${error instanceof Error ? error.message : String(error)}`
-                ),
-              });
-              continue;
-            }
-          }
-
-          // Execute the tool as a step
-          const toolContext: ExpressionContext = {
-            ...context,
-            args, // Use args to pass parameters to tool execution
-          };
-
-          const result = await executeStepFn(toolInfo.execution, toolContext);
-
-          messages.push({
-            role: 'tool',
-            tool_call_id: toolCall.id,
-            name: toolCall.function.name,
-            content: formatToolContent(
-              result.status === 'success'
-                ? safeJsonStringify(result.output)
-                : `Error: ${result.error}`
-            ),
-          });
-        }
-      }
-
-      messages = truncateMessages(messages, maxHistory, maxConversationBytes);
-    }
-
-    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();
-    }
-  }
-}