Dev.to AI πŸ€– Ai πŸ‘ 0 πŸ“– 9 min read

Off-Trace Token Counts Cannot Gate a Run

A usage number that never becomes a span attribute cannot fail a build. It is a caption on a chat UI, and captions do not belong in CI. Agent runs hide cost in three places at once: the model call, the tool payload, and

A usage number that never becomes a span attribute cannot fail a build. It is a caption on a chat UI, and captions do not belong in CI.

Agent runs hide cost in three places at once: the model call, the tool payload, and the retries the panel collapses into one turn. Dashboards mix those numbers after the fact. A trace can keep them on the span that caused them, which is the only representation you can reject in a job without arguing about screenshots.

Think of the vendor usage page as a till receipt you glance at while walking out. The trace is the ledger. If the ledger does not foot, you do not negotiate with the receipt. You fail the run.

This is the slice of loop engineering that stays deterministic when a model writes the glue. Generating the next tool call is cheap. Declaring what that call was allowed to spend is not, unless the spend lives on the span that performed the work.

The schema stays small on purpose. Each span has a kind, a parent pointer, a closed flag, and three integers that default to absent rather than zero. Absent is not free. Zero is a claim that you measured nothing and found nothing.

# trace_budget.py
from __future__ import annotations

import json
import sys
from dataclasses import dataclass
from typing import Optional

REQUIRED_KINDS = {"run", "llm", "tool"}


@dataclass(frozen=True)
class Span:
    span_id: str
    parent_id: Optional[str]
    kind: str
    name: str
    input_tokens: Optional[int]
    output_tokens: Optional[int]
    tool_payload_bytes: Optional[int]
    closed: bool


def _opt_int(value, line_no: int) -> Optional[int]:
    if value is None:
        return None
    if isinstance(value, bool) or not isinstance(value, int):
        raise SystemExit(f"line {line_no}: usage fields must be int or null")
    if value < 0:
        raise SystemExit(f"line {line_no}: usage fields must be >= 0")
    return value


def load_jsonl(path: str) -> list[Span]:
    spans: list[Span] = []
    with open(path, encoding="utf-8") as fh:
        for line_no, raw in enumerate(fh, start=1):
            raw = raw.strip()
            if not raw:
                continue
            row = json.loads(raw)
            try:
                parent = row.get("parent_id")
                spans.append(
                    Span(
                        span_id=str(row["span_id"]),
                        parent_id=None if parent in (None, "") else str(parent),
                        kind=str(row["kind"]),
                        name=str(row["name"]),
                        input_tokens=_opt_int(row.get("input_tokens"), line_no),
                        output_tokens=_opt_int(row.get("output_tokens"), line_no),
                        tool_payload_bytes=_opt_int(
                            row.get("tool_payload_bytes"), line_no
                        ),
                        closed=bool(row.get("closed", False)),
                    )
                )
            except KeyError as exc:
                raise SystemExit(f"line {line_no}: missing {exc}") from exc
    return spans

Two rules keep the sums honest. The run root records no tokens and no payload bytes. Only llm and tool leaves do. If the root copies its children, the auditor double-counts, and the ceiling you tuned yesterday becomes a random number tomorrow.

Missing fields stay None. Treating a missing token count as zero makes a broken exporter look efficient. That is how a budget rots: the cheapest-looking run is the one that forgot to emit usage.

class BudgetError(Exception):
    pass


