How to Test Multi-Agent Systems (Handoffs and Orchestration)

"""Surface 1: the handoff contract.

An orchestrator hands work to a specialist agent as a message. That message is
an API even though nobody wrote an OpenAPI spec for it. This test pins the
shape of that message so a prompt edit upstream cannot silently change what
downstream agents receive.

The orchestrator here is a deterministic fake. No LLM calls, no API keys.

Run with:
    pytest tests/test_handoff_contract.py
"""

from __future__ import annotations

import pytest
from pydantic import BaseModel, ConfigDict, Field, ValidationError


# ---------------------------------------------------------------------------
# The contract
# ---------------------------------------------------------------------------


class HandoffConstraints(BaseModel):
    """Limits the orchestrator imposes on the specialist agent.

    These are the fields that get quietly dropped when someone rewrites the
    orchestrator prompt. Making them required is the whole point.
    """

    model_config = ConfigDict(extra="forbid")

    max_tool_calls: int = Field(ge=1, le=100)
    timeout_seconds: float = Field(gt=0, le=600)
    allowed_tools: list[str] = Field(min_length=1)


class HandoffContext(BaseModel):
    """Everything the specialist needs that it cannot rediscover on its own."""

    model_config = ConfigDict(extra="forbid")

    correlation_id: str = Field(min_length=1)
    user_goal: str = Field(min_length=1)
    prior_findings: list[str] = Field(default_factory=list)


class HandoffMessage(BaseModel):
    """The message an orchestrator sends a specialist agent.

    `extra="forbid"` is deliberate: an agent that invents a new top-level field
    is a contract break, not a harmless addition. You want that to fail loudly
    in CI rather than get ignored by the recipient at runtime.
    """

    model_config = ConfigDict(extra="forbid")

    task_id: str = Field(min_length=1)
    recipient: str = Field(min_length=1)
    constraints: HandoffConstraints
    context: HandoffContext


# ---------------------------------------------------------------------------
# A deterministic fake orchestrator
# ---------------------------------------------------------------------------

ALLOWED_TOOLS_BY_AGENT: dict[str, list[str]] = {
    "researcher": ["web_search", "read_url"],
    "coder": ["read_file", "write_file", "run_tests"],
    "reviewer": ["read_file", "run_tests"],
}


def build_handoff(task_id: str, recipient: str, user_goal: str) -> dict:
    """Stand-in for the payload a real orchestrator would emit.

    In your codebase this is the function that serializes an LLM's delegation
    decision into a message. Swap the body, keep the test.
    """
    return {
        "task_id": task_id,
        "recipient": recipient,
        "constraints": {
            "max_tool_calls": 8,
            "timeout_seconds": 45.0,
            "allowed_tools": ALLOWED_TOOLS_BY_AGENT[recipient],
        },
        "context": {
            "correlation_id": f"corr-{task_id}",
            "user_goal": user_goal,
            "prior_findings": [],
        },
    }


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_orchestrator_payload_validates_against_the_contract():
    payload = build_handoff("t-001", "coder", "Add a retry to the upload path")

    message = HandoffMessage.model_validate(payload)

    assert message.task_id == "t-001"
    assert message.recipient == "coder"
    assert message.constraints.max_tool_calls == 8
    assert message.context.correlation_id == "corr-t-001"
    assert "run_tests" in message.constraints.allowed_tools


@pytest.mark.parametrize("recipient", sorted(ALLOWED_TOOLS_BY_AGENT))
def test_every_recipient_gets_a_valid_handoff(recipient: str):
    """One broken recipient branch should not hide behind a passing default."""
    payload = build_handoff("t-002", recipient, "Investigate the failing job")

    message = HandoffMessage.model_validate(payload)

    assert message.recipient == recipient
    assert message.constraints.allowed_tools == ALLOWED_TOOLS_BY_AGENT[recipient]


def test_missing_required_constraint_is_rejected():
    """The regression this test exists for: a dropped timeout.

    Without the contract, the specialist agent inherits some default timeout
    and the bug shows up three weeks later as an unexplained cost spike.
    """
    payload = build_handoff("t-003", "researcher", "Summarize the incident")
    del payload["constraints"]["timeout_seconds"]

    with pytest.raises(ValidationError) as excinfo:
        HandoffMessage.model_validate(payload)

    errors = excinfo.value.errors()
    assert [e["loc"] for e in errors] == [("constraints", "timeout_seconds")]
    assert errors[0]["type"] == "missing"


