How to create custom middleware for a LangChain agent

LangChain middleware is the extension layer inside an agent where tracing, request shaping, counters, and guardrails can run outside prompt text and tool functions. A small custom hook keeps that behavior attached to the agent factory instead of scattering it through the application code that calls the agent.

The smoke test uses an AgentMiddleware subclass with a custom state_schema, a before_model hook, and a wrap_model_call hook. before_model runs just before each model request, while wrap_model_call receives the model request and a handler so the middleware can inspect or modify the request before handing it to the model.

A deterministic fake chat model keeps the test independent of provider credentials while still exercising create_agent() and agent.invoke(). The printed hook lines show that the middleware ran, and the returned state fields confirm that the hook update traveled through the agent state. Middleware can read raw user messages, so production traces should be sanitized before they reach logs or observability tools.

1. Open an activated Python project environment.

   Use a virtual environment so the middleware test imports the same LangChain package set that the project will use.
   Related: How to create a virtual environment for LangChain

2. Install or upgrade LangChain in the project environment.

   $ python -m pip install --upgrade langchain

   LangChain v1 requires Python 3.10 or newer. Provider integrations such as langchain-openai are separate packages when a real model provider is used.
   Related: How to install LangChain with pip

3. Save the custom middleware smoke test.

   custom_middleware_agent.py

   from collections.abc import Callable
   from typing import Any
    
   from langchain.agents import create_agent
   from langchain.agents.middleware import (
       AgentMiddleware,
       AgentState,
       ModelRequest,
       ModelResponse,
   )
   from langchain.messages import SystemMessage
   from langchain_core.language_models.fake_chat_models import FakeListChatModel
   from langgraph.runtime import Runtime
   from typing_extensions import NotRequired
    
    
   class AgentReadyFakeModel(FakeListChatModel):
       def bind_tools(self, tools, *, tool_choice=None, **kwargs):
           if tools:
               raise NotImplementedError("This smoke test does not exercise tools.")
           return self
    
    
   class TraceState(AgentState):
       model_call_count: NotRequired[int]
       last_user_message: NotRequired[str]
    
    
   class TraceModelCallMiddleware(AgentMiddleware[TraceState]):
       state_schema = TraceState
    
       def before_model(
           self, state: TraceState, runtime: Runtime
       ) -> dict[str, Any] | None:
           count = state.get("model_call_count", 0) + 1
           user_messages = [
               message
               for message in state["messages"]
               if getattr(message, "type", None) == "human"
           ]
           latest = str(user_messages[-1].content) if user_messages else ""
           print(f"middleware before_model: call={count}, latest={latest!r}")
           return {
               "model_call_count": count,
               "last_user_message": latest,
           }
    
       def wrap_model_call(
           self,
           request: ModelRequest,
           handler: Callable[[ModelRequest], ModelResponse],
       ) -> ModelResponse:
           content_blocks = list(request.system_message.content_blocks)
           content_blocks.append(
               {"type": "text", "text": "Add the middleware trace marker."}
           )
           patched_request = request.override(
               system_message=SystemMessage(content=content_blocks)
           )
           print(f"middleware wrap_model_call: messages={len(request.messages)}")
           return handler(patched_request)
    
    
   model = AgentReadyFakeModel(
       responses=["middleware trace marker: custom hook ran."]
   )
    
   agent = create_agent(
       model=model,
       tools=[],
       system_prompt="Answer briefly.",
       middleware=[TraceModelCallMiddleware()],
   )
    
   result = agent.invoke(
       {
           "messages": [
               {"role": "user", "content": "Test the custom middleware."}
           ],
           "model_call_count": 0,
       }
   )
    
   print(f"Agent reply: {result['messages'][-1].content}")
   print(f"Model calls recorded: {result['model_call_count']}")
   print(f"Last user message: {result['last_user_message']}")

   TraceState declares the custom fields that the middleware writes. before_model() records the user message and call count, and wrap_model_call() modifies the model request before calling the handler.

4. Run the middleware script.

   $ python custom_middleware_agent.py
   middleware before_model: call=1, latest='Test the custom middleware.'
   middleware wrap_model_call: messages=1
   Agent reply: middleware trace marker: custom hook ran.
   Model calls recorded: 1
   Last user message: Test the custom middleware.

   The two middleware lines come from the custom hooks. The final state fields show that the model_call_count and last_user_message updates were returned through agent.invoke().

   Related: How to call a chat model in LangChain
   Related: How to build a tool-calling agent in LangChain