keystone-cli 0.5.1 → 0.6.1

package/src/runner/llm-adapter.ts

@@ -1,5 +1,7 @@
+import { pipeline } from '@xenova/transformers';
 import { AuthManager, COPILOT_HEADERS } from '../utils/auth-manager';
 import { ConfigLoader } from '../utils/config-loader';
+import { processOpenAIStream } from './stream-utils';
 
 // Maximum response size to prevent memory exhaustion (1MB)
 const MAX_RESPONSE_SIZE = 1024 * 1024;
@@ -48,6 +50,7 @@ export interface LLMAdapter {
       onStream?: (chunk: string) => void;
     }
   ): Promise<LLMResponse>;
+  embed?(text: string, model?: string): Promise<number[]>;
 }
 
 export class OpenAIAdapter implements LLMAdapter {
@@ -94,72 +97,51 @@ export class OpenAIAdapter implements LLMAdapter {
 
     if (isStreaming) {
       if (!response.body) throw new Error('Response body is null');
-      const reader = response.body.getReader();
-      const decoder = new TextDecoder();
-      let fullContent = '';
-      const toolCalls: LLMToolCall[] = [];
-
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) break;
-
-        const chunk = decoder.decode(value);
-        const lines = chunk.split('\n').filter((line) => line.trim() !== '');
-
-        for (const line of lines) {
-          if (line.includes('[DONE]')) continue;
-          if (!line.startsWith('data: ')) continue;
-
-          try {
-            const data = JSON.parse(line.slice(6));
-            const delta = data.choices[0].delta;
-
-            if (delta.content) {
-              if (fullContent.length + delta.content.length > MAX_RESPONSE_SIZE) {
-                throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
-              }
-              fullContent += delta.content;
-              options.onStream?.(delta.content);
-            }
-
-            if (delta.tool_calls) {
-              for (const tc of delta.tool_calls) {
-                if (!toolCalls[tc.index]) {
-                  toolCalls[tc.index] = {
-                    id: tc.id,
-                    type: 'function',
-                    function: { name: '', arguments: '' },
-                  };
-                }
-                const existing = toolCalls[tc.index];
-                if (tc.function?.name) existing.function.name += tc.function.name;
-                if (tc.function?.arguments) existing.function.arguments += tc.function.arguments;
-              }
-            }
-          } catch (e) {
-            // Ignore parse errors for incomplete chunks
-          }
-        }
-      }
-
-      return {
-        message: {
-          role: 'assistant',
-          content: fullContent || null,
-          tool_calls: toolCalls.length > 0 ? toolCalls.filter(Boolean) : undefined,
-        },
-      };
+      return processOpenAIStream(response, options, 'OpenAI');
     }
 
     const data = (await response.json()) as {
       choices: { message: LLMMessage }[];
       usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
     };
+
+    // Validate response size to prevent memory exhaustion
+    const contentLength = data.choices[0]?.message?.content?.length ?? 0;
+    if (contentLength > MAX_RESPONSE_SIZE) {
+      throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
+    }
+
     return {
       message: data.choices[0].message,
       usage: data.usage,
     };
   }
+
+  async embed(text: string, model = 'text-embedding-3-small'): Promise<number[]> {
+    const response = await fetch(`${this.baseUrl}/embeddings`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify({
+        model,
+        input: text,
+      }),
+    });
+
+    if (!response.ok) {
+      const error = await response.text();
+      throw new Error(
+        `OpenAI Embeddings API error: ${response.status} ${response.statusText} - ${error}`
+      );
+    }
+
+    const data = (await response.json()) as {
+      data: { embedding: number[] }[];
+    };
+    return data.data[0].embedding;
+  }
 }
 
 export class AnthropicAdapter implements LLMAdapter {
@@ -348,7 +330,15 @@ export class AnthropicAdapter implements LLMAdapter {
               }
             }
           } catch (e) {
-            // Ignore parse errors
+            // Log non-SyntaxError exceptions at warning level (they indicate real issues)
+            if (!(e instanceof SyntaxError)) {
+              console.warn(`[Anthropic Stream] Error processing chunk: ${e}`);
+            } else if (process.env.DEBUG || process.env.LLM_DEBUG) {
+              // SyntaxErrors are normal for incomplete chunks - only log in debug mode
+              process.stderr.write(
+                `[Anthropic Stream] Incomplete chunk parse: ${line.slice(0, 50)}...\n`
+              );
+            }
           }
         }
       }
