DEV Community

Cover image for Put Authentication Before AI Budgets
James LIN
James LIN

Posted on

Put Authentication Before AI Budgets

Put Authentication Before AI Budgets: oauth2-proxy for a Ten-Person Team

At 3:00 AM, an engineer deploys a retry loop. Authentication is working perfectly, the session is valid, and the AI provider accepts every request—until the team wakes up to a blown budget and an invoice nobody expected. The failure is not identity. It is admission control.

For a small engineering team, I would evaluate oauth2-proxy as the identity boundary, not the quota engine. It authenticates people through an identity provider and forwards their identity upstream. A separate, private gateway must decide whether each request can spend money.

The architecture is deliberately small:

Developer → oauth2-proxy → local budget gateway → AI relay

The example uses https://b-lost.com/v1 as the relay endpoint. That configuration does not establish its retention policy or pricing; both require separate verification.

Reserve Before Forwarding

Checking yesterday's usage does not stop today's concurrent requests. Even checking a balance immediately before forwarding fails if ten requests read the same remaining allowance.

The gateway below serializes admission through SQLite transactions. Each accepted request consumes a conservative $0.15 allowance against four limits: the engineer's daily and monthly caps, and the team's daily and monthly caps. Reservations are not refunded, including after failures.

That intentionally sacrifices utilization for predictable admission. It is a spending envelope, not invoice reconciliation.

The example permits one model, one bounded user message, and 1,024 output tokens. Its reservation assumes a verified upper price envelope of $10 per million input tokens and $30 per million output tokens. With an 8,192-byte message and protocol overhead, $0.15 leaves margin. Confirm the selected provider's tokenization, billing rules, and rates before treating that envelope as a financial guarantee.

One Executable Configuration

Install Python 3 and an independently verified oauth2-proxy release. Register an OIDC client with callback https://ai.team.example:8443/oauth2/callback, provision a certificate, and export the referenced secrets. Change the example domain to your verified organization domain.

Run this on a dedicated host without untrusted local users: the budget server trusts identity headers only because it listens on loopback.

#!/usr/bin/env bash
set -euo pipefail
: "${OAUTH2_PROXY_CLIENT_ID:?}" "${OAUTH2_PROXY_CLIENT_SECRET:?}"
: "${OAUTH2_PROXY_COOKIE_SECRET:?}" "${OIDC_ISSUER_URL:?}"
: "${RELAY_KEY:?}" "${ALERT_WEBHOOK:?}"
umask 077

cat > budget.py <<'PY'
import datetime as dt
import json, os, sqlite3, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

base_url = "https://b-lost.com/v1"
with sqlite3.connect("budget.db") as db:
    db.execute("""CREATE TABLE IF NOT EXISTS spend
                  (bucket TEXT PRIMARY KEY, cents INTEGER NOT NULL)""")

def alert(day, cents):
    # Use a dedicated webhook accepting JSON; never include request content.
    try:
        body = json.dumps({"day": day, "reserved_cents": cents}).encode()
        req = urllib.request.Request(
            os.environ["ALERT_WEBHOOK"], data=body,
            headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=3):
            pass
    except Exception:
        print("budget alert delivery failed", flush=True)

