DEV Community

Cover image for Designing an AI Gateway from Scratch: Routing, Failover, and Cost Control
mihir mohapatra
mihir mohapatra

Posted on

Designing an AI Gateway from Scratch: Routing, Failover, and Cost Control

Why You Need an AI Gateway

The first LLM feature ships fast: one API key, one provider, one SDK call. Six months later you have:

  • 4 services calling 3 different providers
  • API keys scattered across repos
  • No idea which team is burning the budget
  • An outage at one provider that took down your product

An AI gateway is a single entry point between your applications and LLM providers. It's the same idea as an API gateway, but it handles problems specific to LLMs: token-based billing, slow streaming responses, flaky providers, and non-deterministic outputs.

In this post we'll design one from scratch.

Requirements

Functional

  • One unified API (OpenAI-compatible) for all providers
  • Route requests across providers and models
  • Automatic retries and failover
  • Per-team API keys, quotas, and rate limits
  • Response caching
  • Streaming (SSE) support

Non-functional

  • Added latency overhead under 10 ms (p99)
  • No single point of failure
  • Full visibility into cost, latency, and errors per request
  • Provider keys never exposed to clients

High-Level Architecture

            ┌──────────────┐
 Clients ──▶│  Load Balancer│
            └──────┬───────┘
                   ▼
        ┌─────────────────────┐
        │   Gateway (stateless)│
        │  ┌───────────────┐  │
        │  │ Auth + Quotas │  │
        │  ├───────────────┤  │
        │  │ Cache lookup  │  │
        │  ├───────────────┤  │
        │  │ Router        │  │
        │  ├───────────────┤  │
        │  │ Provider      │  │
        │  │ adapters      │  │
        │  └───────────────┘  │
        └───┬─────┬─────┬─────┘
            ▼     ▼     ▼
        OpenAI Anthropic  Self-hosted
                          (vLLM)

   Side channels: Redis (limits, cache) · Postgres (keys, config)
                  Kafka/queue (usage events) · Metrics + traces
Enter fullscreen mode Exit fullscreen mode

The key decision: keep the gateway stateless. All state (rate-limit counters, cache, key metadata) lives in Redis or Postgres, so you can scale gateway nodes horizontally behind a load balancer.

Core Components

1. Unified API Layer

Pick one schema and translate everything to it. OpenAI's chat completions format is the de facto standard, so clients can use existing SDKs by changing only the base URL.

Each provider gets an adapter that implements the same trait:

#[async_trait]
pub trait Provider: Send + Sync {
    fn name(&self) -> &str;

    async fn chat(&self, req: ChatRequest) -> Result<ChatResponse, ProviderError>;

    async fn chat_stream(
        &self,
        req: ChatRequest,
    ) -> Result<BoxStream<'static, Result<ChatChunk, ProviderError>>, ProviderError>;
}
Enter fullscreen mode Exit fullscreen mode

Adding a new provider means writing one adapter. Nothing else changes.

2. Authentication and Quotas

Clients get virtual keys issued by the gateway. Real provider keys stay server-side.

A virtual key maps to:

  • the owning team or project
  • allowed models
  • rate limits (requests/min and tokens/min)
  • a monthly budget

Store hashed keys in Postgres and cache hot keys in memory with a short TTL.

3. Rate Limiting

LLM traffic needs two limits, not one:

Limit Why
Requests per minute Protects against bursts and abuse
Tokens per minute Matches how providers actually bill and throttle

A token bucket in Redis, updated atomically with a Lua script, works well. Since output tokens aren't known up front, estimate on the way in, then reconcile with actual usage when the response completes.

4. The Router

The router picks a provider and model for each request. Strategies, from simple to advanced:

  1. Static mapping: gpt-4o goes to OpenAI, claude-sonnet goes to Anthropic.
  2. Weighted: split traffic 80/20 across providers for the same model class.
  3. Latency-based: prefer the provider with the lowest recent p50.
  4. Cost-based: send low-complexity prompts to a cheaper model.
  5. Fallback chains: try A, then B, then C.

A config-driven fallback chain is a good place to start:

