Implement OAuth2/OIDC authentication for LLM APIs

Introduction

When you expose LLM API endpoints to external clients, securing them with a shared API key quickly becomes a liability — one leaked key grants unlimited access to every client in your system. OAuth2 Client Credentials flow solves this by issuing short-lived, cryptographically signed JWTs per client, each carrying its own scopes, rate-limit tier, and identity claims. By the end of this lesson, you'll be able to build a FastAPI JWT validation middleware backed by cached JWKS keys, wire it into a LiteLLM gateway, and emit asynchronous audit records that capture client identity, model requested, token consumption, authentication method, and request outcome.

Key Terminology

• OAuth2 Client Credentials Flow — The machine-to-machine OAuth2 grant type that issues short-lived, cryptographically signed JWTs per API client; each token carries the client's identity in the sub claim, its permitted scopes, and any custom claims such as rate-limit tier, replacing shared API keys with per-client credentials.
• JWKS (JSON Web Key Set) — The set of public RSA keys published by an OIDC provider at /.well-known/jwks.json; JWTValidator._refresh_jwks fetches this document and caches it so that RS256 signature verification does not require a network round-trip on every request.
• Cache-aside JWKS refresh — A caching strategy in which _refresh_jwks returns _jwks_cache immediately when time.time() < _cache_expiry, and only re-fetches from jwks_uri when the TTL (default 3600 s) has elapsed, trading full key freshness for sub-millisecond validation latency.
• FastAPI callable dependency — A class whose async __call__ method accepts Security() or Depends() arguments, enabling it to be passed directly to Depends(); JWTValidator.__call__ extracts the Bearer token, validates it, and injects the decoded payload dict into every protected route handler without additional boilerplate.
• AuditEvent — A structured dataclass capturing client_id (from the JWT sub claim), model, prompt_tokens, completion_tokens, auth_method ("jwt" or "api_key"), and outcome ("success", "auth_failed", or "rate_limited") for each authenticated LLM request.
• Non-blocking audit emission — The pattern of calling await queue.put(event) to place an AuditEvent on an asyncio.Queue instead of writing synchronously to a database; the route handler returns to the client immediately while a background coroutine drains the queue into the observability pipeline.

Concepts

JWT Validation as a Reusable FastAPI Dependency

When FastAPI evaluates a route that declares Depends(llm_auth), it calls llm_auth.__call__, which uses Security(security_scheme) to extract the raw Bearer string from the Authorization header. Because __call__ is an async method, FastAPI treats any instantiated JWTValidator as a first-class injectable — no wrapper function required. The decoded payload dict it returns becomes a named argument in the route handler, giving the handler immediate access to sub, scope, and custom claims without extra parsing logic. This is why a single llm_auth = JWTValidator(issuer=..., audience=...) at module scope is enough to protect any number of routes.

The validation step itself has three sequential responsibilities: fetch the OIDC provider's public keys, use them to verify the token's RS256 signature and standard claims, and reject anything that fails. jwt.decode checks signature, expiration (exp), audience (aud), and issuer (iss) atomically — if any check fails, python-jose raises JWTError, which validate_token converts into an HTTP 401 before the route body executes (see Code Walkthrough).

Cache-Aside JWKS Fetching and the Freshness Tradeoff

Fetching /.well-known/jwks.json on every request would add a full network round-trip to every protected endpoint. _refresh_jwks avoids this with a cache-aside check: it compares time.time() against _cache_expiry and only issues an HTTP GET when the in-process cache is stale. The default TTL of 3600 seconds sits well inside the grace period that OIDC providers maintain when rotating signing keys — the old key stays valid long enough that a one-hour cache window will not cause spurious 401 errors during rotation. The tradeoff is deliberate: consistent sub-millisecond validation latency is worth accepting up to one hour of key-rotation lag in typical production environments.

Decoupling Audit Records from the Response Path

Every authenticated LLM request must leave a record of what happened — which client connected, which model they requested, how many tokens were consumed, and whether authentication succeeded. Writing this record synchronously inside the route handler (a direct database insert) blocks the HTTP response until the write completes, adding 5–20 ms of latency on every call.

The solution is to place an AuditEvent on an asyncio.Queue and return to the client immediately. emit_audit_event does exactly this: await queue.put(event) yields control back to the event loop rather than waiting on any I/O. A separate background coroutine drains the queue and forwards records to the observability pipeline — a SIEM, a time-series store, or a dedicated LLM observability platform. The six fields in AuditEvent map directly to the four audit requirements introduced in the lesson: client_id satisfies client identity (extracted from payload["sub"]), model satisfies model tracking, prompt_tokens and completion_tokens satisfy token consumption attribution, and outcome satisfies request result. The auth_method field adds a sixth dimension — distinguishing "jwt" from "api_key" traffic — which is useful for tracking the migration progress from legacy shared keys to per-client OAuth2 tokens (see Code Walkthrough).

Code Walkthrough

Now that you understand how per-client audit trails connect client identity to token consumption and authentication outcome, let's build the two components that make it work: a JWT validation dependency and an async audit emitter.

The JWTValidator class below encapsulates JWKS fetching with cache-aside logic. It fetches the OIDC provider's public keys once, holds them for a configurable TTL (3600 seconds by default), and verifies each incoming Bearer token's RS256 signature, expiration, audience, and issuer. The __call__ method makes instances usable directly as a FastAPI Depends() target, so every protected route receives the decoded JWT payload without additional boilerplate.

