DEV Community

Ali Abdalla
Ali Abdalla

Posted on

I built a spending gate for AI agents. Then I proved it doesn't hold under load

A2A — the open protocol that lets one AI agent hand off work to another — solves discovery and messaging well. It does not solve a question that matters the moment money is involved: when Agent A delegates a task to Agent B, what stops that from quietly running up a bill nobody capped?

I built a small gateway to answer that. It sits between two agents, estimates what each delegated task costs, and can refuse to forward a request once an agent's spending limit is exhausted. It works. It's also got a real, unresolved bug in exactly the part that's supposed to keep you safe — and I want to walk through both the design and the bug, because the bug is more interesting than the feature.

The setup

Two toy agents talk over real A2A: a Coordinator that delegates a task, and a Specialist that completes it. A gateway sits between them, transparently forwarding traffic — with one deliberate exception. It rewrites the Specialist's advertised address to point at itself, so the Coordinator discovers the gateway instead of the real agent. That's the only place the "transparent" proxy touches the wire on purpose; everything else passes through byte-for-byte.

Sitting in that position, the gateway can see every request and response. That visibility is what makes cost estimation and budget enforcement possible at all.

Estimating cost without lying about precision

The gateway never sees what a Specialist does internally — no visibility into system prompts, retries, or tool calls. All it ever sees is the input message and the output artifact. So "cost estimation" here means one of three things, ranked by trust and labeled accordingly in every log record:

  1. Self-reported — a cooperating agent voluntarily attaches real usage data. Ground truth, when available.
  2. Provider-tokenizer estimate — the gateway knows which model backs an agent, and runs the real tokenizer for that provider (verified against OpenAI's tiktoken, Mistral's Tekken tokenizer, and Google's experimental local Gemini tokenizer — each checked by actually installing the library and running it, not assumed from documentation).
  3. Generic fallback — a rough, provider-agnostic character count, used when nothing else applies.

Every logged record says which tier produced it. An estimate is never allowed to look like a measurement. That distinction matters more once you're deciding whether to reject a real request based on the number.

Gating on the past, not the future

Here's the part that isn't obvious until you sit with it: the gateway cannot know a task's true cost before forwarding it. Cost is only computable after the Specialist responds — that's the entire premise of the estimation system above. So a budget check can't ask "will this cost too much?" There's no answer to that question yet.

The only honest thing a budget check can do is ask a different question: has this agent already spent enough that a new request shouldn't be allowed to start? The gateway keeps a running total of an agent's accumulated spend, drawn from real completed-task estimates, and checks that total — plus a cheap pre-flight estimate of just the new request's input side, which is knowable before sending — against a configured ceiling.

This means a single task can still overshoot the budget if its actual response turns out far more expensive than its input suggested. Nothing can prevent that from outside — the overspend is only knowable after the response has already been delivered. The gate stops new expensive-looking work once the well is dry; it can't cap what a request already in flight will cost. I don't think that's a flaw to hide — it's a structural fact about what a network-level proxy can and can't know, and pretending otherwise would be worse than admitting it.

The bug: what "the past" means gets fuzzy under concurrency

The budget check does three things in sequence: read the agent's current spend, decide whether to forward, and — after the response comes back — record the new spend. Between the read and the record, there's a real await on the network call to the Specialist. Python's asyncio can interleave other work during any await.

That's a check-then-act race, and I didn't just suspect it existed — I reproduced it on purpose, in tests/test_budget_race_condition.py. I configured a budget sized to allow exactly one request, then fired five concurrent delegations at the same agent. Sequentially, the math is unambiguous: request one exhausts the budget, and every request after it should be rejected.

Concurrently, all five were allowed through.

Each of the five reads the "current spend" value before any of the others has finished recording theirs — so each one sees a budget that looks unspent, passes its check, and proceeds. The gate isn't wrong about any individual decision; it's reasoning from state that's already stale by the time the decision matters.

Why I'm not fixing it in this post

The real fix is straightforward to describe and non-trivial to add correctly: a per-agent lock, or a compare-and-set, around the entire check-forward-record sequence, so concurrent requests to the same agent serialize at exactly the point that currently races. I know what the fix is. I haven't shipped it yet, on purpose — this project's whole discipline has been "verify before you build," and I wanted the bug documented, reproduced, and provably real before touching the code that's supposed to prevent it. A test now exists that fails safe: it currently passes by confirming the bug is present, and it's written so that fixing the race will flip that assertion — which is the point. If a future version of this makes that test start failing, that's the fix working, not a regression.

Why this is worth reading even though it's not solved

Most "I built an AI tool" posts show the happy path. I think the more useful thing to publish is the place where a system's own logic runs into a wall it can't will its way past — because that's the part someone building something similar will actually hit too, and knowing it's coming is worth more than another polished demo.

If you're building on A2A and have run into this — or solved it — I'd genuinely like to hear how.


Code: github.com/AliAbdallah21/a2a-cost-gateway. The race-condition test lives at tests/test_budget_race_condition.py if you want to see it reproduce the bug yourself — the full design reasoning, including why it's fine as-is, is in the repo's arch.md.

Top comments (0)