benchmarks/bench-agents/benchmark_multi_agent.py

#!/usr/bin/env python3
"""
Multi-Agent Orchestration Benchmark

Simulates a real-world scenario:
- 20 parallel sub-agents executing tasks
- Each agent makes tool calls through MCP
- Measures coordination overhead

Compares:
1. Traditional: Each agent has own MCP connections (JSON-RPC)
2. ZAP: Shared ZAP router with zero-copy message passing
"""

import json
import time
import struct
import asyncio
import random
from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
from concurrent.futures import ThreadPoolExecutor
import threading

NUM_AGENTS = 20
TOOL_CALLS_PER_AGENT = 50
ITERATIONS = 5

# Tool definitions (simulating real MCP tools)
TOOLS = [
    "read_file", "write_file", "list_directory",
    "execute_command", "search_code", "git_status",
    "git_diff", "git_commit", "create_file", "delete_file",
    "http_get", "http_post", "database_query", "cache_get",
    "cache_set", "send_notification", "parse_json", "format_code",
    "run_tests", "deploy_service"
]

@dataclass
class AgentTask:
    agent_id: int
    tool_name: str
    arguments: Dict[str, Any]
    timestamp: float = field(default_factory=time.time)

@dataclass
class AgentResult:
    agent_id: int
    task_count: int
    total_time_ms: float
    avg_latency_us: float
    memory_bytes: int

