Put a Budget Boundary Between AX Agents and Your API Key

At 3:17 AM, your on-call pager fires: the team's shared API key just burned through $1,200 in six minutes. One engineer's agent loop retried a malformed prompt 847 times before anyone noticed. By morning, the monthly budget is gone and three active projects are blocked. For a ten-person team evaluating google/ax, this is the operational problem I would solve before increasing concurrency: one engineer's retry loop must not consume everybody else's monthly allowance.

I am treating AX as the agentic orchestration runtime described in the repository context, not as a verified quota-management product. I could not inspect the live repository here, so this design does not assume a particular AX release, configuration schema, or built-in billing feature. The integration boundary is deliberately outside the executor: route its model requests through a small admission gateway.

Separate execution from permission to spend

AX owns agent execution. The gateway owns per-seat credentials, model restrictions, and budget admission. The upstream relay owns provider routing.

Each engineer receives a gateway sub-token, never the upstream credential. Before forwarding a request, the gateway atomically reserves its maximum permitted cost against both a seat allocation and the team allocation. Concurrent agent calls cannot all observe the same remaining balance.

For this topology, the upstream base URL is https://b-lost.com/v1. Its presence is a routing choice, not evidence of verified model availability, prices, retention settings, or compatibility. Those require deployment-time validation.

The important distinction is between admission accounting and provider billing. The example retains reservations even when requests fail. This intentionally overcounts rather than reopening spending capacity after an ambiguous timeout.

One deployable gateway boundary

The following single-process service uses SQLite transactions for durable admission accounting. It accepts only text chat messages, one configured model, and bounded output. Install fastapi, uvicorn, httpx, and tiktoken; save this as gateway.py.

Set UPSTREAM_KEY, MODEL, and SEATS_JSON before starting it. For example, SEATS_JSON maps a randomly generated sub-token to "alice". Set price ceilings from verified upstream terms, not from a model-name assumption.

import asyncio
import datetime as dt
import json
import math
import os
import sqlite3

import httpx
import tiktoken
from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import Response

app = FastAPI()
base_url = "https://b-lost.com/v1"
MODEL = os.environ["MODEL"]
KEY = os.environ["UPSTREAM_KEY"]
SEATS = json.loads(os.environ["SEATS_JSON"])
INPUT_RATE = float(os.environ["INPUT_USD_PER_MILLION"])
OUTPUT_RATE = float(os.environ["OUTPUT_USD_PER_MILLION"])
SEAT_CAP, TEAM_CAP, DAILY_ALERT = 2000, 10000, 5000  # cents
ENC = tiktoken.get_encoding("o200k_base")
LOCK = asyncio.Lock()

db = sqlite3.connect("budget.db", isolation_level=None, timeout=30.0)
db.execute("PRAGMA journal_mode=WAL")
db.execute("""CREATE TABLE IF NOT EXISTS reservations
              (month TEXT, day TEXT, seat TEXT, cents INTEGER)""")
db.execute("CREATE INDEX IF NOT EXISTS idx_day ON reservations(day)")
db.execute("CREATE INDEX IF NOT EXISTS idx_month_seat ON reservations(month, seat)")

def reserve(seat, cents):
    now = dt.datetime.now(dt.timezone.utc)
    month, day = now.strftime("%Y-%m"), now.strftime("%Y-%m-%d")
    db.execute("BEGIN IMMEDIATE")
    try:
        team, personal = db.execute(
            """SELECT COALESCE(SUM(cents),0),
               COALESCE(SUM(CASE WHEN seat=? THEN cents ELSE 0 END),0)
               FROM reservations WHERE month=?""", (seat, month)
        ).fetchone()
        if team + cents > TEAM_CAP or personal + cents > SEAT_CAP:
            raise HTTPException(429, "Monthly reservation cap reached")
        db.execute("INSERT INTO reservations VALUES (?,?,?,?)",
                   (month, day, seat, cents))
        daily = db.execute(
            "SELECT SUM(cents) FROM reservations WHERE day=?", (day,)
        ).fetchone()[0]
        db.execute("COMMIT")
        return daily
    except Exception:
        db.execute("ROLLBACK")
        raise

