Chapter 12 · Middleware and Testing

Chapter goals:

• Understand the positioning of Language Model Middleware: intercept and enhance language model calls in a model-agnostic way
• Master the usage of wrapLanguageModel and the execution order of multiple middlewares
• Become familiar with six built-in middlewares and three custom implementation entry points (transformParams / wrapGenerate / wrapStream)
• Learn to use Mock models and testing helpers from ai/test for deterministic unit tests
• Master techniques for simulating streaming responses with simulateReadableStream

12.1 What Is Language Model Middleware

Language model middleware is a way to enhance language model behavior — by intercepting and modifying calls to the language model. It can add guardrails, RAG, caching, logging, and other capabilities to a model, and is independent of the specific language model. Such middleware can be developed and distributed independently of the model it is applied to.

The yourModel in this chapter's examples can be either a model constructed via Vercel AI Gateway or a custom OpenAI-compatible Provider instance — middleware operates above the model abstraction layer, so both access methods are fully通用:

import { createGateway } from 'ai';

const gateway = createGateway({ apiKey: process.env.AI_GATEWAY_API_KEY ?? '' });
const yourModel = gateway('openai/gpt-5');

// Method 2: Custom OpenAI-compatible Provider (choose one)
// import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
// const yourModel = createOpenAICompatible({
//   name: 'my-provider',
//   baseURL: process.env.OPENAI_COMPATIBLE_BASE_URL ?? '',
//   apiKey: process.env.OPENAI_COMPATIBLE_API_KEY ?? '',
// })('gpt-4o-mini');

12.2 Using wrapLanguageModel

wrapLanguageModel takes a language model and a middleware, returning a new model that incorporates that middleware:

import { wrapLanguageModel, streamText } from 'ai';

const wrappedLanguageModel = wrapLanguageModel({
  model: yourModel,
  middleware: yourLanguageModelMiddleware,
});

The wrapped model can be used like any ordinary model:

const result = streamText({
  model: wrappedLanguageModel,
  prompt: 'What cities are in the United States?',
});

Multiple Middlewares

You can provide multiple middlewares to wrapLanguageModel, applied in the order provided:

const wrappedLanguageModel = wrapLanguageModel({
  model: yourModel,
  middleware: [firstMiddleware, secondMiddleware],
});

// Application order: firstMiddleware(secondMiddleware(yourModel))

12.3 Built-in Middlewares

AI SDK ships with several built-in middlewares:

Middleware	Purpose
extractReasoningMiddleware	Extract reasoning info from generated text (e.g., ``):

import { wrapLanguageModel, extractReasoningMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: extractReasoningMiddleware({ tagName: 'think' }),
});

The startWithReasoning option is also available: when set to true, the reasoning tag is prepended before generation begins, suitable for models that do not place tags at the start of the reply.

Extracting JSON

Some models wrap JSON in markdown code fences even when you request structured output. extractJsonMiddleware strips these fences to make them compatible with Output.object():

import {
  wrapLanguageModel,
  extractJsonMiddleware,
  Output,
  generateText,
} from 'ai';
import { z } from 'zod';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: extractJsonMiddleware(),
});

const result = await generateText({
  model,
  output: Output.object({
    schema: z.object({
      name: z.string(),
      ingredients: z.array(z.string()),
    }),
  }),
  prompt: 'Generate a recipe.',
});

You can also provide a custom transform function for models with特殊 formats:

const model = wrapLanguageModel({
  model: yourModel,
  middleware: extractJsonMiddleware({
    transform: text => text.replace(/^PREFIX/, '').replace(/SUFFIX$/, ''),
  }),
});

Default Settings and Default Instructions

import { wrapLanguageModel, defaultSettingsMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: defaultSettingsMiddleware({
    settings: {
      temperature: 0.5,
      maxOutputTokens: 800,
      providerOptions: { openai: { store: false } },
    },
  }),
});

defaultInstructionsMiddleware only applies default instructions to calls that do not have a system message; instructions given directly in the call take priority over defaults:

import { wrapLanguageModel, defaultInstructionsMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: defaultInstructionsMiddleware({
    instructions: 'You are a concise technical assistant.',
  }),
});

12.4 Implementing Custom Middleware

⚠️ Implementing middleware is an advanced feature that requires understanding the language model spec (LanguageModelV4).

You can implement one of three functions to modify model behavior:

1. transformParams: Transform parameters before they are sent to the model (applies to both doGenerate and doStream);
2. wrapGenerate: Wrap the model's doGenerate method; can modify parameters, call the model, and modify the result;
3. wrapStream: Wrap the model's doStream method; can transform stream parts.

Logging Middleware

import type {
  LanguageModelV4Middleware,
  LanguageModelV4StreamPart,
} from 'ai';

const logMiddleware: LanguageModelV4Middleware = {
  wrapGenerate({ doGenerate }) {
    return async args => {
      console.log('[generate] prompt:', args.prompt);
      const result = await doGenerate(args);
      console.log('[generate] response:', result.text);
      return result;
    };
  },
  wrapStream({ doStream }) {
    return async args => {
      console.log('[stream] prompt:', args.prompt);
      const { stream, ...rest } = await doStream(args);
      return {
        stream: (async function* () {
          for await (const part of stream) {
            console.log('[stream part]', part);
            yield part;
          }
        })(),
        ...rest,
      };
    };
  },
};

12.5 Testing with Mock Models

The ai/test package provides mock models and testing helpers for deterministic unit tests:

import { mockLanguageModel } from 'ai/test';
import { generateText } from 'ai';

const model = mockLanguageModel({
  onGenerateText: async () => ({
    text: 'Hello from mock!',
    finishReason: 'stop',
    usage: { promptTokens: 10, completionTokens: 5 },
  }),
});

const { text } = await generateText({
  model,
  prompt: 'Say hello',
});

console.log(text); // 'Hello from mock!'

Simulating Streaming Responses

Use simulateReadableStream to模拟 stream responses in tests:

import { simulateReadableStream } from 'ai/test';
import { streamText } from 'ai';

const model = mockLanguageModel({
  onStreamText: async () => ({
    stream: simulateReadableStream({
      chunks: [
        { type: 'text-delta', textDelta: 'Hello' },
        { type: 'text-delta', textDelta: ' world' },
        { type: 'finish', finishReason: 'stop', usage: { promptTokens: 5, completionTokens: 2 } },
      ],
    }),
  }),
});

const { textStream } = streamText({ model, prompt: 'Say hello' });
// Consume the stream...

Chapter Summary

• Middleware intercepts and enhances language model calls in a model-agnostic way;
• wrapLanguageModel applies middleware; multiple middlewares are composed in order;
• Six built-in middlewares cover reasoning extraction, JSON extraction, streaming simulation, default settings/instructions, and tool input examples;
• Custom middleware implements transformParams, wrapGenerate, or wrapStream;
• ai/test provides mock models and simulateReadableStream for deterministic unit tests.

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. What is the application order when wrapping multiple middlewares?

2. When does defaultInstructionsMiddleware apply its instructions?

3. What is the purpose of simulateReadableStream in testing?

4. Which built-in middleware strips markdown code fences from JSON output?

🛠️ Hands-On Exercises

1. Write a custom logging middleware that prints the prompt and response for every generate and stream call; verify it logs correctly.
2. Use mockLanguageModel to write a unit test for a generateText call that asserts the returned text; run the test to confirm it is deterministic.
3. Use simulateReadableStream to test a streamText call chunk by chunk; assert each chunk's type and content.

Next: Chapter 13 · Telemetry and Error Handling