class JSONRPCSimulator:
    """Simulates traditional JSON-RPC MCP communication."""

    def __init__(self):
        self.message_counter = 0
        self.lock = threading.Lock()

    def encode_request(self, tool_name: str, arguments: Dict[str, Any]) -> bytes:
        with self.lock:
            self.message_counter += 1
            msg_id = self.message_counter

        message = {
            "jsonrpc": "2.0",
            "id": msg_id,
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        return json.dumps(message).encode()

    def decode_request(self, data: bytes) -> Dict[str, Any]:
        return json.loads(data)

    def encode_response(self, msg_id: int, result: Any) -> bytes:
        response = {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": result
        }
        return json.dumps(response).encode()

    def decode_response(self, data: bytes) -> Any:
        msg = json.loads(data)
        return msg.get("result")

class ZAPSimulator:
    """Simulates ZAP zero-copy message passing."""

    # Tool name to ID mapping (compiled into both sides)
    TOOL_IDS = {name: i for i, name in enumerate(TOOLS)}

    def __init__(self):
        self.message_counter = 0
        self.lock = threading.Lock()
        # Pre-allocated arena buffer
        self.arena = bytearray(64 * 1024)  # 64KB arena

    def encode_request(self, tool_name: str, arguments: Dict[str, Any]) -> memoryview:
        with self.lock:
            self.message_counter += 1
            msg_id = self.message_counter

        tool_id = self.TOOL_IDS.get(tool_name, 0)

        # Serialize arguments compactly
        # For simplicity, using a minimal binary format
        args_bytes = self._encode_args(arguments)

        # Header: id(8) + tool_id(4) + flags(4) + args_len(4) = 20 bytes
        offset = 0
        struct.pack_into('<QIIi', self.arena, offset, msg_id, tool_id, 0, len(args_bytes))
        offset += 20

        # Copy args
        self.arena[offset:offset+len(args_bytes)] = args_bytes
        offset += len(args_bytes)

        return memoryview(self.arena)[:offset]

    def decode_request(self, data: memoryview) -> tuple[int, int, bytes]:
        msg_id, tool_id, flags, args_len = struct.unpack_from('<QIIi', data, 0)
        args = data[20:20+args_len]
        return msg_id, tool_id, args

    def encode_response(self, msg_id: int, result: bytes) -> memoryview:
        offset = 0
        struct.pack_into('<QIi', self.arena, offset, msg_id, 0, len(result))
        offset += 16
        self.arena[offset:offset+len(result)] = result
        offset += len(result)
        return memoryview(self.arena)[:offset]

    def decode_response(self, data: memoryview) -> tuple[int, bytes]:
        msg_id, status, result_len = struct.unpack_from('<QIi', data, 0)
        result = data[16:16+result_len]
        return msg_id, result

    def _encode_args(self, arguments: Dict[str, Any]) -> bytes:
        """Minimal binary encoding for arguments."""
        # Simple format: count(4) + [key_len(2) + key + val_len(2) + val]*
        parts = []
        parts.append(struct.pack('<I', len(arguments)))
        for key, value in arguments.items():
            key_bytes = key.encode()
            val_bytes = str(value).encode()
            parts.append(struct.pack('<H', len(key_bytes)))
            parts.append(key_bytes)
            parts.append(struct.pack('<H', len(val_bytes)))
            parts.append(val_bytes)
        return b''.join(parts)

def generate_tool_call() -> tuple[str, Dict[str, Any]]:
    """Generate a random tool call."""
    tool = random.choice(TOOLS)
    args = {
        "path": f"/project/src/file_{random.randint(1, 100)}.py",
        "content": f"data_{random.randint(1, 1000)}",
        "options": {"recursive": True, "verbose": False}
    }
    return tool, args

def run_agent_json(agent_id: int, simulator: JSONRPCSimulator, num_calls: int) -> AgentResult:
    """Run an agent using JSON-RPC."""
    start = time.perf_counter_ns()
    latencies = []

    for _ in range(num_calls):
        tool, args = generate_tool_call()

        # Encode request
        call_start = time.perf_counter_ns()
        request = simulator.encode_request(tool, args)

        # Simulate decode on server side
        decoded = simulator.decode_request(request)

        # Simulate response
        response = simulator.encode_response(decoded["id"], {"success": True, "data": "result"})
        result = simulator.decode_response(response)

        latencies.append(time.perf_counter_ns() - call_start)

    total_time = (time.perf_counter_ns() - start) / 1_000_000  # ms
    avg_latency = sum(latencies) / len(latencies) / 1000  # µs

    # Estimate memory (JSON strings allocated per call)
    memory = num_calls * 500  # ~500 bytes per roundtrip

    return AgentResult(
        agent_id=agent_id,
        task_count=num_calls,
        total_time_ms=total_time,
        avg_latency_us=avg_latency,
        memory_bytes=memory
    )

def run_agent_zap(agent_id: int, simulator: ZAPSimulator, num_calls: int) -> AgentResult:
    """Run an agent using ZAP."""
    start = time.perf_counter_ns()
    latencies = []

    for _ in range(num_calls):
        tool, args = generate_tool_call()

        # Encode request (zero-copy to arena)
        call_start = time.perf_counter_ns()
        request = simulator.encode_request(tool, args)

        # Decode (zero-copy, just read offsets)
        msg_id, tool_id, args_data = simulator.decode_request(request)

        # Simulate response (zero-copy)
        response = simulator.encode_response(msg_id, b'{"success":true}')
        result_id, result_data = simulator.decode_response(response)

        latencies.append(time.perf_counter_ns() - call_start)

    total_time = (time.perf_counter_ns() - start) / 1_000_000  # ms
    avg_latency = sum(latencies) / len(latencies) / 1000  # µs

    # Memory: just the arena, shared across calls
    memory = 64 * 1024  # 64KB arena, reused

    return AgentResult(
        agent_id=agent_id,
        task_count=num_calls,
        total_time_ms=total_time,
        avg_latency_us=avg_latency,
        memory_bytes=memory
    )

def run_multi_agent_benchmark():
    """Run the full multi-agent benchmark."""
    results = {
        "metadata": {
            "num_agents": NUM_AGENTS,
            "tool_calls_per_agent": TOOL_CALLS_PER_AGENT,
            "total_tool_calls": NUM_AGENTS * TOOL_CALLS_PER_AGENT,
            "iterations": ITERATIONS
        },
        "json_rpc": [],
        "zap": [],
        "summary": {}
    }

    print(f"\n{'='*60}")
    print(f"MULTI-AGENT ORCHESTRATION BENCHMARK")
    print(f"{'='*60}")
    print(f"Agents: {NUM_AGENTS}")
    print(f"Tool calls per agent: {TOOL_CALLS_PER_AGENT}")
    print(f"Total tool calls: {NUM_AGENTS * TOOL_CALLS_PER_AGENT}")
    print(f"Iterations: {ITERATIONS}")

    # Run JSON-RPC benchmark
    print(f"\n📊 Running JSON-RPC benchmark...")
    json_times = []
    json_latencies = []
    json_memory = 0

    for iteration in range(ITERATIONS):
        json_sim = JSONRPCSimulator()
        start = time.perf_counter_ns()

        with ThreadPoolExecutor(max_workers=NUM_AGENTS) as executor:
            futures = [
                executor.submit(run_agent_json, i, json_sim, TOOL_CALLS_PER_AGENT)
                for i in range(NUM_AGENTS)
            ]
            agent_results = [f.result() for f in futures]

        total_time = (time.perf_counter_ns() - start) / 1_000_000
        json_times.append(total_time)
        json_latencies.extend([r.avg_latency_us for r in agent_results])
        json_memory = sum(r.memory_bytes for r in agent_results)

        print(f"   Iteration {iteration+1}: {total_time:.2f}ms")

    # Run ZAP benchmark
    print(f"\n📊 Running ZAP benchmark...")
    zap_times = []
    zap_latencies = []
    zap_memory = 0

    for iteration in range(ITERATIONS):
        zap_sim = ZAPSimulator()
        start = time.perf_counter_ns()

        with ThreadPoolExecutor(max_workers=NUM_AGENTS) as executor:
            futures = [
                executor.submit(run_agent_zap, i, zap_sim, TOOL_CALLS_PER_AGENT)
                for i in range(NUM_AGENTS)
            ]
            agent_results = [f.result() for f in futures]

        total_time = (time.perf_counter_ns() - start) / 1_000_000
        zap_times.append(total_time)
        zap_latencies.extend([r.avg_latency_us for r in agent_results])
        zap_memory = 64 * 1024  # Shared arena

        print(f"   Iteration {iteration+1}: {total_time:.2f}ms")

    # Calculate statistics
    json_avg_time = sum(json_times) / len(json_times)
    zap_avg_time = sum(zap_times) / len(zap_times)
    json_avg_latency = sum(json_latencies) / len(json_latencies)
    zap_avg_latency = sum(zap_latencies) / len(zap_latencies)

    speedup = json_avg_time / zap_avg_time if zap_avg_time > 0 else float('inf')
    latency_speedup = json_avg_latency / zap_avg_latency if zap_avg_latency > 0 else float('inf')
    memory_ratio = json_memory / zap_memory if zap_memory > 0 else float('inf')

    results["json_rpc"] = {
        "avg_total_time_ms": json_avg_time,
        "avg_latency_us": json_avg_latency,
        "total_memory_bytes": json_memory,
        "total_memory_mb": json_memory / (1024 * 1024)
    }

    results["zap"] = {
        "avg_total_time_ms": zap_avg_time,
        "avg_latency_us": zap_avg_latency,
        "total_memory_bytes": zap_memory,
        "total_memory_mb": zap_memory / (1024 * 1024)
    }

    results["summary"] = {
        "time_speedup": f"{speedup:.1f}x",
        "latency_speedup": f"{latency_speedup:.1f}x",
        "memory_ratio": f"{memory_ratio:.0f}x"
    }

    # Print summary
    print(f"\n{'='*60}")
    print("RESULTS SUMMARY")
    print(f"{'='*60}")

    print(f"\n⏱️  TOTAL ORCHESTRATION TIME ({NUM_AGENTS} agents, {NUM_AGENTS * TOOL_CALLS_PER_AGENT} tool calls)")
    print(f"   JSON-RPC: {json_avg_time:.2f}ms")
    print(f"   ZAP:      {zap_avg_time:.2f}ms")
    print(f"   → ZAP is {speedup:.1f}x faster")

    print(f"\n⚡ PER-CALL LATENCY")
    print(f"   JSON-RPC: {json_avg_latency:.2f}µs")
    print(f"   ZAP:      {zap_avg_latency:.2f}µs")
    print(f"   → ZAP is {latency_speedup:.1f}x faster")

    print(f"\n💾 MEMORY USAGE")
    print(f"   JSON-RPC: {json_memory/(1024*1024):.2f}MB (allocations per call)")
    print(f"   ZAP:      {zap_memory/1024:.2f}KB (shared arena)")
    print(f"   → ZAP uses {memory_ratio:.0f}x less memory")

    print(f"\n{'='*60}")

    # Output JSON
    print(json.dumps(results, indent=2))

    return results

if __name__ == "__main__":
    run_multi_agent_benchmark()