keystone-cli 1.0.2 → 1.1.0

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

@@ -5,12 +5,14 @@ import { mkdirSync, writeFileSync } from 'node:fs';
 import { join } from 'node:path';
 import { Readable, Writable } from 'node:stream';
 import type { ExpressionContext } from '../expression/evaluator';
+import { ExpressionEvaluator } from '../expression/evaluator';
+import { parseAgent } from '../parser/agent-parser';
 import type { LlmStep, Step } from '../parser/schema';
+import { ConsoleLogger, type Logger } from '../utils/logger';
+import { executeLlmStep } from './executors/llm-executor.ts';
 import type { LLMAdapter, LLMMessage, LLMResponse, LLMTool } from './llm-adapter';
-import { executeLlmStep } from './llm-executor';
 import type { MCPServerConfig } from './mcp-manager';
 import type { StepResult } from './step-executor';
-import type { Logger } from './workflow-runner';
 
 // Mock adapters
 // Instead of mutating prototypes (which causes cross-test contamination),
@@ -187,6 +189,25 @@ tools:
 ---
 You are a test agent.`;
     writeFileSync(join(agentsDir, 'test-agent.md'), agentContent);
+
+    const handoffTargetContent = `---
+name: handoff-target
+model: gpt-4
+tools:
+  - name: specialist-tool
+    execution:
+      type: shell
+      run: echo "specialist"
+---
+You are the specialist for \${{ inputs.topic }}.`;
+    writeFileSync(join(agentsDir, 'handoff-target.md'), handoffTargetContent);
+
+    const contextAgentContent = `---
+name: context-agent
+model: gpt-4
+---
+You are a context-aware agent.`;
+    writeFileSync(join(agentsDir, 'context-agent.md'), contextAgentContent);
   });
 
   afterAll(() => {
@@ -230,7 +251,7 @@ You are a test agent.`;
     };
     const context: ExpressionContext = { inputs: {}, steps: {} };
 
-    const executeStepFn = async (s: Step) => {
+    const executeStepFn = async (s: any) => {
       if (s.type === 'shell') {
         return { status: 'success' as const, output: { stdout: 'tool result' } };
       }
@@ -262,7 +283,7 @@ You are a test agent.`;
     };
     const context: ExpressionContext = { inputs: {}, steps: {} };
 
-    const executeStepFn = async (s: Step) => {
+    const executeStepFn = async (s: any) => {
       if (s.type === 'shell') {
         return { status: 'success' as const, output: { stdout: 'tool result' } };
       }
@@ -273,6 +294,8 @@ You are a test agent.`;
       log: mock(() => {}),
       error: mock(() => {}),
       warn: mock(() => {}),
+      info: mock(() => {}),
+      debug: mock(() => {}),
     };
 
     await executeLlmStep(
@@ -325,6 +348,63 @@ You are a test agent.`;
     expect(result.output).toEqual({ foo: 'bar' });
   });
 
+  it('should accept native structured output tool calls when responseSchema is provided', async () => {
+    const outputSchema = {
+      type: 'object',
+      properties: {
+        foo: { type: 'string' },
+      },
+      required: ['foo'],
+    };
+    let receivedSchema: unknown;
+
+    const chatMock = mock(async (_messages, options) => {
+      receivedSchema = options?.responseSchema;
+      return {
+        message: {
+          role: 'assistant',
+          content: null,
+          tool_calls: [
+            {
+              id: 'call-1',
+              type: 'function',
+              function: { name: 'record_output', arguments: '{"foo":"bar"}' },
+            },
+          ],
+        },
+      };
+    }) as unknown as LLMAdapter['chat'];
+    const getAdapter = createMockGetAdapter(chatMock);
+
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'give me json',
+      needs: [],
+      maxIterations: 5,
+      outputSchema,
+    };
+    const context: ExpressionContext = { inputs: {}, steps: {} };
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    const result = await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(receivedSchema).toEqual(outputSchema);
+    expect(result.status).toBe('success');
+    expect(result.output).toEqual({ foo: 'bar' });
+    expect(executeStepFn).not.toHaveBeenCalled();
+  });
+
   it('should retry if LLM output fails schema validation', async () => {
     const step: LlmStep = {
       id: 'l1',
@@ -468,7 +548,7 @@ You are a test agent.`;
       context,
       executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
       console,
-      mcpManager as unknown as { getClient: () => Promise<unknown> },
+      mcpManager as any,
       undefined,
       undefined,
       mockGetAdapter
@@ -527,7 +607,7 @@ You are a test agent.`;
       context,
       executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
       undefined,
-      mcpManager as unknown as { getClient: () => Promise<unknown> },
+      mcpManager as any,
       undefined,
       undefined,
       getAdapter
@@ -570,10 +650,7 @@ You are a test agent.`;
       context,
       executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
       console,
-      manager as unknown as {
-        getClient: () => Promise<unknown>;
-        getGlobalServers: () => unknown[];
-      },
+      manager as any,
       undefined,
       undefined,
       getAdapter
@@ -603,7 +680,8 @@ You are a test agent.`;
     };
     const context: ExpressionContext = { inputs: {}, steps: {} };
     let toolExecuted = false;
-    const executeStepFn = async (s: Step) => {
+
+    const executeStepFn = async (s: any) => {
       if (s.id === 'adhoc-step') {
         toolExecuted = true;
         return { status: 'success' as const, output: { stdout: 'adhoc result' } };
@@ -691,7 +769,7 @@ You are a test agent.`;
       getAdapter
     );
 
-    expect(capturedStep?.type).toBe('engine');
+    expect((capturedStep as any)?.type).toBe('engine');
     expect(chatCount).toBe(2);
   });
 
