Skip to main content
Jodei
Insights

AI SDK Middleware: Logging, Retries, and Rate-Limiting Without Boilerplate

The Vercel AI SDK's middleware API lets you wrap every LLM call with cross-cutting concerns — token logging, exponential-backoff retries, per-user rate limits — in one place, without touching your route handlers.

6 min read

Every production AI feature eventually needs the same things: log what the model was asked and how many tokens it used, retry on transient API errors, and prevent any single user from hammering your LLM quota. Without a structured approach you end up copy-pasting the same try/catch and rate-limit logic across every route handler. The Vercel AI SDK's middleware API solves this with a composable wrapper that intercepts every call to generateText, streamText, or generateObject in one place.

What middleware does

An AI SDK middleware is an object with optional hooks: wrapGenerate for non-streaming calls and wrapStream for streaming ones. Each hook receives the original parameters and a doGenerate / doStream function you call to run the underlying model. That lets you run code before the call (validate, rate-limit, log the prompt), after (log the response, record usage), or instead of it (return a cached result).

You attach middleware to a model with wrapLanguageModel, then use the wrapped model everywhere else exactly as you would the original.

import { wrapLanguageModel, type LanguageModelMiddleware } from "ai";
import { openai } from "@ai-sdk/openai";

const loggingMiddleware: LanguageModelMiddleware = {
  wrapGenerate: async ({ doGenerate, params }) => {
    console.log("[ai] prompt tokens (est.):", params.prompt.length);
    const result = await doGenerate();
    console.log("[ai] finish reason:", result.finishReason);
    console.log("[ai] usage:", result.usage);
    return result;
  },
};

export const model = wrapLanguageModel({
  model: openai("gpt-4o"),
  middleware: loggingMiddleware,
});

Adding exponential-backoff retries

Rate-limit errors (429) and transient network failures are common enough in production that every AI route needs to handle them. Doing this in middleware means every caller gets retries for free — you write it once.

const retryMiddleware: LanguageModelMiddleware = {
  wrapGenerate: async ({ doGenerate }) => {
    const MAX_ATTEMPTS = 3;
    let attempt = 0;
    while (true) {
      try {
        return await doGenerate();
      } catch (err: unknown) {
        attempt++;
        const isRateLimit =
          err instanceof Error && err.message.includes("429");
        if (!isRateLimit || attempt >= MAX_ATTEMPTS) throw err;
        // exponential back-off: 1s, 2s, 4s …
        await new Promise((r) => setTimeout(r, 2 ** (attempt - 1) * 1000));
      }
    }
  },
};

Keep MAX_ATTEMPTS low — three is usually enough. Retrying non-rate-limit errors (like a malformed prompt) wastes time and tokens, so check the error type before retrying.

Per-user rate limiting

Client-side rate limiting on the model itself is the last line of defence before you incur cost. A simple in-memory map works for single-process deployments; for serverless or multi-region, replace the Map with a Redis counter.

const requestCounts = new Map<string, { count: number; reset: number }>();

function makeRateLimitMiddleware(
  userId: string,
  maxPerMinute = 10,
): LanguageModelMiddleware {
  return {
    wrapGenerate: async ({ doGenerate }) => {
      const now = Date.now();
      const window = requestCounts.get(userId) ?? { count: 0, reset: now + 60_000 };
      if (now > window.reset) {
        window.count = 0;
        window.reset = now + 60_000;
      }
      if (window.count >= maxPerMinute) {
        throw new Error("Rate limit exceeded. Try again in a minute.");
      }
      window.count++;
      requestCounts.set(userId, window);
      return doGenerate();
    },
  };
}

Construct this middleware inside your Server Function or Route Handler where you have access to the authenticated user's ID, then pass it to wrapLanguageModel for that request only.

Composing multiple middlewares

wrapLanguageModel accepts a single middleware object, so combine them into one before passing. A simple compose helper applies them inside-out — the first middleware in the array is the outermost wrapper:

function composeMiddleware(
  ...middlewares: LanguageModelMiddleware[]
): LanguageModelMiddleware {
  return {
    wrapGenerate: async ({ doGenerate, params }) => {
      const composed = middlewares.reduceRight<() => Promise<Awaited<ReturnType<typeof doGenerate>>>>(
        (next, mw) => () =>
          mw.wrapGenerate
            ? mw.wrapGenerate({ doGenerate: next, params })
            : next(),
        doGenerate,
      );
      return composed();
    },
  };
}

// Usage
export function createModel(userId: string) {
  return wrapLanguageModel({
    model: openai("gpt-4o"),
    middleware: composeMiddleware(
      loggingMiddleware,
      retryMiddleware,
      makeRateLimitMiddleware(userId),
    ),
  });
}

Practical wiring in a Route Handler

Call createModel once per request with the authenticated user ID. The rest of the route handler is unchanged — it still calls streamText with the same API surface, and all three cross-cutting behaviours are applied automatically.

// app/api/chat/route.ts
import { streamText } from "ai";
import { createModel } from "@/lib/ai/model";
import { auth } from "@/lib/auth";

export async function POST(req: Request) {
  const { userId } = await auth();
  const { messages } = await req.json();

  const result = await streamText({
    model: createModel(userId),
    messages,
  });

  return result.toDataStreamResponse();
}

What to watch out for

  • wrapStream is separate. If you use streamText, implement wrapStream alongside wrapGenerate. A middleware that only implements wrapGenerate is silently bypassed for streaming calls — log the usage event in wrapStream too, once the stream closes.
  • Serverless resets the Map. The in-memory rate-limit counter above resets on every cold start. For serverless deployments, push counts to an edge KV store like Upstash Redis — the API surface is similar and the overhead is one extra await.
  • Middleware errors propagate as-is. If your rate-limit middleware throws, that error reaches the caller unmodified. Return a well-typed Response from the route handler instead of letting it bubble as an unhandled exception.

The right level of abstraction

Middleware sits at the model level, not the route level — that is the right place for concerns that apply to every LLM call regardless of what the call does. Logging, retries, and rate limiting all belong there. Business logic (which tools to expose, what system prompt to use, how to parse the result) still lives in the route handler or Server Function where it has the context it needs. Keep the two layers separate and your AI codebase scales without accumulating per-route boilerplate.