Debugging a multi-step agent is painful when you can’t trace tool calls AI agent loops make across model invocations. This guide gives you a concrete pattern to instrument an agent runtime with correlation IDs, structured spans, and gateway-aware logging so you can reconstruct any run.
Step 1: Create a run-scoped correlation context
Every agent run needs a single identifier that threads through model calls and tool executions. Use a contextvars variable so the ID survives across async boundaries and thread pools without explicit parameter passing.
import contextvars
import uuid
run_id_ctx = contextvars.ContextVar("run_id")
def start_run() -> str:
rid = uuid.uuid4().hex
run_id_ctx.set(rid)
return rid
def get_run_id() -> str:
return run_id_ctx.get("unknown")Call start_run() at the entry point of your agent. Any code that later logs or records a span can pull get_run_id() without plumbing it through signatures.
Step 2: Wrap tool execution in a span
Agents call functions based on model output. You need timing, inputs, outputs, and errors for each invocation. A decorator keeps the instrumentation consistent and avoids copy-paste logging.
import time
import functools
import logging
logger = logging.getLogger("agent.tools")
def trace_tool(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
run_id = get_run_id()
start = time.monotonic()
try:
result = func(*args, **kwargs)
dur = time.monotonic() - start
logger.info({
"type": "tool_span",
"run_id": run_id,
"tool": func.__name__,
"ok": True,
"duration_ms": round(dur*1000, 2),
"args": args,
"kwargs": kwargs,
"result": result,
})
return result
except Exception as e:
dur = time.monotonic() - start
logger.error({
"type": "tool_span",
"run_id": run_id,
"tool": func.__name__,
"ok": False,
"duration_ms": round(dur*1000, 2),
"error": str(e),
})
raise
return wrapper
@trace_tool
def get_weather(city: str) -> str:
return f"Sunny in {city}"The log record is structured, not a string, so you can ship it to JSON sinks or OpenTelemetry without parsing.
Step 3: Extract tool-call metadata from the model response
The chat completions API returns tool_calls with IDs. Capture them alongside token usage so you can link a model turn to the exact tool invocations it triggered.
from openai import OpenAI
client = OpenAI() # base_url set later in Step 4
def model_turn(messages):
resp = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
tools=[{
"type": "function",
"function": {
"name": "get_weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
},
}],
)
msg = resp.choices[0].message
run_id = get_run_id()
logger.info({
"type": "model_span",
"run_id": run_id,
"model": resp.model,
"prompt_tokens": resp.usage.prompt_tokens,
"completion_tokens": resp.usage.completion_tokens,
"tool_calls": [
{"id": tc.id, "name": tc.function.name, "args": tc.function.arguments}
for tc in (msg.tool_calls or [])
],
})
return msgStore the tool_calls[].id values. They are the join key between the model span and the subsequent tool span if you propagate them into the tool call.
Step 4: Route through an OpenAI-compatible gateway
Point the SDK at a gateway that normalizes provider differences. An OpenAI-compatible endpoint such as n4n.ai addresses 240+ models, honors client routing directives, and forwards provider cache-control hints, so your trace can record cache hits without extra plumbing. Use extra_headers to pass routing or cache directives.
client = OpenAI(
base_url="https://api.n4n.ai/v1",
api_key="YOUR_KEY",
)
# Request cache usage reporting from the provider
resp = client.chat.completions.create(
model="anthropic/claude-3.5-sonnet",
messages=messages,
extra_headers={"x-cache-control": "read+write"},
)If a provider is degraded, the gateway’s automatic fallback keeps the run alive; your run_id stays constant, so the trace shows the switch as a new model_span with a different model field rather than a broken run.
Step 5: Emit traces to a store you can query
Structured logs are enough for local debugging, but for multi-step agents you want a queryable store. Write spans as JSON lines to a file or pipe them to an OpenTelemetry collector.
import json
import logging
class JsonFileHandler(logging.Handler):
def __init__(self, path):
super().__init__()
self.f = open(path, "a")
def emit(self, record):
if isinstance(record.msg, dict):
self.f.write(json.dumps(record.msg) + "\n")
self.f.flush()
logging.basicConfig(level=logging.INFO)
logging.getLogger().addHandler(JsonFileHandler("agent_traces.jsonl"))After a run, grep or jq by run_id:
jq 'select(.run_id == "abc123")' agent_traces.jsonlYou will see an ordered list of model_span and tool_span entries that reconstruct the agent’s path.
Step 6: Verify the trace end to end
Write a test that runs a two-step agent with a fake tool and asserts the trace contains the expected entries. This catches regressions in context propagation.
def test_trace_tool_calls_ai_agent(tmp_path):
trace_file = tmp_path / "traces.jsonl"
logging.getLogger().addHandler(JsonFileHandler(trace_file))
rid = start_run()
messages = [{"role": "user", "content": "Weather in Berlin?"}]
msg = model_turn(messages)
if msg.tool_calls:
for tc in msg.tool_calls:
# simulate executor
get_weather(**json.loads(tc.function.arguments))
lines = [json.loads(l) for l in trace_file.read_text().splitlines()]
spans = [s for s in lines if s.get("run_id") == rid]
assert any(s["type"] == "model_span" for s in spans)
assert any(s["type"] == "tool_span" and s["tool"] == "get_weather" for s in spans)Run it with pytest. If both assertions pass, your instrumentation correctly captures the full chain.
What a correct trace looks like
{"type":"model_span","run_id":"def456","model":"gpt-4o-mini","prompt_tokens":42,"completion_tokens":18,"tool_calls":[{"id":"call_1","name":"get_weather","args":"{\"city\":\"Berlin\"}"}]}
{"type":"tool_span","run_id":"def456","tool":"get_weather","ok":true,"duration_ms":1.2,"args":(),"kwargs":{"city":"Berlin"},"result":"Sunny in Berlin"}The run_id ties them together. The tool_calls[].id from the model span matches the invocation context if you extend the tool span to include it.
Common pitfalls
- Thread pools drop contextvars. If you execute tools in
concurrent.futures.ThreadPoolExecutor, copy the context withcontextvars.copy_context()before submitting. - Missing tool_call_id mapping. Always log the model-assigned
idalongside the tool execution. Without it, you cannot prove which model turn triggered which call. - String logging. Avoid
logger.info(f"called {tool}"). Structured dicts are queryable; strings are not. - Ignoring cache hits. If your gateway reports
prompt_tokenswith a cache discount, recordresp.usage.prompt_tokens_detailsif present. That data explains why a step was cheap.
Following these steps gives you a defensible way to trace tool calls AI agent runs make, turning opaque multi-step loops into auditable event streams.