AI gateway sounds like a router: something that picks a model, maybe falls back to a backup when one is down, and gets out of the way. That's the least interesting thing a gateway does.
The actual job is governance. Cost control, security, and audit. Routing is one line item on that list, not the headline.
This post walks through all four and shows the exact config behind each.
Five services, five ways to get it wrong
Picture a company with five things calling LLMs directly: a support bot, a RAG pipeline, an internal LangGraph agent, a nightly batch summarizer, and a customer-facing chat feature. No gateway. Each one talks to the provider on its own.
Here's what that actually means:
- Every service carries its own copy of the provider API key. A leak in the batch summarizer is a leak of the entire budget.
- Rate limiting and retries are whatever each team got around to. Usually that means "none" until the first incident.
- Nobody has a live view of spend. Finance learns about a runaway agent when the invoice arrives, not when the loop starts burning tokens at 2am.
- PII scanning is per-team folklore. Some scrub, some don't, nobody agrees on what counts.
- Logging is five different formats, if it exists at all.
You don't have one problem five times. You have five slightly different, half-finished versions of the same five problems. Put them all in front of an auditor and none of them hold up.
Collapse it into one control plane
One gateway. Every service talks to it with an OpenAI-compatible API, and every rule gets enforced in that one box before a request ever reaches a provider. The four controls below all live inside it.
I used LiteLLM as the proxy. The theme running through all four controls: don't rebuild what the platform already does well, just wire it correctly and make it enforce.
Control 1: cost governance that counts tokens, not requests
A normal API gateway counts requests. For LLMs that's the wrong unit, and it's wrong in a way that hurts. A ten-word prompt and a ten-thousand-token agent trajectory both register as "one request" while costing wildly different amounts. If you want to govern spend, you track tokens.
And you track them per identity, not as one global number. Per API key, per team, per customer. A single global counter tells you the ship is sinking but not which compartment.
So the design is hierarchical budgets:
- A team budget. This is the unit that gets enforced. Every user belongs to a team.
- Rolled up into an org / customer budget above it.
- The request gets rejected before it hits the provider if either level is over.
Users still exist as identities, so spend is attributed to a person (the audit trail in Control 4 needs that). But there's no enforced budget on the user tier, and that's on purpose. LiteLLM only enforces a user's own budget when that user has no team, so for anyone on a team it does nothing. The hierarchy that actually bites is team then org. The user is for attribution.
Budgets reset on a window you pick: daily, weekly, monthly. This is the control that keeps one runaway agent loop from becoming a five-figure surprise. If any of your traffic is agentic rather than one-shot completions, that's the failure mode that will eventually find you.
Two planes, and LiteLLM owns the money
I didn't rebuild token tracking. Per-model pricing and pre-flight spend accounting is the one thing LiteLLM genuinely does well, so the system splits into two planes:
-
Identity plane (mine). Cognito plus an
authservice. Source of truth for who the caller is:Org → Team → User, in Postgres. -
Cost plane (LiteLLM). Budgets, spend, and reject-before-spend all live in LiteLLM's own DB, on its
OrganizationandTeamrecords. I never store a dollar or a token myself.
The two planes meet at two points: provisioning and request time.
Provisioning. Seed the budgeted Org and Team records in LiteLLM with a max_budget and a reset window, the team linked to its org. The key detail is that the IDs are the auth DB's own primary keys, not names. Names collide and aren't stable, so they're never the join key.
# provision.py — mirror each Org/Team into LiteLLM with a budget.
requests.post(f"{LITELLM_URL}/organization/new", headers=_headers, json={
"organization_id": str(org.id), # our Postgres PK, stringified
"organization_alias": org.name, # display name only, not a key
"max_budget": ORG_BUDGET, # e.g. 500
"budget_duration": BUDGET_DURATION, # e.g. "30d"
})
requests.post(f"{LITELLM_URL}/team/new", headers=_headers, json={
"team_id": str(team.id),
"team_alias": team.name,
"organization_id": str(team.org_id), # links team -> org so the org budget rolls up
"max_budget": TEAM_BUDGET, # e.g. 100
"budget_duration": BUDGET_DURATION,
})
Full file:
services/gateway/provision.py
Request time. Callers send their Cognito JWT to the gateway. LiteLLM runs a custom_auth hook that verifies the JWT and resolves it to user_id / team_id / org_id, using the same IDs provisioning used. LiteLLM then checks the team and org budgets and rejects before calling the provider if either is over.
# custom_auth.py — the last few lines of the hook.
# Verify the Cognito access token, resolve the sub to our team/org PKs, and
# hand LiteLLM an identity. The IDs match what provision.py created, so the
# budget lookup lands on the right records.
claims = _verify_access_token(api_key)
team_id, org_id = _resolve_identity(claims["sub"])
return UserAPIKeyAuth(
api_key=api_key,
user_id=claims["sub"], # Cognito sub
team_id=team_id, # our Postgres PK
org_id=org_id,
)
Full file:
services/gateway/custom_auth.py
Wiring the hook in is two lines of proxy config. The second one matters more than it looks:
# config.yaml
general_settings:
custom_auth: custom_auth.user_api_key_auth
# Without this, the hook authenticates but the budgets are never checked.
custom_auth_run_common_checks: true
Full file:
services/gateway/config.yaml
One question worth answering directly: why a custom hook instead of LiteLLM's JWT-to-virtual-key mapping? Because the native mapping is an Enterprise feature. The open-source custom_auth hook does the same job, turning a Cognito identity into a governed LiteLLM identity, in about twenty lines you write yourself.
The demo shrinks a team budget, logs in a real user, and loops until the gateway answers HTTP 429 Budget has been exceeded! Team=1. That 429 comes before any provider call, which is the whole point.
Control 2: security at the edge
An API gateway checks who is calling. An AI gateway also has to check what they're sending, and that second check is the one people skip.
A prompt is free text, and services stuff whatever they have into it. The support bot pastes in the customer's email. The agent quotes a tool result that happens to contain a card number. Send that as-is and the PII lands in a third-party provider's logs. That's a data-egress incident, and the uncomfortable part is that nobody did anything obviously wrong to cause it.
So security at the edge is two layers. The who is already handled by Control 1: the same custom_auth hook rejects any request without a valid JWT. The new part is the what: inspect the prompt's content before it leaves the network, and mask PII pre-flight.
Presidio as a sidecar, masking before the provider call
I didn't write a PII scanner either. LiteLLM ships a guardrail that integrates Microsoft Presidio, so I stood Presidio up as two sidecars and wired it in. The config is small, and every knob in it is a decision:
# config.yaml
guardrails:
- guardrail_name: "presidio-pii"
litellm_params:
guardrail: presidio
mode: "pre_call" # scrub the request BEFORE the provider is called
default_on: true # every request, not opt-in — opt-in isn't enforcement
presidio_language: "en"
pii_entities_config:
PERSON: "MASK" # names -> <PERSON>
EMAIL_ADDRESS: "MASK" # emails -> <EMAIL_ADDRESS>
CREDIT_CARD: "MASK" # cards -> <CREDIT_CARD>
Full file:
services/gateway/config.yaml
Three things worth calling out:
- Mask, don't block. Detected entities are replaced with typed placeholders and the now-clean request continues. The caller still gets a useful answer, the provider never sees the raw PII.
-
default_on: true. The scan runs on every request. An opt-in guardrail is a suggestion, not a control. -
Fail-closed, for free. Because
pii_entities_configis set, LiteLLM's Presidio integration raises if the analyzer is unreachable or errors. A degraded scanner rejects traffic instead of forwarding it unscrubbed, so a misconfiguration or an outage can't quietly leak PII.
The two Presidio containers run right next to the gateway in the same compose file:
# docker-compose.yaml
presidio-analyzer: # detects entities with an NLP model (slow first boot)
image: mcr.microsoft.com/presidio-analyzer:latest
ports: ["5002:3000"]
presidio-anonymizer: # applies the masking
image: mcr.microsoft.com/presidio-anonymizer:latest
ports: ["5001:3000"]
Full file:
docker-compose.yaml
The demo sends a prompt full of PII and asks the model to echo it back verbatim, so you can see exactly what left the network: only the placeholders come back. Stop the analyzer and re-run, and the request is rejected instead of leaked.
Control 3: routing and failover
The caller shouldn't have to know which model to use, and shouldn't ever see a provider outage. Both are the gateway's job, and they sit at opposite ends of a request's timeline.
- Routing is a cost decision made before the call. Not every prompt needs the expensive model. "What's the capital of France?" and a distributed-systems design question both look like one request, but sending the first to gpt-4o is money set on fire.
- Failover is a reliability decision made after a call fails. A provider throws a 500 or times out. Instead of handing that to the caller, the gateway retries on a healthy backup.
Unlike the first two controls, there's no hook code here. LiteLLM does both natively, so Control 3 is configuration. The work was in picking a mechanism that isn't smoke and mirrors.
Routing: the complexity router
LiteLLM ships two content-based routers. A semantic one embeds every prompt and matches it against example utterances, which costs an embedding call per request. A complexity one scores keywords, code, and reasoning markers with rules, in under a millisecond, no extra call. I used the complexity router. It's free per request and deterministic, and paying an embedding call to decide how to save money undercuts the whole point.
Callers hit one logical model, smart-router, and the gateway maps the prompt to a cost tier:
# config.yaml
- model_name: smart-router
litellm_params:
model: auto_router/complexity_router
complexity_router_config:
tiers:
SIMPLE: gpt-4o-mini # greetings, definitions -> cheapest
MEDIUM: gpt-4o-mini # standard queries -> still cheap
COMPLEX: gpt-4o # technical, multi-part -> pricier
REASONING: o3-mini # chain-of-thought -> reasoning model
default_model: gpt-4o-mini
Failover
LiteLLM's fallbacks list reroutes a failed request to a backup. I wired it so the demo can't cheat: a deployment named flaky-provider points its endpoint at a closed port, so it always fails with a connection error, standing in for a provider outage.
# config.yaml
- model_name: flaky-provider
litellm_params:
model: openai/gpt-4o-mini
api_key: os.environ/OPENAI_API_KEY
api_base: http://localhost:1 # nothing listening -> connection refused
litellm_settings:
fallbacks: [{"flaky-provider": ["gpt-4o-mini"]}]
num_retries: 2
request_timeout: 30
Full file:
services/gateway/config.yaml
A call to flaky-provider fails fast, the fallback sends it to a healthy gpt-4o-mini, and the caller gets a normal 200.
One thing I deliberately left out: cascading. Cascading means calling the cheap model, inspecting its answer, and escalating if it isn't good enough. It gates on the response, so no LiteLLM router does it, because they all decide from the prompt before the call. It's a real technique, but it's a different shape (a post_call hook with a quality gate) and belongs in its own control, not smuggled into this one.
The demo sends three prompts through one smart-router endpoint, then resolves the x-litellm-model-id response header to prove they landed on three different models, and finishes by surviving the flaky-provider outage.
Control 4: the audit trail
The first three controls each make a decision per request: who you are, whether the prompt is safe to send, which model runs it, whether you're over budget. To a security or finance team, none of that is worth much without a record. The audit trail is the record. For every request: who called, what model ran, how many tokens, what it cost, how long it took, and the field everyone forgets, whether it was allowed or rejected.
Trace to Langfuse, success and failure, tagged with identity
No logging pipeline here either. LiteLLM traces every call to Langfuse natively: model, tokens, cost, latency, input, and output, out of the box. Three choices turn that from observability into an audit control, and all three are in this block:
# config.yaml
litellm_settings:
success_callback: ["langfuse"] # completed calls
failure_callback: ["langfuse"] # rejections and errors — the events that matter most
langfuse_default_tags:
- user_api_key_user_id # Cognito sub
- user_api_key_team_id # our Postgres PKs, straight from custom_auth
- user_api_key_org_id
Full file:
services/gateway/config.yaml
- Both callbacks. An audit trail that only records successes hides the exact events that matter: the over-budget 429 from Control 1, a provider failure. Send both.
-
Identity tags. Out of the box a trace tells you what ran and what it cost, but not who. The tags pull the identity resolved by
custom_authin Control 1 straight onto every trace, server-side, so the caller can't forge them. Now you can filter the whole trail to one team and see everything it spent. - PII already handled. Control 2 masks PII before the provider call, and the trace stores that same scrubbed text. The audit log doesn't reintroduce the leak by parking raw PII in a third-party tool.
One thing this is not: a SIEM. Langfuse is observability. What you get is a queryable, per-request trail you can filter and export, a genuine first cut of an audit story, not a compliance-grade immutable log pipeline. Worth saying plainly so nobody oversells it in a meeting.
The demo generates one success and one failure, both traced and tagged with the caller. Open the Traces view, filter on user_api_key_team_id:1, and you see everything that team spent, rejections included.
Wrapping up
That's the four controls. Identity and cost, content security, routing and failover, and the audit trail. Each one a working LiteLLM configuration with a demo, and together they collapse five inconsistent per-service versions of these concerns into one control plane.
The through-line is simple: an AI gateway earns its place by governing, not routing. Routing is the easy 10%. The budgets, the PII masking, the identity-tagged audit trail — that's the part that keeps you out of the incident channel and in front of the auditor with something that holds up.
Full source, config, and demos: github.com/mohsinsheikhani/ai-gateway-governance. If you'd wire any of it differently, I'd genuinely like to hear it in the comments.











Top comments (0)