class Gateway(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass  # Do not log prompts, credentials, or identities.

    def reply(self, code, body):
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        if self.path != "/v1/chat/completions":
            return self.reply(404, b'{"error":"unsupported route"}')
        seat = self.headers.get("X-Forwarded-Email", "").lower()
        if not seat.endswith("@team.example"):
            return self.reply(403, b'{"error":"unknown seat"}')
        try:
            size = int(self.headers.get("Content-Length", "0"))
            if not 0 < size <= 16384:
                raise ValueError()
            obj = json.loads(self.rfile.read(size))
            prompt = obj["prompt"]
            if not isinstance(prompt, str) or len(prompt.encode()) > 8192:
                raise ValueError()
        except Exception:
            return self.reply(400, b'{"error":"invalid bounded prompt"}')

        day = dt.datetime.now(dt.timezone.utc).date().isoformat()
        buckets = [
            (f"seat:{seat}:{day}", 500),
            (f"seat:{seat}:{day[:7]}", 3000),
            (f"team:{day}", 5000),
            (f"team:{day[:7]}", 30000),
        ]
        with sqlite3.connect("budget.db", timeout=10) as db:
            db.execute("BEGIN IMMEDIATE")
            for key, limit in buckets:
                db.execute("INSERT OR IGNORE INTO spend VALUES (?,0)", (key,))
                used = db.execute(
                    "SELECT cents FROM spend WHERE bucket=?", (key,)
                ).fetchone()[0]
                if used + 15 > limit:
                    db.rollback()
                    return self.reply(429, b'{"error":"budget exhausted"}')
            for key, _ in buckets:
                db.execute("UPDATE spend SET cents=cents+15 WHERE bucket=?", (key,))
            total = db.execute(
                "SELECT cents FROM spend WHERE bucket=?", (f"team:{day}",)
            ).fetchone()[0]
        if total >= 4500 and total - 15 < 4500:
            alert(day, total)

        payload = {
            "model": "gpt-4o-mini",
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 1024, "stream": False
        }
        try:
            req = urllib.request.Request(
                base_url + "/chat/completions",
                data=json.dumps(payload).encode(),
                headers={"Authorization": "Bearer " + os.environ["RELAY_KEY"],
                         "Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=60) as response:
                result = response.read()
            self.reply(200, result)
        except Exception:
            self.reply(502, b'{"error":"upstream failure; allowance retained"}')

ThreadingHTTPServer(("127.0.0.1", 9000), Gateway).serve_forever()
PY

python3 budget.py &
gateway_pid=$!
trap 'kill "$gateway_pid" 2>/dev/null || true' EXIT

oauth2-proxy \
  --provider=oidc \
  --oidc-issuer-url="$OIDC_ISSUER_URL" \
  --email-domain=team.example \
  --redirect-url=https://ai.team.example:8443/oauth2/callback \
  --https-address=0.0.0.0:8443 \
  --tls-cert-file=server.crt --tls-key-file=server.key \
  --cookie-secure=true \
  --upstream=http://127.0.0.1:9000/ \
  --pass-user-headers=true \
  --request-logging=false
Enter fullscreen mode Exit fullscreen mode

Clients submit {"prompt":"..."} after establishing an authenticated session. This is intentionally not a transparent OpenAI-compatible interface: rejecting arbitrary messages, tools, models, and streaming keeps the admission calculation bounded.

Where Production Gets Less Convenient

Keep oauth2-proxy's default stripping of incoming authentication headers enabled. Verify in integration tests that a forged X-Forwarded-Email cannot override the authenticated identity. Never expose port 9000; authentication disappears if callers bypass the proxy.

The $45 webhook warns before the $50 daily team ceiling. Alert delivery is best-effort, while enforcement is transactional. A failed webhook must not become permission to spend. Export reservation totals separately for durable monitoring.

SQLite works for one gateway instance on persistent local storage. Multiple replicas with independent databases multiply the allowance. Moving to replicas means moving admission into a shared transactional datastore, not copying this file.

Timeouts are deliberately charged. The relay may finish generation after the connection fails, so automatically releasing reservations would reopen overspending. Blind client retries consume additional allowance.

Finally, disabling access logs is not zero retention. Protect session cookies, database backups, crash dumps, and host diagnostics; separately verify relay and model-provider retention controls. This configuration avoids persisting code locally, but cannot promise what an upstream service retains.

oauth2-proxy solves the identity boundary cleanly. Keeping quota accounting outside it makes the spending decision explicit, testable, and much harder to bypass accidentally.

The hardest operational trade-off is deliberate false rejection: retain reservations after ambiguous failures and charge retries, or refund them and risk reopening the overspend window. I would rather return an honest 429 or 502 than silently exceed a budget. How do you reconcile that trade-off without turning every incident into an accounting dispute?

devops #cloud #api #productivity

Top comments (0)