How to rate limit an MCP server (and handle 429s the right way)

// token-bucket.ts
// A per-client token bucket. Burst = capacity, sustained = refillPerSec.
export class TokenBucket {
  private tokens: number
  private lastRefill: number

  constructor(private capacity: number, private refillPerSec: number) {
    this.tokens = capacity
    this.lastRefill = Date.now()
  }

  take(cost = 1): { ok: boolean; retryAfterMs: number } {
    const now = Date.now()
    const refill = ((now - this.lastRefill) / 1000) * this.refillPerSec
    this.tokens = Math.min(this.capacity, this.tokens + refill)
    this.lastRefill = now

    if (this.tokens >= cost) {
      this.tokens -= cost
      return { ok: true, retryAfterMs: 0 }
    }
    const deficit = cost - this.tokens
    return { ok: false, retryAfterMs: Math.ceil((deficit / this.refillPerSec) * 1000) }
  }
}

const buckets = new Map<string, TokenBucket>()

// 20-request burst, 5 requests/second sustained, per client.
export function bucketFor(clientId: string): TokenBucket {
  let b = buckets.get(clientId)
  if (!b) {
    b = new TokenBucket(20, 5)
    buckets.set(clientId, b)
  }
  return b
}

// server.ts
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import { z } from "zod"
import { bucketFor } from "./token-bucket.js"

const server = new McpServer({ name: "orders", version: "1.0.0" })

server.registerTool(
  "search_orders",
  {
    description: "Search orders by customer email.",
    inputSchema: { email: z.string().email() },
  },
  async ({ email }, extra) => {
    // Prefer the authenticated client from validated auth; fall back to a header.
    const clientId =
      extra.authInfo?.clientId ??
      extra.requestInfo?.headers["mcp-client-id"]?.toString() ??
      "anonymous"

    const { ok, retryAfterMs } = bucketFor(clientId).take()
    if (!ok) {
      const seconds = Math.ceil(retryAfterMs / 1000)
      return {
        isError: true,
        content: [
          {
            type: "text",
            text: `Rate limited. Retry after ${seconds}s. This server allows 5 requests/second per client, with a burst of 20.`,
          },
        ],
      }
    }

    const orders = await findOrders(email) // your real work
    return { content: [{ type: "text", text: JSON.stringify(orders) }] }
  },
)

// sliding-window.ts
export class SlidingWindow {
  private hits: number[] = []

  constructor(private limit: number, private windowMs: number) {}

  take(): { ok: boolean; retryAfterMs: number } {
    const now = Date.now()
    const cutoff = now - this.windowMs
    // Drop timestamps that have aged out of the window.
    this.hits = this.hits.filter((t) => t > cutoff)

    if (this.hits.length < this.limit) {
      this.hits.push(now)
      return { ok: true, retryAfterMs: 0 }
    }
    // The window frees up when the oldest hit ages out.
    const retryAfterMs = this.hits[0] + this.windowMs - now
    return { ok: false, retryAfterMs }
  }
}

// 100 requests per 60 seconds, strict.
const window = new SlidingWindow(100, 60_000)

// upstream.ts
export class UpstreamRateLimit extends Error {
  constructor(public retryAfterSec: number) {
    super(`upstream rate limited; retry after ${retryAfterSec}s`)
  }
}

export async function callUpstream(url: string, token: string) {
  const res = await fetch(url, {
    headers: { authorization: `Bearer ${token}` },
  })

  if (res.status === 429) {
    const retryAfter = Number(res.headers.get("retry-after") ?? "1")
    // Do NOT retry here. Surface it so the agent can back off.
    throw new UpstreamRateLimit(retryAfter)
  }
  if (!res.ok) throw new Error(`upstream ${res.status}`)
  return res.json()
}

// In the tool handler:
//   try { return ok(await callUpstream(url, token)) }
//   catch (e) {
//     if (e instanceof UpstreamRateLimit)
//       return { isError: true, content: [{ type: "text",
//         text: `Upstream is rate limited. Retry after ${e.retryAfterSec}s.` }] }
//     throw e
//   }

npm i @modelcontextprotocol/[email protected] zod@3
npx tsx server.ts