Skip to content

AG2 Middleware#

Middleware lets you intercept and customize how an AG2 agent runs a turn. It's the right tool when you want to add cross-cutting behavior such as logging, retries, history trimming, request mutation, tool auditing, or guardrails without changing the agent, model client, or tools themselves.

At a high level, middleware can wrap four parts of the runtime:

  • the full agent turn
  • each LLM call
  • each tool execution
  • each human input request

This makes it a good fit for behavior that should apply consistently across many runs.

What is Middleware#

Middleware is an object that receives the current turn's initial event and Context, then participates in one or more lifecycle hooks.

Each middleware instance is created at the beginning of a turn and can keep per-turn state on self. That same instance can then observe or modify the turn, the LLM call, tool execution, and human input as the run progresses.

In practice, you use middleware to:

  • add observability such as logging, tracing, and timing
  • enforce policies before a tool runs
  • retry transient model failures
  • trim conversation history before sending it to the model
  • normalize tool inputs or outputs
  • short-circuit or reshape a response

Middleware Hooks#

BaseMiddleware exposes four async hooks. You can implement just one of them or mix several in the same class.

on_turn()#

1
2
3
4
5
6
7
8
class BaseMiddleware:
    async def on_turn(
        self,
        call_next: Callable[[BaseEvent, Context], Awaitable[ModelResponse]],
        event: BaseEvent,
        context: Context,
    ) -> ModelResponse:
        return await call_next(event, context)

on_turn() wraps the whole agent turn. It receives the incoming event and the final ModelResponse.

Use on_turn() when you want to:

  • measure total turn latency
  • inspect or rewrite the initial request before anything else happens
  • inspect or rewrite the final response before it is returned
  • implement turn-level policies, approvals, or short-circuit behavior

Conceptually, this is the outermost hook around a single ask(...) call.

on_llm_call()#

1
2
3
4
5
6
7
8
class BaseMiddleware:
    async def on_llm_call(
        self,
        call_next: Callable[[Sequence[BaseEvent], Context], Awaitable[ModelResponse]],
        events: Sequence[BaseEvent],
        context: Context,
    ) -> ModelResponse:
        return await call_next(events, context)

on_llm_call() wraps the call to the configured model client. It receives the event history that will be sent to the LLM.

Use on_llm_call() when you want to:

  • retry transient client failures
  • log prompts and responses
  • trim history before it reaches the model
  • sanitize context / model response
  • inject additional request-time instructions through event mutation
  • implement caching or request deduplication around model calls

This is the hook used by built-in history and token limiting middleware.

on_tool_execution()#

