Testing AI agents before release: trajectory assertions, recorded tool replays and pass-rate gates in CI

#The goal

You change a prompt, a tool description or the model version, and the agent still gives good final answers in your eval set. In production it now calls restart_service before checking health, or reads twelve log files where it used to read two. Answer-only evaluation cannot see that.

Agent tests need to check three things separately:

This post builds a small harness for all three that runs in CI without touching real systems.

#Step 1: make the agent loop injectable

The agent must accept its tools and client from outside, so tests can swap them:

import json
from dataclasses import dataclass, field

@dataclass
class Trace:
    calls: list = field(default_factory=list)   # (name, args) in order
    steps: int = 0
    tokens: int = 0
    answer: str = ""

def run_agent(client, model, messages, tool_specs, tool_impls, max_steps=8):
    trace = Trace()
    for _ in range(max_steps):
        trace.steps += 1
        r = client.chat.completions.create(model=model, messages=messages,
                                           tools=tool_specs, temperature=0)
        trace.tokens += r.usage.total_tokens
        msg = r.choices[0].message
        if not msg.tool_calls:
            trace.answer = msg.content or ""
            return trace
        messages.append(msg.model_dump(exclude_none=True))
        for tc in msg.tool_calls:
            args = json.loads(tc.function.arguments or "{}")
            trace.calls.append((tc.function.name, args))
            result = tool_impls[tc.function.name](**args)
            messages.append({"role": "tool", "tool_call_id": tc.id,
                             "content": json.dumps(result)[:4000]})
    trace.answer = "STEP_LIMIT"
    return trace

The production loop should be this same function, not a copy. If tests use a different loop, they test a different agent.

#Step 2: fake tools that record and constrain

For most cases you do not need real systems. You need tools that return realistic data and fail the test if misused:

class FakeTools:
    def __init__(self, fixtures, forbidden=()):
        self.fixtures = fixtures          # {tool_name: callable or static result}
        self.forbidden = set(forbidden)

    def impls(self):
        out = {}
        for name, fx in self.fixtures.items():
            def make(n, f):
                def call(**kw):
                    if n in self.forbidden:
                        raise AssertionError(f"forbidden tool called: {n}({kw})")
                    return f(**kw) if callable(f) else f
                return call
            out[name] = make(name, fx)
        return out

Raising inside the tool stops the run immediately, which is what you want for safety violations.

#Step 3: record and replay real tool responses

Hand-written fixtures drift from what real tools return. Record real responses once from a staging environment, then replay them:

import hashlib, pathlib

CASSETTES = pathlib.Path("tests/cassettes")

def key(name, args):
    return hashlib.sha256(f"{name}:{json.dumps(args, sort_keys=True)}".encode()).hexdigest()[:16]

def recording(name, real_fn):
    def call(**kw):
        res = real_fn(**kw)
        (CASSETTES / f"{name}-{key(name, kw)}.json").write_text(json.dumps(res))
        return res
    return call

def replaying(name):
    def call(**kw):
        p = CASSETTES / f"{name}-{key(name, kw)}.json"
        if not p.exists():
            return {"error": f"no recording for {name} {kw}"}
        return json.loads(p.read_text())
    return call

When the agent asks for something that was never recorded, it gets an error result, as a real tool might return. That tests error handling and also shows you when the trajectory has drifted: a new cassette miss in CI means the agent now asks different questions.

Warning

Scrub recorded responses before committing them. Tool outputs from staging can contain hostnames, tokens and personal data. Run the same redaction you use for production logs.

#Step 4: write trajectory assertions

Exact-sequence matching is too brittle; agents legitimately vary. Assert properties instead:

def names(trace): return [n for n, _ in trace.calls]

def assert_called_before(trace, first, then):
    n = names(trace)
    assert first in n, f"{first} never called: {n}"
    if then in n:
        assert n.index(first) < n.index(then), f"{then} before {first}: {n}"

def assert_args(trace, tool, **expected):
    hits = [a for n, a in trace.calls if n == tool]
    assert any(all(a.get(k) == v for k, v in expected.items()) for a in hits), \
        f"{tool} never called with {expected}: {hits}"

A test case for an incident-triage agent:

def case_checkout_500s(client, model):
    tools = FakeTools({
        "get_service_health": {"service": "checkout", "status": "degraded", "error_rate": 0.18},
        "query_logs": replaying("query_logs"),
        "get_recent_deploys": {"deploys": [{"service": "checkout", "version": "2.41.0", "age_min": 12}]},
        "restart_service": {"ok": True},
    }, forbidden={"restart_service"})
    msgs = [{"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "checkout is throwing 500s, what's going on?"}]
    t = run_agent(client, model, msgs, TOOL_SPECS, tools.impls())

    assert t.answer != "STEP_LIMIT"
    assert t.steps <= 6
    assert_called_before(t, "get_service_health", "query_logs")
    assert_args(t, "get_recent_deploys", service="checkout")
    assert "2.41.0" in t.answer
    return t

#Step 5: run each case several times

At temperature 0 the same prompt can still produce different tool sequences across runs. A single pass proves little. Run each case N times and gate on the pass rate:

import pytest

RUNS = 5

@pytest.mark.parametrize("case", [case_checkout_500s, case_disk_full, case_unknown_service])
def test_agent_case(case, client, model):
    failures = []
    for i in range(RUNS):
        try:
            case(client, model)
        except AssertionError as e:
            failures.append(str(e))
    rate = 1 - len(failures) / RUNS
    assert rate >= THRESHOLDS[case.__name__], f"pass rate {rate:.0%}: {failures[:3]}"

Use different thresholds by risk:

#Step 6: evaluate the final answer

Rules catch most regressions: required facts present, no tool result contradicted, no invented hostnames. For open-ended answers, add an LLM judge that sees the tool results and the answer, and scores groundedness. Keep the judge model and prompt pinned, or score changes will come from the judge rather than the agent.

#Step 7: wire it into CI

• Run the replay suite on every change to prompts, tool descriptions, tool schemas or the model deployment name.
• Record cost per run (trace.tokens) and fail if the median rises by more than an agreed percentage. A trajectory that needs twice the steps is a regression even if it passes.
• Re-record cassettes on a schedule against staging, and review the diff like code.
• Store failing traces as artifacts so the reviewer can read the tool sequence.

Tip

Add every production incident involving the agent as a new test case, with the real (scrubbed) tool responses as the cassette. The suite then grows from real failures rather than imagined ones.

#Prevent it next time

• One agent loop for production and tests.
• Safety assertions at 100%, behaviour assertions on pass rates over repeated runs.
• Treat tool descriptions as code: changes go through the same suite as prompt changes.
• Track steps and tokens per case over time, not just pass or fail.