def test_out_of_range_constraint_is_rejected():
    """A budget of zero tool calls is a bug, not a strict budget."""
    payload = build_handoff("t-004", "coder", "Fix the flaky test")
    payload["constraints"]["max_tool_calls"] = 0

    with pytest.raises(ValidationError) as excinfo:
        HandoffMessage.model_validate(payload)

    assert excinfo.value.errors()[0]["loc"] == ("constraints", "max_tool_calls")


def test_unexpected_field_is_rejected():
    """An agent inventing `priority` is a contract break, not a feature."""
    payload = build_handoff("t-005", "reviewer", "Check the migration")
    payload["priority"] = "urgent"

    with pytest.raises(ValidationError) as excinfo:
        HandoffMessage.model_validate(payload)

    assert excinfo.value.errors()[0]["type"] == "extra_forbidden"

"""Surface 2: the delegation decision.

The orchestrator picks which specialist handles a task. That decision is the
single highest-leverage thing to test, because a wrong route means every
downstream assertion is measuring the wrong agent.

Two kinds of assertion live here:

1. Deterministic routes. Same input, same agent, every time. Assert exactly.
2. Stochastic routes. Ambiguous input where a real model would waver. Assert on
   the majority outcome over N runs, not on a single pass/fail, otherwise the
   test is a coin flip you have wired into CI.

Run with:
    pytest tests/test_delegation_routing.py
"""

from __future__ import annotations

import random
from collections import Counter

import pytest

# The allowed set. A route outside this set is a hallucinated agent name, which
# is a different and worse bug than a merely suboptimal route.
CANDIDATE_AGENTS = frozenset({"researcher", "coder", "reviewer", "summarizer"})

# Weighted tie-break for genuinely ambiguous work. A real orchestrator gets this
# spread from model sampling; here it is an explicit distribution so the test
# runs offline and the intent is visible.
AMBIGUOUS_SPLIT: dict[str, float] = {"coder": 0.8, "reviewer": 0.2}


def route(task: str, rng: random.Random) -> str:
    """Deterministic fake of an orchestrator's routing step.

    Replace the body with a call to your real router. The tests below do not
    care how the decision is made, only that it lands on a known agent.
    """
    text = task.lower()

    if any(word in text for word in ("implement", "refactor", "fix the bug")):
        return "coder"
    if any(word in text for word in ("research", "look up", "find out")):
        return "researcher"
    if any(word in text for word in ("summarize", "tl;dr", "recap")):
        return "summarizer"
    if any(word in text for word in ("review", "audit", "critique")):
        return "reviewer"

    population = list(AMBIGUOUS_SPLIT)
    weights = [AMBIGUOUS_SPLIT[name] for name in population]
    return rng.choices(population, weights=weights, k=1)[0]


def majority_route(task: str, trials: int = 101, seed: int = 1337) -> tuple[str, float]:
    """Run the same input N times and return (winning agent, its share).

    This is the assertion shape you want for any non-deterministic step. A
    single call tells you nothing; the distribution tells you whether the
    router actually has an opinion.
    """
    if trials < 1:
        raise ValueError("trials must be >= 1")

    rng = random.Random(seed)
    counts = Counter(route(task, rng) for _ in range(trials))
    winner, hits = counts.most_common(1)[0]

    unknown = set(counts) - CANDIDATE_AGENTS
    if unknown:
        raise AssertionError(f"router produced unknown agents: {sorted(unknown)}")

    return winner, hits / trials


# ---------------------------------------------------------------------------
# Deterministic routes
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    ("task", "expected_agent"),
    [
        ("Implement pagination on the orders endpoint", "coder"),
        ("Refactor the retry helper", "coder"),
        ("Fix the bug in the CSV importer", "coder"),
        ("Research how competitors price seats", "researcher"),
        ("Look up the Postgres lock semantics", "researcher"),
        ("Summarize yesterday's incident thread", "summarizer"),
        ("Review the auth middleware diff", "reviewer"),
        ("Audit the new IAM policy", "reviewer"),
    ],
)
def test_deterministic_routes_pick_the_expected_agent(task: str, expected_agent: str):
    selected = route(task, random.Random(0))

    assert selected == expected_agent
    assert selected in CANDIDATE_AGENTS


