@jackchen_me/open-multi-agent 0.2.0 → 1.0.0

package/tests/trace.test.ts

@@ -0,0 +1,453 @@
+import { describe, it, expect, vi } from 'vitest'
+import { z } from 'zod'
+import { Agent } from '../src/agent/agent.js'
+import { AgentRunner, type RunOptions } from '../src/agent/runner.js'
+import { ToolRegistry, defineTool } from '../src/tool/framework.js'
+import { ToolExecutor } from '../src/tool/executor.js'
+import { executeWithRetry } from '../src/orchestrator/orchestrator.js'
+import { emitTrace, generateRunId } from '../src/utils/trace.js'
+import { createTask } from '../src/task/task.js'
+import type {
+  AgentConfig,
+  AgentRunResult,
+  LLMAdapter,
+  LLMResponse,
+  TraceEvent,
+} from '../src/types.js'
+
+// ---------------------------------------------------------------------------
+// Mock adapters
+// ---------------------------------------------------------------------------
+
+function mockAdapter(responses: LLMResponse[]): LLMAdapter {
+  let callIndex = 0
+  return {
+    name: 'mock',
+    async chat() {
+      return responses[callIndex++]!
+    },
+    async *stream() {
+      /* unused */
+    },
+  }
+}
+
+function textResponse(text: string): LLMResponse {
+  return {
+    id: `resp-${Math.random().toString(36).slice(2)}`,
+    content: [{ type: 'text' as const, text }],
+    model: 'mock-model',
+    stop_reason: 'end_turn',
+    usage: { input_tokens: 10, output_tokens: 20 },
+  }
+}
+
+function toolUseResponse(toolName: string, input: Record<string, unknown>): LLMResponse {
+  return {
+    id: `resp-${Math.random().toString(36).slice(2)}`,
+    content: [
+      {
+        type: 'tool_use' as const,
+        id: `tu-${Math.random().toString(36).slice(2)}`,
+        name: toolName,
+        input,
+      },
+    ],
+    model: 'mock-model',
+    stop_reason: 'tool_use',
+    usage: { input_tokens: 15, output_tokens: 25 },
+  }
+}
+
+function buildMockAgent(
+  config: AgentConfig,
+  responses: LLMResponse[],
+  registry?: ToolRegistry,
+  executor?: ToolExecutor,
+): Agent {
+  const reg = registry ?? new ToolRegistry()
+  const exec = executor ?? new ToolExecutor(reg)
+  const adapter = mockAdapter(responses)
+  const agent = new Agent(config, reg, exec)
+
+  const runner = new AgentRunner(adapter, reg, exec, {
+    model: config.model,
+    systemPrompt: config.systemPrompt,
+    maxTurns: config.maxTurns,
+    maxTokens: config.maxTokens,
+    temperature: config.temperature,
+    agentName: config.name,
+  })
+  ;(agent as any).runner = runner
+
+  return agent
+}
+
+// ---------------------------------------------------------------------------
+// emitTrace helper
+// ---------------------------------------------------------------------------
+
+describe('emitTrace', () => {
+  it('does nothing when fn is undefined', () => {
+    // Should not throw
+    emitTrace(undefined, {
+      type: 'agent',
+      runId: 'r1',
+      agent: 'a',
+      turns: 1,
+      tokens: { input_tokens: 0, output_tokens: 0 },
+      toolCalls: 0,
+      startMs: 0,
+      endMs: 0,
+      durationMs: 0,
+    })
+  })
+
+  it('calls fn with the event', () => {
+    const fn = vi.fn()
+    const event: TraceEvent = {
+      type: 'agent',
+      runId: 'r1',
+      agent: 'a',
+      turns: 1,
+      tokens: { input_tokens: 0, output_tokens: 0 },
+      toolCalls: 0,
+      startMs: 0,
+      endMs: 0,
+      durationMs: 0,
+    }
+    emitTrace(fn, event)
+    expect(fn).toHaveBeenCalledWith(event)
+  })
+
+  it('swallows errors thrown by callback', () => {
+    const fn = () => { throw new Error('boom') }
+    expect(() =>
+      emitTrace(fn, {
+        type: 'agent',
+        runId: 'r1',
+        agent: 'a',
+        turns: 1,
+        tokens: { input_tokens: 0, output_tokens: 0 },
+        toolCalls: 0,
+        startMs: 0,
+        endMs: 0,
+        durationMs: 0,
+      }),
+    ).not.toThrow()
+  })
+
+  it('swallows rejected promises from async callbacks', async () => {
+    // An async onTrace that rejects should not produce unhandled rejection
+    const fn = async () => { throw new Error('async boom') }
+    emitTrace(fn as unknown as (event: TraceEvent) => void, {
+      type: 'agent',
+      runId: 'r1',
+      agent: 'a',
+      turns: 1,
+      tokens: { input_tokens: 0, output_tokens: 0 },
+      toolCalls: 0,
+      startMs: 0,
+      endMs: 0,
+      durationMs: 0,
+    })
+    // If the rejection is not caught, vitest will fail with unhandled rejection.
+    // Give the microtask queue a tick to surface any unhandled rejection.
+    await new Promise(resolve => setTimeout(resolve, 10))
+  })
+})
+
+describe('generateRunId', () => {
+  it('returns a UUID string', () => {
+    const id = generateRunId()
+    expect(id).toMatch(/^[0-9a-f-]{36}$/)
+  })
+
+  it('returns unique IDs', () => {
+    const ids = new Set(Array.from({ length: 100 }, generateRunId))
+    expect(ids.size).toBe(100)
+  })
+})
+
+// ---------------------------------------------------------------------------
+// AgentRunner trace events
+// ---------------------------------------------------------------------------
+
+describe('AgentRunner trace events', () => {
+  it('emits llm_call trace for each LLM turn', async () => {
+    const traces: TraceEvent[] = []
+    const registry = new ToolRegistry()
+    const executor = new ToolExecutor(registry)
+    const adapter = mockAdapter([textResponse('Hello!')])
+
+    const runner = new AgentRunner(adapter, registry, executor, {
+      model: 'test-model',
+      agentName: 'test-agent',
+    })
+
+    const runOptions: RunOptions = {
+      onTrace: (e) => traces.push(e),
+      runId: 'run-1',
+      traceAgent: 'test-agent',
+    }
+
+    await runner.run(
+      [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
+      runOptions,
+    )
+
+    const llmTraces = traces.filter(t => t.type === 'llm_call')
+    expect(llmTraces).toHaveLength(1)
+
+    const llm = llmTraces[0]!
+    expect(llm.type).toBe('llm_call')
+    expect(llm.runId).toBe('run-1')
+    expect(llm.agent).toBe('test-agent')
+    expect(llm.model).toBe('test-model')
+    expect(llm.turn).toBe(1)
+    expect(llm.tokens).toEqual({ input_tokens: 10, output_tokens: 20 })
+    expect(llm.durationMs).toBeGreaterThanOrEqual(0)
+    expect(llm.startMs).toBeLessThanOrEqual(llm.endMs)
+  })
+
+  it('emits tool_call trace with correct fields', async () => {
+    const traces: TraceEvent[] = []
+    const registry = new ToolRegistry()
+    registry.register(
+      defineTool({
+        name: 'echo',
+        description: 'echoes',
+        inputSchema: z.object({ msg: z.string() }),
+        execute: async ({ msg }) => ({ data: msg }),
+      }),
+    )
+    const executor = new ToolExecutor(registry)
+    const adapter = mockAdapter([
+      toolUseResponse('echo', { msg: 'hello' }),
+      textResponse('Done'),
+    ])
+
+    const runner = new AgentRunner(adapter, registry, executor, {
+      model: 'test-model',
+      agentName: 'tooler',
+    })
+
+    await runner.run(
+      [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
+      { onTrace: (e) => traces.push(e), runId: 'run-2', traceAgent: 'tooler' },
+    )
+
+    const toolTraces = traces.filter(t => t.type === 'tool_call')
+    expect(toolTraces).toHaveLength(1)
+
+    const tool = toolTraces[0]!
+    expect(tool.type).toBe('tool_call')
+    expect(tool.runId).toBe('run-2')
+    expect(tool.agent).toBe('tooler')
+    expect(tool.tool).toBe('echo')
+    expect(tool.isError).toBe(false)
+    expect(tool.durationMs).toBeGreaterThanOrEqual(0)
+  })
+
+  it('tool_call trace has isError: true on tool failure', async () => {
+    const traces: TraceEvent[] = []
+    const registry = new ToolRegistry()
+    registry.register(
+      defineTool({
+        name: 'boom',
+        description: 'fails',
+        inputSchema: z.object({}),
+        execute: async () => { throw new Error('fail') },
+      }),
+    )
+    const executor = new ToolExecutor(registry)
+    const adapter = mockAdapter([
+      toolUseResponse('boom', {}),
+      textResponse('Handled'),
+    ])
+
+    const runner = new AgentRunner(adapter, registry, executor, {
+      model: 'test-model',
+      agentName: 'err-agent',
+    })
+
+    await runner.run(
+      [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
+      { onTrace: (e) => traces.push(e), runId: 'run-3', traceAgent: 'err-agent' },
+    )
+
+    const toolTraces = traces.filter(t => t.type === 'tool_call')
+    expect(toolTraces).toHaveLength(1)
+    expect(toolTraces[0]!.isError).toBe(true)
+  })
+
+  it('does not call Date.now for LLM timing when onTrace is absent', async () => {
+    // This test just verifies no errors occur when onTrace is not provided
+    const registry = new ToolRegistry()
+    const executor = new ToolExecutor(registry)
+    const adapter = mockAdapter([textResponse('hi')])
+
+    const runner = new AgentRunner(adapter, registry, executor, {
+      model: 'test-model',
+    })
+
+    const result = await runner.run(
+      [{ role: 'user', content: [{ type: 'text', text: 'test' }] }],
+      {},
+    )
+
+    expect(result.output).toBe('hi')
+  })
+})
+
+// ---------------------------------------------------------------------------
+// Agent-level trace events
+// ---------------------------------------------------------------------------
+
+describe('Agent trace events', () => {
+  it('emits agent trace with turns, tokens, and toolCalls', async () => {
+    const traces: TraceEvent[] = []
+    const config: AgentConfig = {
+      name: 'my-agent',
+      model: 'mock-model',
+      systemPrompt: 'You are a test.',
+    }
+
+    const agent = buildMockAgent(config, [textResponse('Hello world')])
+
+    const runOptions: Partial<RunOptions> = {
+      onTrace: (e) => traces.push(e),
+      runId: 'run-agent-1',
+      traceAgent: 'my-agent',
+    }
+
+    const result = await agent.run('Say hello', runOptions)
+    expect(result.success).toBe(true)
+
+    const agentTraces = traces.filter(t => t.type === 'agent')
+    expect(agentTraces).toHaveLength(1)
+
+    const at = agentTraces[0]!
+    expect(at.type).toBe('agent')
+    expect(at.runId).toBe('run-agent-1')
+    expect(at.agent).toBe('my-agent')
+    expect(at.turns).toBe(1) // one assistant message
+    expect(at.tokens).toEqual({ input_tokens: 10, output_tokens: 20 })
+    expect(at.toolCalls).toBe(0)
+    expect(at.durationMs).toBeGreaterThanOrEqual(0)
+  })
+
+  it('all traces share the same runId', async () => {
+    const traces: TraceEvent[] = []
+    const registry = new ToolRegistry()
+    registry.register(
+      defineTool({
+        name: 'greet',
+        description: 'greets',
+        inputSchema: z.object({ name: z.string() }),
+        execute: async ({ name }) => ({ data: `Hi ${name}` }),
+      }),
+    )
+    const executor = new ToolExecutor(registry)
+    const config: AgentConfig = {
+      name: 'multi-trace-agent',
+      model: 'mock-model',
+      tools: ['greet'],
+    }
+
+    const agent = buildMockAgent(
+      config,
+      [
+        toolUseResponse('greet', { name: 'world' }),
+        textResponse('Done'),
+      ],
+      registry,
+      executor,
+    )
+
+    const runId = 'shared-run-id'
+    await agent.run('test', {
+      onTrace: (e) => traces.push(e),
+      runId,
+      traceAgent: 'multi-trace-agent',
+    })
+
+    // Should have: 2 llm_call, 1 tool_call, 1 agent
+    expect(traces.length).toBeGreaterThanOrEqual(4)
+
+    for (const trace of traces) {
+      expect(trace.runId).toBe(runId)
+    }
+  })
+
+  it('onTrace error does not break agent execution', async () => {
+    const config: AgentConfig = {
+      name: 'resilient-agent',
+      model: 'mock-model',
+    }
+
+    const agent = buildMockAgent(config, [textResponse('OK')])
+
+    const result = await agent.run('test', {
+      onTrace: () => { throw new Error('callback exploded') },
+      runId: 'run-err',
+      traceAgent: 'resilient-agent',
+    })
+
+    // The run should still succeed despite the broken callback
+    expect(result.success).toBe(true)
+    expect(result.output).toBe('OK')
+  })
+
+  it('per-turn token usage in llm_call traces', async () => {
+    const traces: TraceEvent[] = []
+    const registry = new ToolRegistry()
+    registry.register(
+      defineTool({
+        name: 'noop',
+        description: 'noop',
+        inputSchema: z.object({}),
+        execute: async () => ({ data: 'ok' }),
+      }),
+    )
+    const executor = new ToolExecutor(registry)
+
+    // Two LLM calls: first triggers a tool, second is the final response
+    const resp1: LLMResponse = {
+      id: 'r1',
+      content: [{ type: 'tool_use', id: 'tu1', name: 'noop', input: {} }],
+      model: 'mock-model',
+      stop_reason: 'tool_use',
+      usage: { input_tokens: 100, output_tokens: 50 },
+    }
+    const resp2: LLMResponse = {
+      id: 'r2',
+      content: [{ type: 'text', text: 'Final answer' }],
+      model: 'mock-model',
+      stop_reason: 'end_turn',
+      usage: { input_tokens: 200, output_tokens: 100 },
+    }
+
+    const adapter = mockAdapter([resp1, resp2])
+    const runner = new AgentRunner(adapter, registry, executor, {
+      model: 'mock-model',
+      agentName: 'token-agent',
+    })
+
+    await runner.run(
+      [{ role: 'user', content: [{ type: 'text', text: 'go' }] }],
+      { onTrace: (e) => traces.push(e), runId: 'run-tok', traceAgent: 'token-agent' },
+    )
+
+    const llmTraces = traces.filter(t => t.type === 'llm_call')
+    expect(llmTraces).toHaveLength(2)
+
+    // Each trace carries its own turn's token usage, not the aggregate
+    expect(llmTraces[0]!.tokens).toEqual({ input_tokens: 100, output_tokens: 50 })
+    expect(llmTraces[1]!.tokens).toEqual({ input_tokens: 200, output_tokens: 100 })
+
+    // Turn numbers should be sequential
+    expect(llmTraces[0]!.turn).toBe(1)
+    expect(llmTraces[1]!.turn).toBe(2)
+  })
+})

package/vitest.config.ts

@@ -0,0 +1,9 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    coverage: {
+      include: ['src/**'],
+    },
+  },
+})