def audit(
    spans: list[Span],
    max_input_tokens: int,
    max_output_tokens: int,
    max_tool_bytes: int,
) -> dict:
    by_id: dict[str, Span] = {}
    for span in spans:
        if span.span_id in by_id:
            raise BudgetError(f"duplicate span_id {span.span_id}")
        if span.kind not in REQUIRED_KINDS:
            raise BudgetError(f"{span.span_id}: unknown kind {span.kind}")
        by_id[span.span_id] = span

    roots = [s for s in spans if s.parent_id is None]
    if len(roots) != 1 or roots[0].kind != "run":
        raise BudgetError("forest must have exactly one run root")
    if roots[0].input_tokens or roots[0].output_tokens or roots[0].tool_payload_bytes:
        raise BudgetError("run root must not carry usage integers")

    for span in spans:
        if span.parent_id is None:
            continue
        if span.parent_id not in by_id:
            raise BudgetError(f"{span.span_id}: missing parent")
        if not span.closed:
            raise BudgetError(f"{span.span_id}: open span")
        if span.kind == "llm" and (
            span.input_tokens is None or span.output_tokens is None
        ):
            raise BudgetError(f"{span.span_id}: llm span missing token attributes")
        if span.kind == "tool" and span.tool_payload_bytes is None:
            raise BudgetError(f"{span.span_id}: tool span missing payload bytes")

    input_tokens = sum(s.input_tokens or 0 for s in spans)
    output_tokens = sum(s.output_tokens or 0 for s in spans)
    tool_bytes = sum(s.tool_payload_bytes or 0 for s in spans)

    if input_tokens > max_input_tokens:
        raise BudgetError(f"input_tokens {input_tokens} > {max_input_tokens}")
    if output_tokens > max_output_tokens:
        raise BudgetError(f"output_tokens {output_tokens} > {max_output_tokens}")
    if tool_bytes > max_tool_bytes:
        raise BudgetError(f"tool_payload_bytes {tool_bytes} > {max_tool_bytes}")

    return {
        "spans": len(spans),
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "tool_payload_bytes": tool_bytes,
    }

The fail cases are mechanical, not stylistic. An open span means the exporter died mid-tool, so the cost is unknown. A missing parent means spend leaked out of the run. A missing integer means you are about to crown a partial trace as a saving. Duplicate span_id values are not a merge. They are a broken primary key.

Calibrate ceilings from traces the process actually wrote. Emit JSONL from the agent, run the auditor, then tighten the integers. A bound copied from a chat panel will miss the first hidden retry, because the panel is allowed to collapse turns and the trace is not.