@@ -728,6 +806,194 @@ You are a test agent.`;
     consoleSpy.mockRestore();
   });
 
+  it('should summarize messages when history is too long', async () => {
+    let summaryAttempted = false;
+    const chatMock = mock(async (messages: LLMMessage[]) => {
+      if (messages.find((m) => m.name === 'context_summary')) {
+        summaryAttempted = true;
+      }
+      return { message: { role: 'assistant', content: 'Resuming' } };
+    }) as unknown as LLMAdapter['chat'];
+
+    const getAdapter = (modelString: string) => {
+      const mockAdapter: LLMAdapter = {
+        chat: async (messages, options) => {
+          if (messages[0].role === 'system' && messages[0].content?.includes('Summarize')) {
+            return { message: { role: 'assistant', content: 'Summary text' } };
+          }
+          return chatMock(messages, options);
+        },
+      };
+      return { adapter: mockAdapter, resolvedModel: 'gpt-4' };
+    };
+
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'continue',
+      needs: [],
+      maxIterations: 1,
+      maxMessageHistory: 4, // Allow at least one non-system message before summarization
+      contextStrategy: 'summary',
+    };
+
+    const context: ExpressionContext = {
+      inputs: {},
+      steps: {
+        l1: {
+          output: {
+            messages: [
+              { role: 'user', content: 'm1' },
+              { role: 'assistant', content: 'm2' },
+              { role: 'user', content: 'm3' },
+            ],
+          },
+        },
+      },
+    };
+
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(summaryAttempted).toBe(true);
+  });
+
+  it('should fall back to truncation if summarization fails', async () => {
+    const logger: Logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+      debug: mock(() => {}),
+    };
+
+    const getAdapter = (modelString: string) => {
+      const mockAdapter: LLMAdapter = {
+        chat: async (messages) => {
+          if (messages[0].role === 'system' && messages[0].content?.includes('Summarize')) {
+            throw new Error('Summary failed');
+          }
+          return { message: { role: 'assistant', content: 'Truncated response' } };
+        },
+      };
+      return { adapter: mockAdapter, resolvedModel: 'gpt-4' };
+    };
+
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'continue',
+      needs: [],
+      maxIterations: 1,
+      maxMessageHistory: 4,
+      contextStrategy: 'summary',
+    };
+
+    const context: ExpressionContext = {
+      inputs: {},
+      steps: {
+        l1: {
+          output: {
+            messages: [
+              { role: 'user', content: 'm1' },
+              { role: 'assistant', content: 'm2' },
+              { role: 'user', content: 'm3' },
+            ],
+          },
+        },
+      },
+    };
+
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      logger,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(logger.warn).toHaveBeenCalledWith(
+      expect.stringContaining('Context summarization failed')
+    );
+  });
+
+  it('should extract thought blocks and emit thought events', async () => {
+    const logger: Logger = {
+      log: mock(() => {}),
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+      debug: mock(() => {}),
+    };
+
+    const emitEvent = mock(() => {});
+    const eventContext = { runId: 'run-1', workflow: 'wf-1' };
+
+    const chatMock = mock(async () => {
+      return {
+        message: {
+          role: 'assistant',
+          content: '<thinking>I should do X</thinking>Final answer',
+        },
+      };
+    }) as unknown as LLMAdapter['chat'];
+
+    const getAdapter = () => ({
+      adapter: { chat: chatMock },
+      resolvedModel: 'gpt-4',
+    });
+
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'hello',
+      needs: [],
+      maxIterations: 10,
+    };
+
+    await executeLlmStep(
+      step,
+      { inputs: {}, steps: {} },
+      mock(async () => ({ status: 'success' as const, output: 'ok' })) as any,
+      logger,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter as any,
+      emitEvent,
+      eventContext
+    );
+
+    expect(logger.info).toHaveBeenCalledWith(
+      expect.stringContaining('Thought (thinking): I should do X')
+    );
+    expect(emitEvent).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: 'llm.thought',
+        content: 'I should do X',
+        source: 'thinking',
+      })
+    );
+  });
+
   it('should not add global MCP server if already explicitly listed', async () => {
     const mockClient = createMockMcpClient();
     const manager = createMockMcpManager({
@@ -757,10 +1023,7 @@ You are a test agent.`;
       context,
       executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
       console,
