Tool middleware#

Tool middleware lets you wrap one function tool with async hooks that run immediately around its execution—the same idea as an agent's on_tool_execution() Middleware, but attached at tool definition time with plain callables (no BaseMiddleware subclass).

Use this pattern when behavior is specific to a single tool. For policies that apply to every tool on an agent, register BaseMiddleware with on_tool_execution() instead.

Why use it#

Tool-scoped hooks are optional. They help when:

Colocation — Validation, redaction, or metrics for one tool live next to that implementation instead of in shared agent middleware.
Clear contracts — Libraries can ship a tool with hooks that always run (normalize arguments, scrub secrets) without requiring consumers to register matching agent middleware.
Simpler agents — You avoid a large on_tool_execution full of if event.name == ... when only a few tools need special handling.

Use agent middleware=[...] when the policy is global or shared across most tools. Use middleware=[...] on @tool, @agent.tool, or @toolkit.tool when the behavior belongs to that tool only.

Typical cases#

Scenario	Reason to use tool-scoped hooks
Normalize or validate arguments for one function	Only that tool’s schema needs the transform; keeps agent middleware small.
Redact or reshape results before they return to the model	Per-tool privacy or formatting (for example strip internal IDs).
Light auditing or metrics for a sensitive action	The hook is bundled with the tool so it is hard to forget at agent setup.
Retry or fallback tied to one integration	Failure handling stays next to the API client without naming the tool in global middleware.
Approve or reject before the tool body runs	Gate dangerous or irreversible tools on policy, session flags, or a human decision without scattering checks inside the implementation.

API#

The public type alias is ToolMiddleware in ag2.middleware. A hook is an async callable with the same parameters as BaseMiddleware.on_tool_execution, except the first argument is the inner ToolExecution (the next step in the chain), not self.

Pass middleware=[hook, ...] to @tool, Agent.tool, or Toolkit.tool.
Multiple hooks use the same nesting as agent tool middleware: the first entry in the list is the outermost layer around the tool body.
If the agent also registers BaseMiddleware with on_tool_execution, agent middleware runs outside tool-scoped hooks (it sees the full execution, including hooks).

from typing import Annotated

from ag2 import Agent, Context, tool, Variable
from ag2.config import OpenAIConfig
from ag2.events import ToolCallEvent, ToolResultEvent
from ag2.middleware import ToolExecution

async def add_request_id(
    call_next: ToolExecution,
    event: ToolCallEvent,
    context: Context,
) -> ToolResultEvent:
    context.variables.setdefault("request_id", "unknown")
    return await call_next(event, context)

@tool(middleware=[add_request_id])
def search(query: str, request_id: Annotated[str, Variable()]) -> str:
    """Runs a search. request_id may be injected by middleware."""
    return f"results-for-{query}-{request_id}"

agent = Agent("assistant", config=OpenAIConfig("gpt-4o-mini"))

Note

Tool-scoped hooks are plain callables. They do not use the Middleware(...) factory or BaseMiddleware.

Making a hook describable#

A plain function like add_request_id above works, but cannot report its configuration: it has no describe(), so AG2 reports it as complete=False rather than guessing, because its settings live in closure cells that AG2 does not inspect.

For a describable hook, write a class with async def __call__ and a describe() method. It is still a plain callable, so how you pass it does not change:

from ag2 import Context, tool
from ag2.events import ToolCallEvent
from ag2.middleware import MiddlewareDescription, ToolExecution, ToolResultType

class AddRequestId:
    def __init__(self, default: str = "unknown") -> None:
        self._default = default

    async def __call__(
        self,
        call_next: ToolExecution,
        event: ToolCallEvent,
        context: Context,
    ) -> ToolResultType:
        context.variables.setdefault("request_id", self._default)
        return await call_next(event, context)

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

@tool(middleware=[AddRequestId()])
def search(query: str) -> str:
    """Runs a search."""
    return f"results-for-{query}"

approval_required() is the built-in example: it returns an ApprovalRequired instance that describes its prompt, timeout, and allow_always settings. Keep run-time state out of the hook — config is for settings only.

Which carrier to use for run-time state

context.variables is rebuilt for every ask() and is copied into a sub-task without being copied back, so state kept there is scoped to one run. For state that must outlive a single ask() or count work done by sub-tasks, use context.dependencies: values are passed by reference into every run and into sub-tasks, and are never copied or serialised.

Adding middleware to an existing tool#

Use tool.with_middleware() to wrap a tool with additional hooks without modifying the original. The returned tool is an independent copy with the new middleware as the outermost layer:

from ag2 import tool
from ag2.events import ToolCallEvent, ToolResultEvent
from ag2.middleware import ToolExecution

async def audit(
    call_next: ToolExecution, event: ToolCallEvent, context: Context,
) -> ToolResultEvent:
    print(f"audit: {event.name}")
    return await call_next(event, context)

@tool
def delete_record(record_id: str) -> str:
    """Deletes a record."""
    return f"Deleted {record_id}"

audited_delete = delete_record.with_middleware(audit)

# delete_record is unchanged; audited_delete runs audit → delete_record

Attaching one hook instance to several tools does not give those tools a combined allowance. Registering a tool deep-copies it, so a class-based hook is duplicated and each tool ends up counting on its own instance. A plain function is left as-is by deepcopy and stays shared, which is why the same hook behaves differently depending on how it was written.

Do not rely on either. Put the shared state in context.dependencies and keep the hook itself a settings object — then it does not matter how many copies of the hook exist:

from ag2 import Agent, Context, tool
from ag2.config import OpenAIConfig
from ag2.events import ToolCallEvent, ToolResultEvent
from ag2.middleware import MiddlewareDescription, ToolExecution, ToolResultType

class Counter:
    def __init__(self) -> None:
        self.used = 0

class CallBudget:
    def __init__(self, limit: int, *, key: str = "default") -> None:
        self._limit = limit
        self._key = key

    def describe(self) -> MiddlewareDescription:
        return MiddlewareDescription(
            kind=type(self).__qualname__,
            config={"limit": self._limit, "key": self._key},
        )

    async def __call__(
        self, call_next: ToolExecution, event: ToolCallEvent, context: Context,
    ) -> ToolResultType:
        counter: Counter = context.dependencies[("budget", self._key)]
        if counter.used >= self._limit:
            return ToolResultEvent.from_call(event, result="budget exhausted")
        counter.used += 1
        return await call_next(event, context)

def search(query: str) -> str:
    """Searches the web."""
    return f"results-for-{query}"

def fetch(url: str) -> str:
    """Fetches a URL."""
    return f"body-of-{url}"

budget = CallBudget(limit=3, key="web")
agent = Agent(
    "assistant",
    config=OpenAIConfig("gpt-4o-mini"),
    tools=[tool(search, middleware=[budget]), tool(fetch, middleware=[budget])],
    dependencies={("budget", "web"): Counter()},
)

search and fetch draw on the same three calls, and the allowance survives every ask() and counts calls made by sub-tasks. The hook stays describable, because the limit is a setting and only the counter lives in the context.

Toolkit-level middleware#

Pass middleware=[...] to a Toolkit constructor to apply hooks to all tools in the set. Toolkit middleware is the outermost layer — it runs before any per-tool hooks. See Toolkit middleware.