routes:
  - name: default-chat
    strategy: fallback
    targets:
      - provider: openai
        model: gpt-4o
        timeout_ms: 20000
      - provider: anthropic
        model: claude-sonnet
        timeout_ms: 20000
      - provider: self-hosted
        model: llama-3-70b
        timeout_ms: 30000
Enter fullscreen mode Exit fullscreen mode

5. Retries, Timeouts, and Circuit Breakers

Failure handling is where a gateway earns its keep.

  • Retry only what's safe: 429s, 5xx, and connection errors. Never retry 4xx validation errors.
  • Exponential backoff with jitter to avoid thundering herds.
  • Respect Retry-After headers from providers.
  • Circuit breaker per provider: after N failures in a window, stop sending traffic for a cool-down period, then probe with a few requests.
  • Streaming caveat: once you've sent the first token to the client, you can't transparently fail over. Fail over only before the first byte.
closed ──(failures > threshold)──▶ open
  ▲                                  │
  │                          (cool-down elapsed)
  └──(probe succeeds)── half-open ◀──┘
Enter fullscreen mode Exit fullscreen mode

6. Caching

Two layers:

  • Exact-match cache: hash of (model, messages, params). Cheap, safe, and effective for repeated prompts. Only cache when temperature is 0 or the caller opts in.
  • Semantic cache: embed the prompt and look up near neighbors in a vector store. Higher hit rate, but risk of returning a subtly wrong answer. Use a strict similarity threshold and make it opt-in per route.

7. Observability

Every request should emit one structured event:

{
  "request_id": "req_8f2c",
  "team": "search",
  "provider": "anthropic",
  "model": "claude-sonnet",
  "latency_ms": 1840,
  "ttft_ms": 420,
  "input_tokens": 912,
  "output_tokens": 304,
  "cost_usd": 0.0071,
  "cache": "miss",
  "retries": 1,
  "status": 200
}
Enter fullscreen mode Exit fullscreen mode

Push these to a queue and process them asynchronously. Never block the response path on logging or billing writes. Track:

  • TTFT (time to first token), which is what users actually feel
  • Cost per team, model, and route
  • Error rate and fallback rate per provider

Handling Streaming

Streaming adds real complexity:

  • Proxy SSE chunks as they arrive. Don't buffer the full response.
  • Normalize each provider's chunk format into your unified schema.
  • Count tokens as they stream so quotas stay accurate.
  • Handle client disconnects by cancelling the upstream request. Otherwise you pay for tokens nobody reads.

Scaling and Reliability

  • Stateless nodes behind a load balancer; scale on CPU and open connections.
  • Redis for counters and cache. Run it replicated, and fail open on limits if Redis is down (a deliberate trade-off: availability over strict enforcement).
  • Multi-region: deploy gateways close to your apps. Keep config in Postgres with read replicas.
  • Config hot-reload: change routes without redeploying.
  • Backpressure: cap concurrent upstream requests per provider and shed load gracefully with a clear 429.

Trade-offs to Call Out

Decision Upside Downside
Single OpenAI-compatible schema Easy adoption Provider-specific features get lost or need passthrough
Gateway in the hot path Central control Extra hop and a critical dependency
Semantic caching Big cost savings Risk of wrong or stale answers
Fail-open rate limits High availability Brief over-consumption possible

Build vs. Buy

Existing options (LiteLLM, Portkey, Cloudflare AI Gateway, and others) cover a lot. Build your own when you need tight control over routing logic, strict data-residency rules, or very low overhead. Otherwise, start with an existing tool and replace pieces as you hit limits.

Wrapping Up

An AI gateway comes down to a few ideas:

  1. A unified API with provider adapters
  2. Virtual keys, quotas, and token-aware rate limits
  3. Routing with fallbacks and circuit breakers
  4. Caching where it's safe
  5. Per-request observability for cost and latency

Start with a static router and exact-match cache, then add sophistication as real traffic shows you where the pain is.

Next in the series: RAG Pipelines in Production. Follow along so you don't miss it.

What does your gateway setup look like? Are you routing by cost, latency, or something else? Let me know in the comments.

Top comments (0)