-      manager as unknown as {
-        getClient: () => Promise<unknown>;
-        getGlobalServers: () => unknown[];
-      },
+      manager as any,
       undefined,
       undefined,
       getAdapter
@@ -788,7 +1051,6 @@ You are a test agent.`;
 
     let capturedPrompt = '';
     const chatMock = mock(async (messages: LLMMessage[]) => {
-      // console.log('MESSAGES:', JSON.stringify(messages, null, 2));
       capturedPrompt = messages.find((m) => m.role === 'user')?.content || '';
       return { message: { role: 'assistant', content: 'Response' } };
     }) as unknown as LLMAdapter['chat'];
@@ -810,4 +1072,257 @@ You are a test agent.`;
     expect(capturedPrompt).toContain('"key": "value"');
     expect(capturedPrompt).not.toContain('[object Object]');
   });
+
+  it('should evaluate expressions in agent system prompts', async () => {
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'handoff-target',
+      prompt: 'hello',
+      needs: [],
+      maxIterations: 3,
+    };
+    const context: ExpressionContext = { inputs: { topic: 'payments' }, steps: {} };
+    let capturedSystem = '';
+
+    const chatMock = mock(async (messages: LLMMessage[]) => {
+      const systemMessages = messages.filter((m) => m.role === 'system');
+      capturedSystem =
+        (systemMessages.find((m) => typeof m.content === 'string')?.content as string) || '';
+      return { message: { role: 'assistant', content: 'ok' } };
+    }) as unknown as LLMAdapter['chat'];
+    const getAdapter = createMockGetAdapter(chatMock);
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    const result = await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(result.status).toBe('success');
+    expect(capturedSystem).toContain('payments');
+    expect(capturedSystem).not.toContain('${{');
+  });
+
+  it('should handle streaming chunks with thoughts', async () => {
+    const step = {
+      id: 'l-stream',
+      type: 'llm' as const,
+      agent: 'test-agent',
+      prompt: 'stream this',
+      needs: [],
+      maxIterations: 1,
+    };
+
+    // We can't easily add 'stream' to LlmStep without changing schema,
+    // but we can mock the adapter to stream if onStream is provided.
+
+    const chatMock = mock(async (messages: LLMMessage[], options: any) => {
+      if (options.onStream) {
+        options.onStream('<thinking>thought</thinking>done');
+      }
+      return { message: { role: 'assistant', content: '<thinking>thought</thinking>done' } };
+    }) as unknown as LLMAdapter['chat'];
+
+    const adapter = {
+      chat: chatMock,
+    } as any;
+
+    const context: ExpressionContext = { inputs: {}, steps: {} };
+    spyOn(process.stdout, 'write').mockImplementation(() => true);
+
+    const emitThought = mock(() => {});
+
+    await executeLlmStep(
+      step as any,
+      context,
+      mock(async () => ({ status: 'success' as const, output: 'ok' })) as any,
+      new ConsoleLogger(),
+      undefined,
+      undefined,
+      undefined,
+      () => ({ adapter, resolvedModel: 'test-model' }),
+      emitThought,
+      { runId: 'test-run', workflow: 'test-wf' }
+    );
+
+    expect(emitThought).toHaveBeenCalled();
+  });
+
+  it('should transfer to allowed agent and swap system prompt/tools', async () => {
+    let callCount = 0;
+    let sawTransferTool = false;
+    let sawOriginalTool = false;
+    let sawTargetToolAfter = false;
+    let sawOriginalToolAfter = false;
+    let sawTargetPrompt = false;
+
+    const chatMock = mock(async (messages: LLMMessage[], options: { tools?: LLMTool[] }) => {
+      callCount++;
+      const toolNames = options.tools?.map((t) => t.function.name) || [];
+
+      if (callCount === 1) {
+        sawTransferTool = toolNames.includes('transfer_to_agent');
+        sawOriginalTool = toolNames.includes('test-tool');
+        return {
+          message: {
+            role: 'assistant',
+            content: null,
+            tool_calls: [
+              {
+                id: 'call-transfer',
+                type: 'function',
+                function: {
+                  name: 'transfer_to_agent',
+                  arguments: '{"agent_name":"handoff-target"}',
+                },
+              },
+            ],
+          },
+        };
+      }
+
+      const systemMessages = messages.filter((m) => m.role === 'system');
+      sawTargetPrompt = systemMessages.some(
+        (m) => typeof m.content === 'string' && m.content.includes('specialist for billing')
+      );
+      sawTargetToolAfter = toolNames.includes('specialist-tool');
+      sawOriginalToolAfter = toolNames.includes('test-tool');
+
+      return {
+        message: { role: 'assistant', content: 'done' },
+      };
+    }) as unknown as LLMAdapter['chat'];
+    const getAdapter = createMockGetAdapter(chatMock);
+
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'test-agent',
+      prompt: 'handoff',
+      needs: [],
+      maxIterations: 4,
+      allowedHandoffs: ['handoff-target'],
+    };
+    const context: ExpressionContext = { inputs: { topic: 'billing' }, steps: {} };
+    const executeStepFn = mock(async () => ({ status: 'success' as const, output: 'ok' }));
+
+    const result = await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(result.status).toBe('success');
+    expect(sawTransferTool).toBe(true);
+    expect(sawOriginalTool).toBe(true);
+    expect(sawTargetToolAfter).toBe(true);
+    expect(sawOriginalToolAfter).toBe(false);
+    expect(sawTargetPrompt).toBe(true);
+  });
+
+  it('should apply context updates from tool output', async () => {
+    const step: LlmStep = {
+      id: 'l1',
+      type: 'llm',
+      agent: 'context-agent',
+      prompt: 'update context',
+      needs: [],
+      maxIterations: 4,
+      tools: [
+        {
+          name: 'update-context',
+          execution: {
+            id: 'update-step',
+            type: 'shell',
+            run: 'echo update',
+          },
+        },
+        {
+          name: 'read-context',
+          execution: {
+            id: 'read-step',
+            type: 'shell',
+            run: 'echo read',
+          },
+        },
+      ],
+    };
+    const context: ExpressionContext = { inputs: {}, steps: {} };
+    let sawEnvUpdate = false;
+    let sawMemoryUpdate = false;
+
+    const executeStepFn = async (_step: any, toolContext: ExpressionContext) => {
+      if (_step.id === 'update-step') {
+        return {
+          status: 'success' as const,
+          output: {
+            __keystone_context: {
+              env: { USER_ID: '123' },
+              memory: { user: 'Ada' },
+            },
+            ok: true,
+          },
+        };
+      }
+      if (_step.id === 'read-step') {
+        sawEnvUpdate = toolContext.env?.USER_ID === '123';
+        sawMemoryUpdate = toolContext.memory?.user === 'Ada';
+        return { status: 'success' as const, output: { seen: true } };
+      }
+      return { status: 'success' as const, output: 'ok' };
+    };
+
+    let callCount = 0;
+    const chatMock = mock(async () => {
+      callCount++;
+      if (callCount === 1) {
+        return {
+          message: {
+            role: 'assistant',
+            content: null,
+            tool_calls: [
+              {
+                id: 'call-update',
+                type: 'function',
+                function: { name: 'update-context', arguments: '{}' },
+              },
+              {
+                id: 'call-read',
+                type: 'function',
+                function: { name: 'read-context', arguments: '{}' },
+              },
+            ],
+          },
+        };
+      }
+      return { message: { role: 'assistant', content: 'done' } };
+    }) as unknown as LLMAdapter['chat'];
+    const getAdapter = createMockGetAdapter(chatMock);
+
+    await executeLlmStep(
+      step,
+      context,
+      executeStepFn as unknown as (step: Step, context: ExpressionContext) => Promise<StepResult>,
+      undefined,
+      undefined,
+      undefined,
+      undefined,
+      getAdapter
+    );
+
+    expect(sawEnvUpdate).toBe(true);
+    expect(sawMemoryUpdate).toBe(true);
+  });
 });

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

@@ -17,8 +17,7 @@ describe('MCPClient Audit Fixes', () => {
       }),
       kill: () => {},
       on: () => {},
-      // biome-ignore lint/suspicious/noExplicitAny: Mocking complex object
-    } as any);
+    } as unknown as child_process.ChildProcess);
   });
 
   afterEach(() => {