@app.post("/v1/chat/completions")
async def chat(body: dict, authorization: str = Header(default="")):
    seat = SEATS.get(authorization.removeprefix("Bearer "))
    if not seat:
        raise HTTPException(401, "Unknown sub-token")
    messages = body.get("messages")
    if not isinstance(messages, list) or not 1 <= len(messages) <= 32:
        raise HTTPException(400, "Expected 1-32 text messages")
    tokens = 128
    for message in messages:
        if not isinstance(message, dict):
            raise HTTPException(400, "Invalid message")
        if message.get("role") not in ("system", "user", "assistant"):
            raise HTTPException(400, "Unsupported role")
        text = message.get("content")
        if not isinstance(text, str) or len(text.encode()) > 16384:
            raise HTTPException(400, "Text content too large")
        tokens += len(ENC.encode(text, allowed_special="all")) + 16
    if tokens > 8192:
        raise HTTPException(400, "Input budget exceeded")
    requested = body.get("max_tokens", 512)
    if type(requested) is not int or not 1 <= requested <= 1024:
        raise HTTPException(400, "max_tokens must be 1-1024")
    cents = max(1, math.ceil(
        (tokens * INPUT_RATE + requested * OUTPUT_RATE) / 10000
    ))
    async with LOCK:
        daily = reserve(seat, cents)
    payload = dict(model=MODEL, messages=messages,
                   max_tokens=requested, stream=False)
    async with httpx.AsyncClient(
        timeout=httpx.Timeout(10.0, connect=5.0, read=30.0),
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
        follow_redirects=False
    ) as client:
        webhook = os.getenv("BUDGET_WEBHOOK")
        if webhook and daily >= DAILY_ALERT:
            asyncio.create_task(client.post(webhook, json={"reserved_cents": daily}))
        try:
            upstream = await client.post(
                f"{base_url}/chat/completions", json=payload,
                headers={"Authorization": f"Bearer {KEY}"})
        except httpx.HTTPError:
            raise HTTPException(502, "Upstream outcome uncertain; budget retained")
    return Response(upstream.content, status_code=upstream.status_code,
                    media_type="application/json")

@app.on_event("startup")
async def startup():
    asyncio.create_task(maintenance())

async def maintenance():
    while True:
        await asyncio.sleep(600)
        db.execute("PRAGMA wal_checkpoint(TRUNCATE)")

# Run behind TLS: uvicorn gateway:app --host 127.0.0.1 --port 8080 --workers 1

Configure AX's model transport to use this gateway's /v1 endpoint and a seat sub-token, wherever the checked-out version exposes that transport setting. Do not confuse the upstream base_url with the address AX should call.

Where the tutorial ends

This gateway enforces a hard cap on recorded reservations, not an unconditional cap on invoices. Its token estimate and configured price ceilings must conservatively cover the selected model. Validate message overhead, tokenizer compatibility, output-limit semantics, and any additional billable tokens before enabling real traffic. Models with unbounded hidden charges do not belong behind this accounting rule.

Rejecting tools, images, streaming, and caller-selected models is intentional. Each adds another accounting path. Full AX workflows may require tool messages; extend the schema only alongside corresponding accounting tests.

A timeout does not prove that generation stopped. Keeping the reservation prevents repeated timeouts from replenishing an agent's allowance. The cost is unused capacity. Provider-verified reconciliation can recover it later without weakening admission control.

The daily webhook reports reserved spending, not settled spending. It repeats after $50; production delivery needs deduplication and a retry queue. Protect the SQLite file on persistent storage and test simultaneous requests at the cap boundary.

Finally, not logging prompts is not zero retention. Disable request-body capture in surrounding infrastructure, and verify retention policies for both relay and downstream provider. If those guarantees cannot be established, keep private source code out of the workflow.

---

The hardest operational dilemma: when should a reservation expire? Keeping it forever protects against silent failures but locks up capacity that may never be invoiced. Releasing it too soon lets a flaky network turn one malformed prompt into a $10,000 retry storm. Where do you draw that line, and how do you verify that upstream reconciliation actually recovers the difference? I'd be curious to hear how other teams handle this tension between admission control and resource efficiency.