Code snippetpython
1import time
2import httpx
3from jose import jwt, JWTError
4from fastapi import HTTPException, Security
5from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
6from dataclasses import dataclass, field
7
8security_scheme = HTTPBearer()
9
10@dataclass
11class JWTValidator:
12    issuer: str
13    audience: str
14    jwks_uri: str = ""
15    _jwks_cache: dict = field(default_factory=dict)
16    _cache_expiry: float = 0.0
17    _cache_ttl: int = 3600
18
19    def __post_init__(self):
20        if not self.jwks_uri:
21            self.jwks_uri = f"{self.issuer}/.well-known/jwks.json"
22
23    async def _refresh_jwks(self) -> dict:
24        now = time.time()
25        if self._jwks_cache and now < self._cache_expiry:
26            return self._jwks_cache
27        async with httpx.AsyncClient(timeout=10.0) as client:
28            resp = await client.get(self.jwks_uri)
29            resp.raise_for_status()
30            self._jwks_cache = resp.json()
31            self._cache_expiry = now + self._cache_ttl
32        return self._jwks_cache
33
34    async def validate_token(self, token: str) -> dict:
35        jwks = await self._refresh_jwks()
36        try:
37            payload = jwt.decode(
38                token,
39                jwks,
40                algorithms=["RS256"],
41                audience=self.audience,
42                issuer=self.issuer,
43            )
44        except JWTError as e:
45            raise HTTPException(status_code=401, detail=f"Token invalid: {e}")
46        if payload.get("token_use") not in ("access", None):
47            raise HTTPException(status_code=401, detail="Invalid token_use")
48        return payload
49
50    async def __call__(
51        self, credentials: HTTPAuthorizationCredentials = Security(security_scheme)
52    ) -> dict:
53        return await self.validate_token(credentials.credentials)
54
55llm_auth = JWTValidator(
56    issuer="https://auth.example.com/realms/llm-platform",
57    audience="llm-api-gateway",
58)

With the validator wired in, each route receives a verified payload dict containing sub (the client ID), scope, and any custom claims your OIDC provider injects. The Concepts section establishes that every authenticated request must also produce an audit record capturing client_id, model, prompt_tokens, completion_tokens, auth_method, and outcome. The snippet below shows how to emit that record without blocking the response path — the route drops the event onto an asyncio.Queue and a separate background consumer drains it into your observability pipeline:

Code snippetpython
1import asyncio
2from dataclasses import dataclass
3
4@dataclass
5class AuditEvent:
6    client_id: str
7    model: str
8    prompt_tokens: int
9    completion_tokens: int
10    auth_method: str  # "jwt" or "api_key"
11    outcome: str      # "success", "auth_failed", or "rate_limited"
12
13async def emit_audit_event(event: AuditEvent, queue: asyncio.Queue) -> None:
14    """Non-blocking: places the event on the queue for async consumption."""
15    await queue.put(event)
16
17# Inside a FastAPI route that depends on llm_auth:
18# payload = await llm_auth(credentials)
19# await emit_audit_event(
20#     AuditEvent(
21#         client_id=payload["sub"],
22#         model="gpt-4",
23#         prompt_tokens=150,
24#         completion_tokens=80,
25#         auth_method="jwt",
26#         outcome="success",
27#     ),
28#     audit_queue,
29# )

Using asyncio.Queue keeps the write path non-blocking — the route handler returns to the client immediately while the background task persists records to your SIEM or LLM observability platform.

Confirm that a request carrying a valid JWT returns 200 OK and that a tampered or expired token returns 401 Unauthorized with a "Token invalid" detail string in the response body.

Do's and Don'ts

Now that you have worked through the implementation, the practices below separate a durable approach from a fragile one.

Do's

1. ✓Do cache JWKS keys with a TTL using _refresh_jwks's cache-aside logic — fetching jwks.json on every request adds a blocking network round-trip to your hot path and makes your gateway dependent on OIDC provider availability for every LLM call; a 3600-second in-memory cache absorbs that cost while staying fresh enough for key rotation.
2. ✓Do wire JWTValidator as a FastAPI Depends() target via its __call__ method — this lets every protected route receive the decoded payload dict (including sub, scope, and custom claims) without duplicating token extraction or error handling, and keeps RS256 signature verification, expiration, audience, and issuer checks in one testable place.
3. ✓Do emit AuditEvent records via asyncio.Queue rather than writing directly inside the route handler — placing the event on the queue with emit_audit_event returns control to the client immediately; blocking writes to a SIEM or observability platform inside the handler add latency to every LLM response and can cascade failures if the downstream sink is slow.

Don'ts

1. ✗Don't skip validating token_use, audience, and issuer inside validate_token — omitting any of these checks means a JWT issued for a different service or token type (e.g., an ID token instead of an access token) will pass jwt.decode's signature check and reach your LiteLLM gateway with full client privileges, silently bypassing your per-client scope and rate-limit tier enforcement.
2. ✗Don't hardcode a shared API key in place of the OAuth2 Client Credentials flow — a single leaked key grants unlimited access to every client in your system with no per-client sub claim, making it impossible to attribute prompt_tokens and completion_tokens to a specific client or to revoke access for one client without rotating credentials for all of them.
3. ✗Don't populate AuditEvent.client_id from anything other than payload["sub"] — using a request header or query parameter supplied by the caller lets any authenticated client impersonate another in your audit trail, destroying the integrity of the per-client token-consumption and outcome records that your observability pipeline depends on.