moving

packages/agent/tests/subagent-orchestrator.test.ts

@@ -0,0 +1,480 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { SubagentOrchestrator, type WorkflowTemplate, type OrchestratorConfig } from '../src/subagent-orchestrator.js';
+import type { ToolDefinition } from '../src/tools.js';
+import type { AgentLoopConfig, AgentResponse } from '../src/agent-loop.js';
+
+function makeMockTools(): ToolDefinition[] {
+  return [
+    { name: 'web_search', description: 'Search', parameters: { type: 'object', properties: {} }, execute: async () => 'results' },
+    { name: 'web_fetch', description: 'Fetch', parameters: { type: 'object', properties: {} }, execute: async () => 'content' },
+    { name: 'read_file', description: 'Read', parameters: { type: 'object', properties: {} }, execute: async () => 'file content' },
+    { name: 'write_file', description: 'Write', parameters: { type: 'object', properties: {} }, execute: async () => 'ok' },
+    { name: 'bash', description: 'Shell', parameters: { type: 'object', properties: {} }, execute: async () => 'output' },
+    { name: 'search_memory', description: 'Memory', parameters: { type: 'object', properties: {} }, execute: async () => 'memories' },
+    { name: 'save_memory', description: 'Save', parameters: { type: 'object', properties: {} }, execute: async () => 'saved' },
+    { name: 'search_files', description: 'Search files', parameters: { type: 'object', properties: {} }, execute: async () => 'files' },
+    { name: 'search_content', description: 'Search content', parameters: { type: 'object', properties: {} }, execute: async () => 'content' },
+  ];
+}
+
+function makeMockRunner() {
+  return vi.fn(async (config: AgentLoopConfig): Promise<AgentResponse> => ({
+    content: `Result for: ${config.messages[0]?.content}`,
+    usage: { inputTokens: 100, outputTokens: 50 },
+    toolsUsed: ['web_search'],
+  }));
+}
+
+function makeConfig(runLoop?: ReturnType<typeof makeMockRunner>): OrchestratorConfig {
+  return {
+    availableTools: makeMockTools(),
+    runLoop: runLoop ?? makeMockRunner(),
+    litellmUrl: 'http://localhost:4000',
+    litellmApiKey: 'test-key',
+    defaultModel: 'test-model',
+  };
+}
+
+describe('SubagentOrchestrator', () => {
+  let orchestrator: SubagentOrchestrator;
+  let runner: ReturnType<typeof makeMockRunner>;
+
+  beforeEach(() => {
+    runner = makeMockRunner();
+    orchestrator = new SubagentOrchestrator(makeConfig(runner));
+  });
+
+  it('constructor creates orchestrator', () => {
+    expect(orchestrator).toBeInstanceOf(SubagentOrchestrator);
+    expect(orchestrator.getWorkers()).toEqual([]);
+  });
+
+  it('runWorkflow executes steps in order', async () => {
+    const callOrder: string[] = [];
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      // Extract the step name from the system prompt
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      callOrder.push(match?.[1] ?? 'unknown');
+      return {
+        content: `Result from ${match?.[1]}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: ['web_search'],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'test-workflow',
+      description: 'Test workflow',
+      steps: [
+        { name: 'Step A', role: 'researcher', task: 'Research something' },
+        { name: 'Step B', role: 'writer', task: 'Write something' },
+      ],
+      aggregation: 'concatenate',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    expect(callOrder).toEqual(['Step A', 'Step B']);
+    expect(result.results.size).toBe(2);
+    expect(runner).toHaveBeenCalledTimes(2);
+  });
+
+  it('dependsOn is respected — step B waits for step A', async () => {
+    const callOrder: string[] = [];
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      callOrder.push(match?.[1] ?? 'unknown');
+      return {
+        content: `Done: ${match?.[1]}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: [],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'dep-workflow',
+      description: 'Dependency workflow',
+      steps: [
+        { name: 'Step A', role: 'researcher', task: 'First task' },
+        { name: 'Step B', role: 'writer', task: 'Second task', dependsOn: ['Step A'] },
+      ],
+      aggregation: 'concatenate',
+    };
+
+    await orchestrator.runWorkflow(template);
+    expect(callOrder).toEqual(['Step A', 'Step B']);
+  });
+
+  it('runs each dependency-ready wave concurrently and honors per-step models', async () => {
+    const started: string[] = [];
+    const releases = new Map<string, () => void>();
+    let active = 0;
+    let maxActive = 0;
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const name = config.systemPrompt.match(/Sub-Agent: (.+)/)?.[1] ?? 'unknown';
+      started.push(name);
+      active++;
+      maxActive = Math.max(maxActive, active);
+      if (name !== 'Step C') {
+        await new Promise<void>((resolve) => releases.set(name, resolve));
+      }
+      active--;
+      return {
+        content: `Done: ${name}`,
+        usage: { inputTokens: 1, outputTokens: 1 },
+        toolsUsed: [],
+      };
+    });
+
+    const running = orchestrator.runWorkflow({
+      name: 'parallel-wave',
+      description: 'Parallel work followed by a dependent step',
+      steps: [
+        { name: 'Step A', role: 'researcher', task: 'A', model: 'model-a' },
+        { name: 'Step B', role: 'writer', task: 'B', model: 'model-b' },
+        {
+          name: 'Step C', role: 'analyst', task: 'C', model: 'model-c',
+          dependsOn: ['Step A', 'Step B'], contextFrom: ['Step A', 'Step B'],
+        },
+      ],
+      aggregation: 'last',
+    });
+
+    await vi.waitFor(() => expect(started).toEqual(['Step A', 'Step B']));
+    expect(maxActive).toBe(2);
+    expect(started).not.toContain('Step C');
+    releases.get('Step A')?.();
+    releases.get('Step B')?.();
+
+    const result = await running;
+    expect(started).toEqual(['Step A', 'Step B', 'Step C']);
+    expect(runner.mock.calls.map(([config]) => config.model)).toEqual(['model-a', 'model-b', 'model-c']);
+    expect([...result.results.values()].map((worker) => worker.model)).toEqual(['model-a', 'model-b', 'model-c']);
+  });
+
+  it('contextFrom injects previous results into system prompt', async () => {
+    const systemPrompts: string[] = [];
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      systemPrompts.push(config.systemPrompt);
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      return {
+        content: `Result from ${match?.[1]}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: [],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'context-workflow',
+      description: 'Context injection workflow',
+      steps: [
+        { name: 'Research', role: 'researcher', task: 'Find data' },
+        { name: 'Write', role: 'writer', task: 'Write report', contextFrom: ['Research'] },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+
+    // Second step's system prompt should contain the first step's result
+    expect(systemPrompts[1]).toContain('Result from Research');
+    expect(systemPrompts[1]).toContain('Research');
+    // First step should NOT contain context from others
+    expect(systemPrompts[0]).not.toContain('Previous Results');
+  });
+
+  it('worker status tracking — getWorkers returns correct statuses', async () => {
+    const template: WorkflowTemplate = {
+      name: 'status-workflow',
+      description: 'Status tracking workflow',
+      steps: [
+        { name: 'Worker 1', role: 'researcher', task: 'Do research' },
+        { name: 'Worker 2', role: 'writer', task: 'Write docs' },
+      ],
+      aggregation: 'concatenate',
+    };
+
+    await orchestrator.runWorkflow(template);
+    const workers = orchestrator.getWorkers();
+    expect(workers).toHaveLength(2);
+    expect(workers[0].status).toBe('done');
+    expect(workers[1].status).toBe('done');
+    expect(workers[0].name).toBe('Worker 1');
+    expect(workers[1].name).toBe('Worker 2');
+    expect(workers[0].result).toContain('Result for');
+    expect(workers[0].startedAt).toBeDefined();
+    expect(workers[0].completedAt).toBeDefined();
+    expect(workers[0].completedAt!).toBeGreaterThanOrEqual(workers[0].startedAt!);
+  });
+
+  it('worker:status events emitted', async () => {
+    const events: Array<{ workerId: string; status: string }> = [];
+    orchestrator.on('worker:status', (data) => {
+      events.push({ workerId: data.workerId, status: data.status });
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'event-workflow',
+      description: 'Event workflow',
+      steps: [
+        { name: 'Solo', role: 'researcher', task: 'Do work' },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+
+    // Should have 'pending', 'running', and 'done' events in order
+    const statuses = events.map(e => e.status);
+    expect(statuses).toContain('pending');
+    expect(statuses).toContain('running');
+    expect(statuses).toContain('done');
+    expect(statuses.indexOf('pending')).toBeLessThan(statuses.indexOf('running'));
+    expect(statuses.indexOf('running')).toBeLessThan(statuses.indexOf('done'));
+  });
+
+  it('workers start as pending before execution begins', async () => {
+    const statesObserved: Array<{ name: string; status: string }> = [];
+    orchestrator.on('worker:status', (data) => {
+      statesObserved.push({ name: data.workerState.name, status: data.status });
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'pending-workflow',
+      description: 'Pending state workflow',
+      steps: [
+        { name: 'First', role: 'researcher', task: 'Task 1' },
+        { name: 'Second', role: 'writer', task: 'Task 2', dependsOn: ['First'] },
+      ],
+      aggregation: 'concatenate',
+    };
+
+    await orchestrator.runWorkflow(template);
+
+    // Both workers should emit 'pending' before any 'running'
+    const firstPending = statesObserved.findIndex(s => s.name === 'First' && s.status === 'pending');
+    const secondPending = statesObserved.findIndex(s => s.name === 'Second' && s.status === 'pending');
+    const firstRunning = statesObserved.findIndex(s => s.name === 'First' && s.status === 'running');
+
+    expect(firstPending).toBeGreaterThanOrEqual(0);
+    expect(secondPending).toBeGreaterThanOrEqual(0);
+    // Both pending events should fire before any running
+    expect(firstPending).toBeLessThan(firstRunning);
+    expect(secondPending).toBeLessThan(firstRunning);
+  });
+
+  it('failed worker sets error status', async () => {
+    runner.mockImplementation(async () => {
+      throw new Error('LLM connection failed');
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'fail-workflow',
+      description: 'Failure workflow',
+      steps: [
+        { name: 'Failing Worker', role: 'researcher', task: 'This will fail' },
+      ],
+      aggregation: 'last',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    const workers = orchestrator.getWorkers();
+    expect(workers).toHaveLength(1);
+    expect(workers[0].status).toBe('failed');
+    expect(workers[0].error).toBe('LLM connection failed');
+    expect(workers[0].completedAt).toBeDefined();
+  });
+
+  it('aggregation concatenate joins all results', async () => {
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      return {
+        content: `Output-${match?.[1]}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: [],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'concat-workflow',
+      description: 'Concat workflow',
+      steps: [
+        { name: 'A', role: 'researcher', task: 'Task A' },
+        { name: 'B', role: 'writer', task: 'Task B' },
+      ],
+      aggregation: 'concatenate',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    expect(result.aggregated).toContain('Output-A');
+    expect(result.aggregated).toContain('Output-B');
+    expect(result.aggregated).toContain('A'); // header
+    expect(result.aggregated).toContain('B'); // header
+  });
+
+  it('aggregation last returns only last step result', async () => {
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      return {
+        content: `Output-${match?.[1]}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: [],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'last-workflow',
+      description: 'Last workflow',
+      steps: [
+        { name: 'First', role: 'researcher', task: 'Task 1' },
+        { name: 'Last', role: 'writer', task: 'Task 2' },
+      ],
+      aggregation: 'last',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    expect(result.aggregated).toBe('Output-Last');
+    expect(result.aggregated).not.toContain('Output-First');
+  });
+
+  it('aggregation synthesize spawns extra synthesizer worker', async () => {
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const match = config.systemPrompt.match(/Sub-Agent: (.+)/);
+      const name = match?.[1] ?? 'unknown';
+      return {
+        content: name.startsWith('Synthesizer') ? 'Final synthesis' : `Output-${name}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: [],
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'synth-workflow',
+      description: 'Synthesize workflow',
+      steps: [
+        { name: 'Research', role: 'researcher', task: 'Research' },
+        { name: 'Analysis', role: 'analyst', task: 'Analyze' },
+      ],
+      aggregation: 'synthesize',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    // Should have 3 calls — 2 steps + 1 synthesizer
+    expect(runner).toHaveBeenCalledTimes(3);
+    expect(result.aggregated).toBe('Final synthesis');
+  });
+
+  it('getActiveWorkers returns only running workers', async () => {
+    // Before any workflow, no active workers
+    expect(orchestrator.getActiveWorkers()).toEqual([]);
+
+    // After workflow completes, no active workers (all done)
+    const template: WorkflowTemplate = {
+      name: 'active-workflow',
+      description: 'Active workflow',
+      steps: [
+        { name: 'Worker', role: 'researcher', task: 'Work' },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+    expect(orchestrator.getActiveWorkers()).toEqual([]);
+    expect(orchestrator.getWorkers()).toHaveLength(1);
+  });
+
+  it('empty workflow returns empty results', async () => {
+    const template: WorkflowTemplate = {
+      name: 'empty-workflow',
+      description: 'Empty workflow',
+      steps: [],
+      aggregation: 'concatenate',
+    };
+
+    const result = await orchestrator.runWorkflow(template);
+    expect(result.results.size).toBe(0);
+    expect(result.aggregated).toBe('');
+    expect(runner).not.toHaveBeenCalled();
+  });
+
+  it('role tool presets are accessible', () => {
+    expect(SubagentOrchestrator.ROLE_TOOL_PRESETS.researcher).toContain('web_search');
+    expect(SubagentOrchestrator.ROLE_TOOL_PRESETS.coder).toContain('bash');
+    expect(SubagentOrchestrator.ROLE_TOOL_PRESETS.synthesizer).toContain('save_memory');
+    expect(SubagentOrchestrator.ROLE_TOOL_PRESETS.summarizer).toContain('read_file');
+  });
+
+  it('worker tools are filtered from available tools by role preset', async () => {
+    runner.mockImplementation(async (config: AgentLoopConfig) => {
+      const toolNames = config.tools.map(t => t.name);
+      return {
+        content: `Tools: ${toolNames.join(', ')}`,
+        usage: { inputTokens: 100, outputTokens: 50 },
+        toolsUsed: toolNames,
+      };
+    });
+
+    const template: WorkflowTemplate = {
+      name: 'tool-filter-workflow',
+      description: 'Tool filter workflow',
+      steps: [
+        { name: 'Coder', role: 'coder', task: 'Code something' },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+    const call = runner.mock.calls[0][0];
+    const toolNames = call.tools.map((t: ToolDefinition) => t.name);
+    expect(toolNames).toContain('bash');
+    expect(toolNames).toContain('read_file');
+    expect(toolNames).toContain('write_file');
+    // researcher-only tools should not be present
+    expect(toolNames).not.toContain('web_search');
+  });
+
+  it('step with explicit tools overrides role preset', async () => {
+    runner.mockImplementation(async (config: AgentLoopConfig) => ({
+      content: 'done',
+      usage: { inputTokens: 100, outputTokens: 50 },
+      toolsUsed: [],
+    }));
+
+    const template: WorkflowTemplate = {
+      name: 'custom-tools-workflow',
+      description: 'Custom tools workflow',
+      steps: [
+        { name: 'Custom', role: 'researcher', task: 'Custom task', tools: ['bash', 'read_file'] },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+    const call = runner.mock.calls[0][0];
+    const toolNames = call.tools.map((t: ToolDefinition) => t.name);
+    expect(toolNames).toContain('bash');
+    expect(toolNames).toContain('read_file');
+    expect(toolNames).not.toContain('web_search'); // researcher default, but overridden
+  });
+
+  it('tracks usage tokens per worker', async () => {
+    runner.mockImplementation(async () => ({
+      content: 'done',
+      usage: { inputTokens: 200, outputTokens: 100 },
+      toolsUsed: ['web_search', 'read_file'],
+    }));
+
+    const template: WorkflowTemplate = {
+      name: 'usage-workflow',
+      description: 'Usage workflow',
+      steps: [
+        { name: 'W1', role: 'researcher', task: 'Work' },
+      ],
+      aggregation: 'last',
+    };
+
+    await orchestrator.runWorkflow(template);
+    const workers = orchestrator.getWorkers();
+    expect(workers[0].usage).toEqual({ inputTokens: 200, outputTokens: 100 });
+    expect(workers[0].toolsUsed).toEqual(['web_search', 'read_file']);
+  });
+});