@pytest.mark.parametrize("seed", [0, 1, 2, 7, 99])
def test_deterministic_routes_ignore_the_rng(seed: int):
    """A route that is supposed to be deterministic must not drift with sampling."""
    assert route("Implement pagination", random.Random(seed)) == "coder"


# ---------------------------------------------------------------------------
# Stochastic route
# ---------------------------------------------------------------------------


def test_ambiguous_task_has_a_stable_majority():
    """"Make this better" is ambiguous. The router should still lean one way."""
    winner, share = majority_route("Make this better")

    assert winner == "coder"
    assert share >= 0.6, f"router has no clear preference (winner share {share:.2f})"


def test_ambiguous_task_never_leaves_the_candidate_set():
    """Even when the router wavers, it may only waver between real agents."""
    rng = random.Random(4242)
    seen = {route("Make this better", rng) for _ in range(500)}

    assert seen <= CANDIDATE_AGENTS
    assert len(seen) > 1, "expected this input to be genuinely ambiguous"


def test_majority_route_rejects_a_nonsense_trial_count():
    with pytest.raises(ValueError):
        majority_route("Make this better", trials=0)

"""Surface 3: the executed step sequence.

A plan is what the orchestrator said it would do. The step log is what it
actually did. Multi-agent bugs live in the gap between the two: a step silently
dropped, two steps executed out of order, one step run four times because a
retry had no ceiling.

Every step is logged as structured data (name plus timestamp), never as a free
text line, because you cannot assert on a log line without regexes and you can
assert on a tuple of names in one comparison.

The clock is fake and monotonic, so the log is byte-identical on every run and
in CI. Wall clock timestamps would make these assertions untestable.

Run with:
    pytest tests/test_orchestration_sequence.py
"""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field

import pytest

PLAN: tuple[str, ...] = (
    "parse_request",
    "retrieve_context",
    "draft_answer",
    "review_answer",
    "emit_result",
)


@dataclass(frozen=True)
class StepEvent:
    """One executed step. Structured, so tests compare values not substrings."""

    name: str
    ts: float
    attempt: int = 1


@dataclass
class FakeClock:
    """Monotonic, deterministic, no wall clock anywhere in the suite."""

    now: float = 0.0
    tick: float = 0.5

    def advance(self) -> float:
        self.now = round(self.now + self.tick, 6)
        return self.now


@dataclass
class FakeOrchestrator:
    """Runs a plan and records what it actually executed.

    The `drop`, `swap` and `retries` knobs are how we inject the three failure
    modes on purpose. Your real orchestrator does not need them; it produces
    the same failure modes for free.
    """

    plan: tuple[str, ...] = PLAN
    drop: frozenset[str] = frozenset()
    swap: tuple[str, str] | None = None
    retries: dict[str, int] = field(default_factory=dict)
    clock: FakeClock = field(default_factory=FakeClock)
    log: list[StepEvent] = field(default_factory=list)

    def _ordered_plan(self) -> list[str]:
        steps = list(self.plan)
        if self.swap is not None:
            first, second = self.swap
            i, j = steps.index(first), steps.index(second)
            steps[i], steps[j] = steps[j], steps[i]
        return steps

    def run(self) -> list[StepEvent]:
        self.log.clear()
        for step in self._ordered_plan():
            if step in self.drop:
                continue
            for attempt in range(1, self.retries.get(step, 1) + 1):
                self.log.append(
                    StepEvent(name=step, ts=self.clock.advance(), attempt=attempt)
                )
        return list(self.log)


def step_names(log: list[StepEvent]) -> tuple[str, ...]:
    return tuple(event.name for event in log)


def first_divergence(actual: tuple[str, ...], expected: tuple[str, ...]) -> int | None:
    """Index of the first position where the run left the plan, or None.

    Reporting the index beats reporting "sequences differ" when the plan has
    twenty steps and you are reading a CI log at 2am.
    """
    for index in range(max(len(actual), len(expected))):
        actual_step = actual[index] if index < len(actual) else None
        expected_step = expected[index] if index < len(expected) else None
        if actual_step != expected_step:
            return index
    return None


# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------


def test_workflow_executes_the_planned_sequence():
    log = FakeOrchestrator().run()

    assert step_names(log) == PLAN
    assert first_divergence(step_names(log), PLAN) is None