def main(argv: list[str]) -> int:
    if len(argv) != 5:
        sys.stderr.write(
            "usage: python trace_budget.py RUN.jsonl MAX_IN MAX_OUT MAX_TOOL_BYTES\n"
        )
        return 2
    path, max_in, max_out, max_bytes = argv[1], int(argv[2]), int(argv[3]), int(argv[4])
    try:
        summary = audit(load_jsonl(path), max_in, max_out, max_bytes)
    except BudgetError as exc:
        sys.stderr.write(f"BUDGET {exc}\n")
        return 1
    sys.stdout.write(json.dumps(summary, indent=2) + "\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))

A fixture makes the contract visible. The first file is a legal run. The second file drops output_tokens on the llm span. The first command must exit 0. The second must exit 1. If both pass, the auditor is treating absence as zero, and the gate is theater.

{"span_id":"r1","parent_id":null,"kind":"run","name":"issue-triage","closed":true}
{"span_id":"m1","parent_id":"r1","kind":"llm","name":"plan","input_tokens":412,"output_tokens":86,"closed":true}
{"span_id":"t1","parent_id":"r1","kind":"tool","name":"read_file","tool_payload_bytes":1540,"closed":true}
{"span_id":"m2","parent_id":"r1","kind":"llm","name":"answer","input_tokens":903,"output_tokens":140,"closed":true}
python trace_budget.py good.jsonl 2000 400 8000
python trace_budget.py missing_tokens.jsonl 2000 400 8000; echo $?

Retries are where panel counts lie hardest. One user turn in the UI can be two llm spans with the same name after a swallowed timeout. Summing attributes makes that cost visible. A single β€œtokens used” caption does not, because the caption is allowed to describe the conversation instead of the forest.

Tool payloads need their own integer because they are not tokens. A read_file call can push a hundred kilobytes into the next prompt without a polite usage field from the model vendor. Converting bytes to tokens with a homemade factor will be wrong on binary files, on compressed JSON, and on every encoding edge a demo skips. Keep the units separate. Fail them separately.

Wall-clock duration is a third unit, and it does not belong in this sum. A slow tool on a cold disk can look expensive in time while remaining cheap in tokens. Mixing those into one score produces a number you cannot explain when it moves. If you need latency, put exclusive duration on the span later. Do not smuggle it into a token ceiling.

Generating calibration traces still costs something, even when the gate is local. Disclosure: This article was prepared as part of MonkeyCode's product outreach. MonkeyCode is an open-source coding agent with free model access and a free server option, which is enough to emit JSONL spans while you are still choosing ceilings, without pointing the same loop at a paid API. The checker does not depend on that stack. If you already write spans, run it on those.

Pytest pins the fail-closed policy so a later refactor cannot quietly start treating None as zero.

# test_trace_budget.py
import json
from pathlib import Path

import pytest
import trace_budget as tb


def write_jsonl(path: Path, rows: list[dict]) -> None:
    path.write_text("".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8")


def sample_rows() -> list[dict]:
    return [
        {"span_id": "r1", "parent_id": None, "kind": "run", "name": "issue-triage", "closed": True},
        {
            "span_id": "m1",
            "parent_id": "r1",
            "kind": "llm",
            "name": "plan",
            "input_tokens": 412,
            "output_tokens": 86,
            "closed": True,
        },
        {
            "span_id": "t1",
            "parent_id": "r1",
            "kind": "tool",
            "name": "read_file",
            "tool_payload_bytes": 1540,
            "closed": True,
        },
    ]


def test_sum_and_pass(tmp_path: Path) -> None:
    path = tmp_path / "good.jsonl"
    write_jsonl(path, sample_rows())
    summary = tb.audit(tb.load_jsonl(str(path)), 2000, 400, 8000)
    assert summary["input_tokens"] == 412
    assert summary["output_tokens"] == 86
    assert summary["tool_payload_bytes"] == 1540


def test_missing_output_tokens_fails(tmp_path: Path) -> None:
    rows = sample_rows()
    del rows[1]["output_tokens"]
    path = tmp_path / "bad.jsonl"
    write_jsonl(path, rows)
    with pytest.raises(tb.BudgetError, match="missing token"):
        tb.audit(tb.load_jsonl(str(path)), 2000, 400, 8000)


def test_root_usage_is_rejected(tmp_path: Path) -> None:
    rows = sample_rows()
    rows[0]["input_tokens"] = 412
    path = tmp_path / "double.jsonl"
    write_jsonl(path, rows)
    with pytest.raises(tb.BudgetError, match="run root"):
        tb.audit(tb.load_jsonl(str(path)), 2000, 400, 8000)


def test_open_span_fails(tmp_path: Path) -> None:
    rows = sample_rows()
    rows[2]["closed"] = False
    path = tmp_path / "open.jsonl"
    write_jsonl(path, rows)
    with pytest.raises(tb.BudgetError, match="open span"):
        tb.audit(tb.load_jsonl(str(path)), 2000, 400, 8000)
python -m pytest test_trace_budget.py -q

The method is narrow on purpose. It is not a billing system. Provider invoices remain the financial source of truth, because streaming usage is easy to under-count and because cached prompt tokens are not described the same way across vendors. Sampling will also lie. If you drop most tool spans, the sum is a lower bound, not a budget.

Do not put this checker in front of a pipeline that never closes spans, such as a streaming demo that keeps the root open until the socket dies. Do not use the sum as a quality score. A cheap run can still be wrong. Teams that do not emit parent ids should not start here; fix the forest first. Teams that need per-tenant chargeback should not start here either. This auditor has no customer, no discount, and no reserved capacity. It answers one question: did this run stay inside the bound you declared for this job.

Once the gate is in CI, stop reading the chat UI for spend. The UI is still useful for reading the answer. It is not useful for deciding whether the loop is allowed to ship.

πŸ“° Read the original article on Dev.to AI

Originally published by Dev.to AI. Aggregated on AIWithGhost for educational purposes β€” full credit and traffic to the original publisher.