@@ -383,6 +373,12 @@ export class AnthropicAdapter implements LLMAdapter {
     };
 
     const content = data.content.find((c) => c.type === 'text')?.text || null;
+
+    // Validate response size to prevent memory exhaustion
+    if (content && content.length > MAX_RESPONSE_SIZE) {
+      throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
+    }
+
     const toolCalls = data.content
       .filter((c) => c.type === 'tool_use')
       .map((c) => ({
@@ -455,68 +451,20 @@ export class CopilotAdapter implements LLMAdapter {
     if (isStreaming) {
       // Use the same streaming logic as OpenAIAdapter since Copilot uses OpenAI API
       if (!response.body) throw new Error('Response body is null');
-      const reader = response.body.getReader();
-      const decoder = new TextDecoder();
-      let fullContent = '';
-      const toolCalls: LLMToolCall[] = [];
-
-      while (true) {
-        const { done, value } = await reader.read();
-        if (done) break;
-
-        const chunk = decoder.decode(value);
-        const lines = chunk.split('\n').filter((line) => line.trim() !== '');
-
-        for (const line of lines) {
-          if (line.includes('[DONE]')) continue;
-          if (!line.startsWith('data: ')) continue;
-
-          try {
-            const data = JSON.parse(line.slice(6));
-            if (!data.choices?.[0]?.delta) continue;
-            const delta = data.choices[0].delta;
-
-            if (delta.content) {
-              if (fullContent.length + delta.content.length > MAX_RESPONSE_SIZE) {
-                throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
-              }
-              fullContent += delta.content;
-              options.onStream?.(delta.content);
-            }
-
-            if (delta.tool_calls) {
-              for (const tc of delta.tool_calls) {
-                if (!toolCalls[tc.index]) {
-                  toolCalls[tc.index] = {
-                    id: tc.id,
-                    type: 'function',
-                    function: { name: '', arguments: '' },
-                  };
-                }
-                const existing = toolCalls[tc.index];
-                if (tc.function?.name) existing.function.name += tc.function.name;
-                if (tc.function?.arguments) existing.function.arguments += tc.function.arguments;
-              }
-            }
-          } catch (e) {
-            // Ignore parse errors
-          }
-        }
-      }
-
-      return {
-        message: {
-          role: 'assistant',
-          content: fullContent || null,
-          tool_calls: toolCalls.length > 0 ? toolCalls.filter(Boolean) : undefined,
-        },
-      };
+      return processOpenAIStream(response, options, 'Copilot');
     }
 
     const data = (await response.json()) as {
       choices: { message: LLMMessage }[];
       usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
     };
+
+    // Validate response size to prevent memory exhaustion
+    const contentLength = data.choices[0]?.message?.content?.length ?? 0;
+    if (contentLength > MAX_RESPONSE_SIZE) {
+      throw new Error(`LLM response exceeds maximum size of ${MAX_RESPONSE_SIZE} bytes`);
+    }
+
     return {
       message: data.choices[0].message,
       usage: data.usage,
@@ -524,7 +472,37 @@ export class CopilotAdapter implements LLMAdapter {
   }
 }
 
+export class LocalEmbeddingAdapter implements LLMAdapter {
+  // biome-ignore lint/suspicious/noExplicitAny: transformers pipeline type
+  private static extractor: any = null;
+
+  async chat(): Promise<LLMResponse> {
+    throw new Error(
+      'Local models in Keystone currently only support memory/embedding operations. ' +
+        'To use a local LLM for chat/generation, please use an OpenAI-compatible local server ' +
+        '(like Ollama, LM Studio, or LocalAI) and configure it as an OpenAI provider in your config.'
+    );
+  }
+
+  async embed(text: string, model = 'Xenova/all-MiniLM-L6-v2'): Promise<number[]> {
+    const modelToUse = model === 'local' ? 'Xenova/all-MiniLM-L6-v2' : model;
+    if (!LocalEmbeddingAdapter.extractor) {
+      LocalEmbeddingAdapter.extractor = await pipeline('feature-extraction', modelToUse);
+    }
+    const output = await LocalEmbeddingAdapter.extractor(text, {
+      pooling: 'mean',
+      normalize: true,
+    });
+    return Array.from(output.data);
+  }
+}
+
 export function getAdapter(model: string): { adapter: LLMAdapter; resolvedModel: string } {
+  if (model === 'local' || model.startsWith('local:')) {
+    const resolvedModel = model === 'local' ? 'Xenova/all-MiniLM-L6-v2' : model.substring(6);
+    return { adapter: new LocalEmbeddingAdapter(), resolvedModel };
+  }
+
   const providerName = ConfigLoader.getProviderForModel(model);
   const config = ConfigLoader.load();
   const providerConfig = config.providers[providerName];

package/src/runner/llm-executor.test.ts

@@ -28,7 +28,8 @@ import {
 import { executeLlmStep } from './llm-executor';
 import { MCPClient, type MCPResponse } from './mcp-client';
 import { MCPManager } from './mcp-manager';
-import type { StepResult } from './step-executor';
+import { type StepResult, executeStep } from './step-executor';
+import type { Logger } from './workflow-runner';
 
 // Mock adapters
 const originalOpenAIChat = OpenAIAdapter.prototype.chat;
@@ -129,16 +130,18 @@ describe('llm-executor', () => {
     };
   };
 
-  beforeAll(() => {
+  beforeAll(async () => {
     // Mock spawn to avoid actual process creation
     const mockProcess = Object.assign(new EventEmitter(), {
-      stdout: new Readable({ read() { } }),
+      stdout: new Readable({
+        read() {},
+      }),
       stdin: new Writable({
         write(_chunk, _encoding, cb: (error?: Error | null) => void) {
           cb();
         },
       }),
-      kill: mock(() => { }),
+      kill: mock(() => {}),
     });
     spawnSpy = spyOn(child_process, 'spawn').mockReturnValue(
       mockProcess as unknown as child_process.ChildProcess
@@ -239,6 +242,44 @@ You are a test agent.`;
     expect(result.output).toBe('LLM Response');
   });
 
+  it('should log tool call arguments', async () => {
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'trigger tool',
+      needs: [],
+      maxIterations: 10,
+    };
+    const context: ExpressionContext = { inputs: {}, steps: {} };
+
+    const executeStepFn = async (s: Step) => {
+      if (s.type === 'shell') {
+        return { status: 'success' as const, output: { stdout: 'tool result' } };
+      }
+      return { status: 'success' as const, output: 'ok' };
+    };
+
+    const logger: Logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+    };
+
+    await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      logger
+    );
+
+    // Check if logger.log was called with arguments
+    // The tool call from mockChat is { name: 'test-tool', arguments: '{"val": 123}' }
+    expect(logger.log).toHaveBeenCalledWith(
+      expect.stringContaining('🛠️  Tool Call: test-tool {"val":123}')
+    );
+  });
+
   it('should support schema for JSON output', async () => {
     const step: LlmStep = {
       id: 'l1',
@@ -422,7 +463,7 @@ You are a test agent.`;
       spyOn(client, 'stop').mockReturnValue(undefined);
       return client;
     });
-    const consoleSpy = spyOn(console, 'error').mockImplementation(() => { });
+    const consoleSpy = spyOn(console, 'error').mockImplementation(() => {});
 
     await executeLlmStep(
       step,
@@ -609,7 +650,7 @@ You are a test agent.`;
     };
     const context: ExpressionContext = { inputs: {}, steps: {} };
     const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
-    const consoleSpy = spyOn(console, 'error').mockImplementation(() => { });
+    const consoleSpy = spyOn(console, 'error').mockImplementation(() => {});
 
     await executeLlmStep(
       step,

package/src/runner/llm-executor.ts

@@ -4,12 +4,12 @@ import { ExpressionEvaluator } from '../expression/evaluator';
 import { parseAgent, resolveAgentPath } from '../parser/agent-parser';
 import type { AgentTool, LlmStep, Step } from '../parser/schema';
 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 type { StepResult } from './step-executor';
-import type { Logger } from './workflow-runner';
 
 interface ToolDefinition {
   name: string;
@@ -24,7 +24,7 @@ export async function executeLlmStep(
   step: LlmStep,
   context: ExpressionContext,
   executeStepFn: (step: Step, context: ExpressionContext) => Promise<StepResult>,
-  logger: Logger = console,
+  logger: Logger = new ConsoleLogger(),
   mcpManager?: MCPManager,
   workflowDir?: string
 ): Promise<StepResult> {
@@ -269,7 +269,22 @@ export async function executeLlmStep(
 
       // Execute tools
       for (const toolCall of message.tool_calls) {
-        logger.log(`  🛠️  Tool Call: ${toolCall.function.name}`);
+        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) {

package/src/runner/mcp-client.audit.test.ts

@@ -77,3 +77,72 @@ describe('MCPClient Audit Fixes', () => {
     }
   });
 });
+
+describe('MCPClient SSRF Protection', () => {
+  it('should reject localhost URLs without allowInsecure', async () => {
+    // HTTP localhost is rejected for not using HTTPS
+    await expect(MCPClient.createRemote('http://localhost:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*HTTPS/
+    );
+    // HTTPS localhost is rejected for being localhost
+    await expect(MCPClient.createRemote('https://localhost:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*localhost/
+    );
+  });
+
+  it('should reject 127.0.0.1', async () => {
+    await expect(MCPClient.createRemote('https://127.0.0.1:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*localhost/
+    );
+  });
+
+  it('should reject private IP ranges (10.x.x.x)', async () => {
+    await expect(MCPClient.createRemote('https://10.0.0.1:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*private/
+    );
+  });
+
+  it('should reject private IP ranges (192.168.x.x)', async () => {
+    await expect(MCPClient.createRemote('https://192.168.1.1:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*private/
+    );
+  });
+
+  it('should reject private IP ranges (172.16-31.x.x)', async () => {
+    await expect(MCPClient.createRemote('https://172.16.0.1:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*private/
+    );
+    await expect(MCPClient.createRemote('https://172.31.255.1:8080/sse')).rejects.toThrow(
+      /SSRF Protection.*private/
+    );
+  });
+
+  it('should reject cloud metadata endpoints', async () => {
+    // 169.254.169.254 is caught by link-local IP range check
+    await expect(
+      MCPClient.createRemote('https://169.254.169.254/latest/meta-data/')
+    ).rejects.toThrow(/SSRF Protection.*private/);
+    // Also test the hostname-based metadata detection
+    await expect(MCPClient.createRemote('https://metadata.google.internal/sse')).rejects.toThrow(
+      /SSRF Protection.*metadata/
+    );
+  });
+
+  it('should require HTTPS by default', async () => {
+    await expect(MCPClient.createRemote('http://api.example.com/sse')).rejects.toThrow(
+      /SSRF Protection.*HTTPS/
+    );
+  });
+
+  it('should allow HTTP with allowInsecure option', async () => {
+    // This will fail due to network issues, not SSRF
+    const promise = MCPClient.createRemote(
+      'http://api.example.com/sse',
+      {},
+      100, // short timeout
+      { allowInsecure: true }
+    );
+    // Should NOT throw SSRF error, but will throw timeout/connection error
+    await expect(promise).rejects.not.toThrow(/SSRF Protection/);
+  });
+});

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

@@ -139,7 +139,10 @@ describe('MCPClient', () => {
         return Promise.resolve(new Response(JSON.stringify({ ok: true })));
       });
 
-      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse');
+      // Use allowInsecure for testing with localhost (fetch is mocked anyway)
+      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse', {}, 60000, {
+        allowInsecure: true,
+      });
 
       const client = await clientPromise;
       expect(client).toBeDefined();
@@ -185,7 +188,10 @@ describe('MCPClient', () => {
         return Promise.resolve(new Response(JSON.stringify({ ok: true })));
       });
 
-      const client = await MCPClient.createRemote('http://localhost:8080/sse');
+      // Use allowInsecure for testing with localhost (fetch is mocked anyway)
+      const client = await MCPClient.createRemote('http://localhost:8080/sse', {}, 60000, {
+        allowInsecure: true,
+      });
 
       // We can't easily hook into onMessage without reaching into internals
       // Instead, we'll test that initialize resolves correctly when the response arrives
@@ -228,7 +234,10 @@ describe('MCPClient', () => {
         )
       );
 
-      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse');
+      // Use allowInsecure for testing with localhost (fetch is mocked anyway)
+      const clientPromise = MCPClient.createRemote('http://localhost:8080/sse', {}, 60000, {
+        allowInsecure: true,
+      });
 
       await expect(clientPromise).rejects.toThrow(/SSE connection failed: 500/);