def test_step_timestamps_are_monotonic():
    """Out-of-order timestamps mean your trace is lying about causality."""
    log = FakeOrchestrator().run()

    timestamps = [event.ts for event in log]
    assert timestamps == sorted(timestamps)
    assert len(set(timestamps)) == len(timestamps)


# ---------------------------------------------------------------------------
# Failure mode 1: a silently dropped step
# ---------------------------------------------------------------------------


def test_dropped_step_is_caught():
    """The orchestrator skips review and the answer still looks plausible."""
    log = FakeOrchestrator(drop=frozenset({"review_answer"})).run()
    actual = step_names(log)

    assert actual != PLAN
    assert "review_answer" not in actual
    assert first_divergence(actual, PLAN) == PLAN.index("review_answer")


@pytest.mark.parametrize("dropped", PLAN)
def test_dropping_any_single_step_breaks_the_sequence(dropped: str):
    """No step in the plan is quietly optional."""
    actual = step_names(FakeOrchestrator(drop=frozenset({dropped})).run())

    assert actual != PLAN
    assert len(actual) == len(PLAN) - 1


# ---------------------------------------------------------------------------
# Failure mode 2: out-of-order execution
# ---------------------------------------------------------------------------


def test_out_of_order_step_is_caught():
    """Reviewing before drafting: every step ran, the order is nonsense."""
    log = FakeOrchestrator(swap=("draft_answer", "review_answer")).run()
    actual = step_names(log)

    assert sorted(actual) == sorted(PLAN), "same steps, so a set check would pass"
    assert actual != PLAN, "but the order is wrong"
    assert first_divergence(actual, PLAN) == PLAN.index("draft_answer")


# ---------------------------------------------------------------------------
# Failure mode 3: duplication from an unbounded retry
# ---------------------------------------------------------------------------


def test_duplicated_step_from_unbounded_retry_is_caught():
    """retrieve_context runs four times. Nothing errors. The bill triples."""
    log = FakeOrchestrator(retries={"retrieve_context": 4}).run()
    actual = step_names(log)

    counts = Counter(actual)
    assert counts["retrieve_context"] == 4
    assert actual != PLAN

    repeats = {name: n for name, n in counts.items() if n > 1}
    assert repeats == {"retrieve_context": 4}

    attempts = [e.attempt for e in log if e.name == "retrieve_context"]
    assert attempts == [1, 2, 3, 4]


def test_retry_ceiling_of_one_is_the_normal_case():
    log = FakeOrchestrator(retries={"retrieve_context": 1}).run()

    assert step_names(log) == PLAN
    assert all(event.attempt == 1 for event in log)

#!/usr/bin/env python3
"""Find the first divergent handoff in a multi-agent message log.

When a five-agent run produces a wrong answer, the useful question is not "was
the output wrong" but "which handoff was the first one to be wrong". Everything
after that point is a downstream symptom and reading it is wasted time.

This walks a message log in order and stops at the first message whose actual
payload no longer matches what the sender was supposed to emit.

Standard library only. No third-party dependencies.

Usage:
    python tools/trace_bisect.py sample_trace.json
    python tools/trace_bisect.py sample_trace.json --json
    python tools/trace_bisect.py sample_trace.json --correlation-id corr-9f2

Exit codes:
    0  no divergence found
    1  divergence found (the interesting case, so CI can gate on it)
    2  bad input (missing file, malformed JSON, malformed message)
"""

from __future__ import annotations

import argparse
import json
import sys
from typing import Any

REQUIRED_FIELDS = ("correlation_id", "sender", "recipient", "payload", "expected_payload")

EXIT_CLEAN = 0
EXIT_DIVERGED = 1
EXIT_BAD_INPUT = 2


class TraceError(Exception):
    """Raised when the trace file cannot be interpreted."""


def load_trace(path: str) -> list[dict[str, Any]]:
    """Read and structurally validate a message log."""
    try:
        with open(path, encoding="utf-8") as handle:
            data = json.load(handle)
    except FileNotFoundError as exc:
        raise TraceError(f"trace file not found: {path}") from exc
    except json.JSONDecodeError as exc:
        raise TraceError(f"{path} is not valid JSON: {exc}") from exc

    if not isinstance(data, list):
        raise TraceError(f"{path} must contain a JSON array of messages")
    if not data:
        raise TraceError(f"{path} contains no messages")

    for index, message in enumerate(data):
        if not isinstance(message, dict):
            raise TraceError(f"message {index} is not an object")
        missing = [name for name in REQUIRED_FIELDS if name not in message]
        if missing:
            raise TraceError(f"message {index} is missing fields: {', '.join(missing)}")

    return data