package/examples/09-gemma4-auto-orchestration.ts

@@ -1,162 +0,0 @@
-/**
- * Example 09 — Gemma 4 Auto-Orchestration (runTeam, 100% Local)
- *
- * Demonstrates the framework's key feature — automatic task decomposition —
- * powered entirely by a local Gemma 4 model. No cloud API needed.
- *
- * What happens:
- * 1. A Gemma 4 "coordinator" receives the goal + agent roster
- * 2. It outputs a structured JSON task array (title, description, assignee, dependsOn)
- * 3. The framework resolves dependencies, schedules tasks, and runs agents
- * 4. The coordinator synthesises all task results into a final answer
- *
- * This is the hardest test for a local model — it must produce valid JSON
- * for task decomposition AND do tool-calling for actual task execution.
- * Gemma 4 e2b (5.1B params) handles both reliably.
- *
- * Run:
- *   no_proxy=localhost npx tsx examples/09-gemma4-auto-orchestration.ts
- *
- * Prerequisites:
- *   1. Ollama >= 0.20.0 installed and running: https://ollama.com
- *   2. Pull the model: ollama pull gemma4:e2b
- *   3. No API keys needed!
- *
- * Note: The no_proxy=localhost prefix is needed if you have an HTTP proxy
- * configured, since the OpenAI SDK would otherwise route Ollama requests
- * through the proxy.
- */
-
-import { OpenMultiAgent } from '../src/index.js'
-import type { AgentConfig, OrchestratorEvent, Task } from '../src/types.js'
-
-// ---------------------------------------------------------------------------
-// Configuration
-// ---------------------------------------------------------------------------
-
-// See available tags at https://ollama.com/library/gemma4
-const OLLAMA_MODEL = 'gemma4:e2b'      // or 'gemma4:e4b', 'gemma4:26b'
-const OLLAMA_BASE_URL = 'http://localhost:11434/v1'
-
-// ---------------------------------------------------------------------------
-// Agents — the coordinator is created automatically by runTeam()
-// ---------------------------------------------------------------------------
-
-const researcher: AgentConfig = {
-  name: 'researcher',
-  model: OLLAMA_MODEL,
-  provider: 'openai',
-  baseURL: OLLAMA_BASE_URL,
-  apiKey: 'ollama',
-  systemPrompt: `You are a system researcher. Use bash to run non-destructive,
-read-only commands and report the results concisely.`,
-  tools: ['bash'],
-  maxTurns: 4,
-}
-
-const writer: AgentConfig = {
-  name: 'writer',
-  model: OLLAMA_MODEL,
-  provider: 'openai',
-  baseURL: OLLAMA_BASE_URL,
-  apiKey: 'ollama',
-  systemPrompt: `You are a technical writer. Use file_write to create clear,
-structured Markdown reports based on the information provided.`,
-  tools: ['file_write'],
-  maxTurns: 4,
-}
-
-// ---------------------------------------------------------------------------
-// Progress handler
-// ---------------------------------------------------------------------------
-
-function handleProgress(event: OrchestratorEvent): void {
-  const ts = new Date().toISOString().slice(11, 23)
-  switch (event.type) {
-    case 'task_start': {
-      const task = event.data as Task | undefined
-      console.log(`[${ts}] TASK START    "${task?.title ?? event.task}" → ${task?.assignee ?? '?'}`)
-      break
-    }
-    case 'task_complete':
-      console.log(`[${ts}] TASK DONE     "${event.task}"`)
-      break
-    case 'agent_start':
-      console.log(`[${ts}] AGENT START   ${event.agent}`)
-      break
-    case 'agent_complete':
-      console.log(`[${ts}] AGENT DONE    ${event.agent}`)
-      break
-    case 'error':
-      console.error(`[${ts}] ERROR         ${event.agent ?? ''}  task=${event.task ?? '?'}`)
-      break
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Orchestrator — defaultModel is used for the coordinator agent
-// ---------------------------------------------------------------------------
-
-const orchestrator = new OpenMultiAgent({
-  defaultModel: OLLAMA_MODEL,
-  defaultProvider: 'openai',
-  defaultBaseURL: OLLAMA_BASE_URL,
-  defaultApiKey: 'ollama',
-  maxConcurrency: 1, // local model serves one request at a time
-  onProgress: handleProgress,
-})
-
-const team = orchestrator.createTeam('gemma4-auto', {
-  name: 'gemma4-auto',
-  agents: [researcher, writer],
-  sharedMemory: true,
-})
-
-// ---------------------------------------------------------------------------
-// Give a goal — the framework handles the rest
-// ---------------------------------------------------------------------------
-
-const goal = `Check this machine's Node.js version, npm version, and OS info,
-then write a short Markdown summary report to /tmp/gemma4-auto/report.md`
-
-console.log('Gemma 4 Auto-Orchestration — Zero API Cost')
-console.log('='.repeat(60))
-console.log(`  model        → ${OLLAMA_MODEL} via Ollama (all agents + coordinator)`)
-console.log(`  researcher   → bash`)
-console.log(`  writer       → file_write`)
-console.log(`  coordinator  → auto-created by runTeam()`)
-console.log()
-console.log(`Goal: ${goal.replace(/\n/g, ' ').trim()}`)
-console.log('='.repeat(60))
-
-const start = Date.now()
-const result = await orchestrator.runTeam(team, goal)
-const totalTime = Date.now() - start
-
-// ---------------------------------------------------------------------------
-// Results
-// ---------------------------------------------------------------------------
-
-console.log('\n' + '='.repeat(60))
-console.log('Pipeline complete.\n')
-console.log(`Overall success: ${result.success}`)
-console.log(`Total time: ${(totalTime / 1000).toFixed(1)}s`)
-console.log(`Tokens — input: ${result.totalTokenUsage.input_tokens}, output: ${result.totalTokenUsage.output_tokens}`)
-
-console.log('\nPer-agent results:')
-for (const [name, r] of result.agentResults) {
-  const icon = r.success ? 'OK  ' : 'FAIL'
-  const tools = r.toolCalls.length > 0 ? r.toolCalls.map(c => c.toolName).join(', ') : '(none)'
-  console.log(`  [${icon}] ${name.padEnd(24)} tools: ${tools}`)
-}
-
-// Print the coordinator's final synthesis
-const coordResult = result.agentResults.get('coordinator')
-if (coordResult?.success) {
-  console.log('\nFinal synthesis (from local Gemma 4 coordinator):')
-  console.log('-'.repeat(60))
-  console.log(coordResult.output)
-  console.log('-'.repeat(60))
-}
-
-console.log('\nAll processing done locally. $0 API cost.')