1
2
3
4
5
6
7
8
class BaseMiddleware:
    async def on_tool_execution(
        self,
        call_next: Callable[[ToolCallEvent, Context], Awaitable[ToolResultType]],
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        return await call_next(event, context)

on_tool_execution() wraps each tool invocation triggered during the turn. It receives the current ToolCallEvent and returns a ToolResultType.

Use on_tool_execution() when you want to:

  • validate or rewrite tool arguments before execution
  • log tool usage
  • transform tool results before they go back into the event stream
  • capture tool failures and replace them with safer fallback results
  • enforce access control around specific tools

For wrapping one tool at definition time with async hooks, see Tool middleware.

on_human_input()#

1
2
3
4
5
6
7
8
class BaseMiddleware:
    async def on_human_input(
        self,
        call_next: Callable[[HumanInputRequest, Context], Awaitable[HumanMessage]],
        event: HumanInputRequest,
        context: Context,
    ) -> HumanMessage:
        return await call_next(event, context)

on_human_input() wraps each human-in-the-loop (HITL) request triggered during a turn. It receives the HumanInputRequest emitted by a tool via ctx.input(...) and can intercept or modify both the request and the HumanMessage response.

Use on_human_input() when you want to:

  • log or audit human input requests and responses
  • rewrite or enrich the prompt shown to the human
  • transform the human's reply before it reaches the tool
  • short-circuit the request with an automated response instead of asking a human
  • enforce policies or rate limits on human input requests

Registering Middleware#

On an Agent#

To make middleware apply to every turn for an agent, pass it through the middleware argument when constructing the agent.

from ag2 import Agent
from ag2.config import OpenAIConfig
from ag2.middleware import LoggingMiddleware, RetryMiddleware

agent = Agent(
    "assistant",
    prompt="Be helpful.",
    config=OpenAIConfig("gpt-4o-mini"),
    middleware=[
        LoggingMiddleware(),
        RetryMiddleware(max_retries=2),
    ],
)

Use agent-level registration for behavior that should always be present, such as logging, tracing, or default retry policy.

On a Single Call#

You can also add middleware just for a specific turn. This is useful when you want temporary behavior without changing the agent's defaults.

Both Agent.ask(...) and AgentReply.ask(...) accept a middleware argument.

from ag2 import Agent
from ag2.config import OpenAIConfig
from ag2.middleware import LoggingMiddleware, TokenLimiter

agent = Agent(
    "assistant",
    prompt="Be helpful.",
    config=OpenAIConfig("gpt-4o-mini"),
)

reply = await agent.ask(
    "Summarize the latest messages.",
    middleware=[LoggingMiddleware()],
)

next_turn = await reply.ask(
    "Now answer in one paragraph.",
    middleware=[TokenLimiter(max_tokens=4000)],
)

Call-level middleware is appended after the middleware list defined on the agent.

Middleware Ordering#

Middleware runs in the order you register them. If you register [A, B, C], they enter in the order A -> B -> C and unwind in reverse order C -> B -> A.

This matters when you combine behaviors such as logging, mutation, and retries.

from ag2 import Agent, Context
from ag2.config import OpenAIConfig
from ag2.events import BaseEvent, ModelResponse
from ag2.middleware import AgentTurn, BaseMiddleware

class A(BaseMiddleware):
    async def on_turn(
        self,
        call_next: AgentTurn,
        event: BaseEvent,
        context: Context,
    ) -> ModelResponse:
        print("enter A")
        response = await call_next(event, context)
        print("exit A")
        return response

class B(BaseMiddleware):
    async def on_turn(
        self,
        call_next: AgentTurn,
        event: BaseEvent,
        context: Context,
    ) -> ModelResponse:
        print("enter B")
        response = await call_next(event, context)
        print("exit B")
        return response

class C(BaseMiddleware):
    async def on_turn(
        self,
        call_next: AgentTurn,
        event: BaseEvent,
        context: Context,
    ) -> ModelResponse:
        print("enter C")
        response = await call_next(event, context)
        print("exit C")
        return response

agent = Agent(
    "assistant",
    prompt="Be helpful.",
    config=OpenAIConfig("gpt-4o-mini"),
    middleware=[A, B],
)

await agent.ask(
    "Hello",
    middleware=[C],
)

# Output:
# enter A
# enter B
# enter C
# exit C
# exit B
# exit A

Writing Your Own Middleware#

To create custom middleware, subclass BaseMiddleware and implement the hooks you need.

If your middleware does not need extra constructor arguments, you can register the class directly. If it does need configuration, wrap it with Middleware(...) when registering it.

import logging
from collections.abc import Sequence

from ag2 import Agent, Context
from ag2.config import OpenAIConfig
from ag2.events import BaseEvent, ModelResponse, ToolCallEvent
from ag2.middleware import BaseMiddleware, LLMCall, Middleware, ToolExecution

class AuditMiddleware(BaseMiddleware):
    def __init__(
        self,
        event: BaseEvent,
        context: Context,
        logger: logging.Logger,
    ) -> None:
        super().__init__(event, context)
        self.logger = logger

    async def on_llm_call(
        self,
        call_next: LLMCall,
        events: Sequence[BaseEvent],
        context: Context,
    ) -> ModelResponse:
        self.logger.info("Calling model with %d events", len(events))
        response = await call_next(events, context)
        self.logger.info("Model returned: %s", response)
        return response

    async def on_tool_execution(
        self,
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ):
        self.logger.info("Executing tool: %s", event.name)
        return await call_next(event, context)

agent = Agent(
    "assistant",
    prompt="Be helpful.",
    config=OpenAIConfig("gpt-4o-mini"),
    middleware=[
        Middleware(AuditMiddleware, logger=logging.getLogger("ag2.audit")),
    ],
)

Guidelines for Custom Middleware#

  • Keep hook behavior focused. Middleware that does one job well is easier to reason about than one that handles, for example, logging, retries, mutation, and policy checks together.
  • Prefer on_turn() for whole-run behavior, on_llm_call() for model-facing behavior, on_tool_execution() for tool-facing behavior, and on_human_input() for human-in-the-loop behavior.
  • Be deliberate when mutating event, events, or tool results. Later executing middleware and the rest of the runtime will observe those changes.
  • Register zero-config middleware classes directly, and use Middleware(YourMiddleware, ...) when the constructor needs additional options.

Built-In Middleware#

AG2 currently includes four built-in middleware in ag2.middleware:

LoggingMiddleware#

1
2
3
4
from ag2 import Agent
from ag2.middleware import LoggingMiddleware

agent = Agent(..., middleware=[LoggingMiddleware()])

Logs the lifecycle of a turn, including:

  • when a turn starts and finishes
  • each LLM call and its response time
  • each tool execution and its result

Use it for quick debugging or application-level observability.

RetryMiddleware#

1
2
3
4
from ag2 import Agent
from ag2.middleware import RetryMiddleware

agent = Agent(..., middleware=[RetryMiddleware(max_retries=2)])

Retries failed LLM calls up to max_retries times. By default it retries any Exception, but you can narrow that with retry_on=....

Use it for transient failures such as provider timeouts or flaky network issues.

HistoryLimiter#

1
2
3
4
from ag2 import Agent
from ag2.middleware import HistoryLimiter

agent = Agent(..., middleware=[HistoryLimiter(max_events=100)])

Trims the event history to a maximum number of events before the model call. It preserves the first ModelRequest when possible, and drops whatever the cut orphaned so the trimmed history still replays.

max_events is a target rather than a hard cap, for the same reason it is one on the reduction policies — see Reduction limits are targets.

Use it when you want a simple, deterministic cap on context length by event count.

TokenLimiter#

1
2
3
4
from ag2 import Agent
from ag2.middleware import TokenLimiter

agent = Agent(..., middleware=[TokenLimiter(max_tokens=1000)])

Trims the event history to fit within an approximate token budget before the model call. It uses a character-based estimate controlled by chars_per_token.

max_tokens is a target rather than a hard cap — see Reduction limits are targets.

Use it when you need lightweight context budgeting without depending on a model-specific tokenizer.

Conditional Middleware#

ConditionalMiddleware lets you gate any middleware so each hook only activates when a condition matches the hook's own event. When the condition is not met, that hook passes through to the next middleware in the chain.

This is useful when you have middleware that should only run for certain event types — for example, approval middleware that only fires for a specific tool — without writing condition checks inside every hook method.

from ag2 import Agent
from ag2.config import OpenAIConfig
from ag2.events import ToolCallEvent
from ag2.middleware import ConditionalMiddleware, Middleware

agent = Agent(
    "assistant",
    prompt="Be helpful.",
    config=OpenAIConfig("gpt-4o-mini"),
    middleware=[
        ConditionalMiddleware(
            Middleware(ApprovalMiddleware),
            condition=ToolCallEvent.name == "execute_code",
        ),
    ],
)

In this example, ApprovalMiddleware only activates during on_tool_execution when the tool name is "execute_code". The on_turn hook receives a different event type (ModelRequest), so the condition does not match there and the middleware passes through.

Each hook — on_turn, on_tool_execution, on_human_input — checks the condition against its own event. on_llm_call checks against the initial turn event, since it receives a sequence of events rather than a single one. When a condition targets a specific event type like ToolCallEvent, hooks that receive a different type will naturally pass through.

Composing Conditions#

Conditions support & (and), | (or), and ~ (not) operators, so you can build expressive gates:

from ag2.events import ToolCallEvent
from ag2.middleware import ConditionalMiddleware, Middleware

# Activate only for tool calls named "search" or "browse"
conditional = ConditionalMiddleware(
    Middleware(MyAuditMiddleware, logger=logger),
    condition=(ToolCallEvent.name == "search") | (ToolCallEvent.name == "browse"),
)

# Activate for all tool calls EXCEPT "calculate"
conditional = ConditionalMiddleware(
    Middleware(MyAuditMiddleware, logger=logger),
    condition=~(ToolCallEvent.name == "calculate"),
)

You can also pass a bare event type as the condition — it is automatically wrapped:

1
2
3
4
5
# Activate only when the hook receives a ToolCallEvent
conditional = ConditionalMiddleware(
    Middleware(MyAuditMiddleware, logger=logger),
    condition=ToolCallEvent,
)

ConditionalMiddleware wraps any MiddlewareFactory — including Middleware(...) instances, bare BaseMiddleware subclasses, and other ConditionalMiddleware wrappers.

Describing Middleware#

A middleware can report what it is and how it was configured, so it can be logged, compared against another instance, or asserted on in a test. Every built-in does:

1
2
3
4
5
6
7
from ag2.middleware import TokenLimiter

TokenLimiter(max_tokens=100).describe()
# MiddlewareDescription(kind='TokenLimiter', config={'max_tokens': 100, 'chars_per_token': 4}, complete=True)

TokenLimiter(max_tokens=100).describe() == TokenLimiter(max_tokens=100).describe()
# True

Your own middleware opts in the same way, by implementing describe():

from ag2.middleware import MiddlewareDescription

class RateLimit:
    def __init__(self, per_minute: int) -> None:
        self._per_minute = per_minute

    async def __call__(self, call_next, event, context):
        return await call_next(event, context)

    def describe(self) -> MiddlewareDescription:
        return MiddlewareDescription(
            kind=type(self).__qualname__,
            config={"per_minute": self._per_minute},
        )

Opting in is optional. Where AG2 surfaces middleware, anything that has not opted in is reported honestly rather than guessed at — config is empty and complete is False. AG2 never inspects __closure__ to recover configuration, because cell names and ordering are an implementation detail of whatever function produced the closure. A describe() that fails or returns the wrong type degrades to that same honest unknown, so introspection cannot break the code doing the logging.

Middleware that wraps other middleware reports them in inner, keeping config a flat settings mapping:

1
2
3
4
5
6
7
8
from ag2.middleware import ConditionalMiddleware, TokenLimiter
from ag2.events import ToolCallEvent

ConditionalMiddleware(TokenLimiter(max_tokens=10), ToolCallEvent).describe()
# MiddlewareDescription(
#     kind='ConditionalMiddleware', config={'condition': 'TypeCondition'}, complete=True,
#     inner=(MiddlewareDescription(kind='TokenLimiter', config={'max_tokens': 10, ...}),),
# )

A composite is only as describable as its parts: complete is forced to False when any entry in inner is incomplete, so composing middleware cannot claim completeness its parts do not support. Because each description applies the rule as it is constructed, it holds at any depth without a parent walking the tree.

Descriptions support equality but are deliberately not hashableconfig may hold arbitrary values, so they cannot go in a set or be used as dict keys.

To read the middleware attached to an agent or a tool, use agent.middleware and tool.middleware. Each entry pairs the middleware object with its description, so an entry always has one whether or not that middleware opted in. See Reading an agent's composition.

Note

Middleware(...) reports the wrapped class and which options were set, but never their values, and always complete=False. The options are caller-supplied and the wrapper cannot know whether one is a credential. Write a MiddlewareFactory with its own describe() when you want a complete description. Its .cls and .options expose the class and the option values directly, for when you are inspecting a factory you already hold rather than producing something to log.

Note

config is for settings only. Counters, caches, and anything else that changes during a run belong in context.variables or the per-turn BaseMiddleware instance — a description that carries a moving number produces flaky snapshots.

Warning

Descriptions are meant to be logged and committed as test fixtures, so describe() must not expose credentials. The built-ins report live objects by name or presence — LoggingMiddleware reports its logger's name, and TelemetryMiddleware reports only the keys of its span_attributes.

Choosing the Right Hook#

If you are unsure where a behavior belongs, use this rule of thumb:

  • Use on_turn() when the behavior is about the entire request/response lifecycle.
  • Use on_llm_call() when the behavior is about what goes into or comes out of the model.
  • Use on_tool_execution() when the behavior is about tool safety, auditing, or result shaping across tools (or when branching on event.name is acceptable).
    • Use tool-scoped middleware=[...] on @tool / Agent.tool / Toolkit.tool when the behavior applies only to that tool’s definition; see Tool middleware.
  • Use on_human_input() when the behavior is about intercepting, logging, or transforming human-in-the-loop requests and responses.

For related runtime customization patterns, see Tools, Tool middleware, Prompt Management, and Events Streaming.