def diff_payload(actual: Any, expected: Any, prefix: str = "") -> list[str]:
    """Human-readable list of the specific fields that diverged."""
    if isinstance(actual, dict) and isinstance(expected, dict):
        findings: list[str] = []
        for key in sorted(set(actual) | set(expected)):
            path = f"{prefix}.{key}" if prefix else str(key)
            if key not in actual:
                findings.append(f"{path}: missing (expected {expected[key]!r})")
            elif key not in expected:
                findings.append(f"{path}: unexpected field ({actual[key]!r})")
            else:
                findings.extend(diff_payload(actual[key], expected[key], path))
        return findings

    if actual != expected:
        label = prefix or "payload"
        return [f"{label}: got {actual!r}, expected {expected!r}"]
    return []


def find_first_divergence(
    messages: list[dict[str, Any]],
    correlation_id: str | None = None,
) -> dict[str, Any] | None:
    """Walk the log in order; return the first message that diverged."""
    for index, message in enumerate(messages):
        if correlation_id is not None and message["correlation_id"] != correlation_id:
            continue
        if message["payload"] != message["expected_payload"]:
            return {
                "index": index,
                "correlation_id": message["correlation_id"],
                "sender": message["sender"],
                "recipient": message["recipient"],
                "differences": diff_payload(
                    message["payload"], message["expected_payload"]
                ),
                "payload": message["payload"],
                "expected_payload": message["expected_payload"],
            }
    return None


def format_report(result: dict[str, Any] | None, inspected: int) -> str:
    if result is None:
        return f"No divergence found across {inspected} message(s)."

    lines = [
        f"First divergent handoff at message index {result['index']}.",
        f"  correlation_id : {result['correlation_id']}",
        f"  {result['sender']} -> {result['recipient']}",
        "  differences:",
    ]
    lines.extend(f"    - {item}" for item in result["differences"])
    lines.append("")
    lines.append(
        f"Everything after index {result['index']} is downstream of this. "
        "Fix this handoff first."
    )
    return "\n".join(lines)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="trace_bisect.py",
        description="Find the first divergent handoff in a multi-agent message log.",
    )
    parser.add_argument("trace", help="path to a JSON message log")
    parser.add_argument(
        "--correlation-id",
        default=None,
        help="only inspect messages carrying this correlation id",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        dest="as_json",
        help="emit machine-readable JSON instead of a text report",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    try:
        messages = load_trace(args.trace)
    except TraceError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return EXIT_BAD_INPUT

    if args.correlation_id is not None:
        inspected = sum(
            1 for m in messages if m["correlation_id"] == args.correlation_id
        )
        if inspected == 0:
            print(
                f"error: no messages with correlation_id {args.correlation_id!r}",
                file=sys.stderr,
            )
            return EXIT_BAD_INPUT
    else:
        inspected = len(messages)

    result = find_first_divergence(messages, args.correlation_id)

    if args.as_json:
        print(json.dumps({"inspected": inspected, "divergence": result}, indent=2))
    else:
        print(format_report(result, inspected))

    return EXIT_CLEAN if result is None else EXIT_DIVERGED


if __name__ == "__main__":
    raise SystemExit(main())

name: multi-agent-tests

# The suite is deterministic and offline: no API keys, no live LLM calls, no
# network. That is what makes it safe to run on every pull request.

on:
  pull_request:
  push:
    branches: [main]
  workflow_dispatch:

jobs:
  pytest:
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.11", "3.12"]

    steps:
      - name: Check out the repository
        uses: actions/checkout@v4

      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
          cache: pip

      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install -r requirements.txt

      - name: Run the multi-agent test suite
        run: pytest -v

      - name: Bisect the sample trace
        # Exit code 1 is the expected result here: sample_trace.json ships with
        # a deliberate divergence so the helper has something to find.
        run: |
          python tools/trace_bisect.py sample_trace.json && exit 1 || true
          python tools/trace_bisect.py examples/clean_trace.json