Tool Middleware — Composable Interceptors Around Tool Execution

A ToolMiddleware wraps the call to tool.on_invoke(...) so cross-cutting concerns — logging, metrics, distributed-tracing spans, retries, agent-global arg injection — can run without modifying every tool individually.

The shape is (ctx, tool, args, next) -> result. Implementations can:

• Run code before the tool (timer start, span open, args validation)
• Pass control to the rest of the chain via next(ctx, tool, args)
• Run code after the tool (timer stop, span close, result transformation)
• Mutate args before passing them to the next link
• Transform the result before returning it to the outer middleware
• Short-circuit by returning a result without calling next (or by raising ToolMiddlewareTermination)

Why a separate mechanism

Tool guardrails (tool.input_guardrails, tool.output_guardrails) are per-tool typed verdicts with .allow() / .reject_content() / .raise_exception() semantics. They cannot share state between pre and post (separate functions), cannot transform args, cannot short-circuit inner middleware, and must be re-declared on every tool to apply broadly. RunHooks.on_tool_start / on_tool_end are observation callbacks — they cannot modify args, results, or control flow.

Pydantic-AI ships both per-tool validators and WrapperToolset; Microsoft Agent Framework ships FunctionMiddleware, AgentMiddleware, and ChatMiddleware as three separate Protocols with their own typed contexts. The Augments ADK follows the same split.

Comparison with other frameworks

Different frameworks answer the same question differently. The matrix below shows where each one places safety verdicts (PII detection, content filtering, approval gates) versus plumbing (logging, metrics, retries):

Three notes that explain the Augments ADK position:

1. Why verdicts and plumbing stay separate. LangChain's PIIMiddleware(strategy="block") collapses verdict and plumbing into one polymorphic surface. The cost is that the verdict becomes implicit in the middleware's behaviour rather than an explicit typed return — a reader cannot tell from the signature whether a middleware halts execution, redacts content, or merely observes. The Augments ADK keeps Guardrail as the typed-verdict surface and ToolMiddleware as plumbing only; both are loadable by the same review tooling, but they do not share a registration list.
2. Why there is one Protocol per layer. As with Microsoft's FunctionMiddleware / AgentMiddleware / ChatMiddleware, each layer carries its own typed context, so a turn-scope middleware does not have to pretend it sees a ToolContext. The Middleware config dataclass on Agent mirrors that split with plural slot names (tools, agents, llms, stream_llms).
3. What this means in practice. Reach for a guardrail when the answer is "verdict on a single decision point". Reach for middleware when the answer is "wrap the call with shared state, transforms, or short-circuits". The two surfaces compose: the executor runs input guardrail → middleware chain → tool → middleware unwind → output guardrail, so they cannot conflict.

Quickstart

from augments.adk import Agent, Middleware
from augments.adk.tools import ToolLoggingMiddleware, ToolMetricsMiddleware
 
 
class MyMetrics:
    def record_duration(self, tool_name, duration_seconds):
        # Send to StatsD/Prometheus/OpenTelemetry/etc.
        ...
 
    def record_outcome(self, tool_name, *, success):
        ...
 
 
agent = Agent(
    name="Service",
    system_prompt="...",
    tools=[search, summarise, lookup],
    middleware=Middleware(
        tools=[
            ToolLoggingMiddleware(),
            ToolMetricsMiddleware(recorder=MyMetrics()),
        ],
    ),
)

Agent.middleware is a single typed config object holding per-layer middleware lists: tools (tool calls), agents (turn-scope), llms (non-streaming LLM calls), and stream_llms (streaming LLM calls). This page covers the tools slot; see Scope: three middleware layers for the others.

Every call to search, summarise, or lookup flows through both middleware in the listed order (logging outer, metrics inner).

Authoring a custom middleware

import time
from augments.adk.tools import ToolMiddleware
 
 
class TimingMiddleware:
    """Time every tool call and accumulate per-tool totals."""
 
    def __init__(self):
        self.totals: dict[str, float] = {}
 
    async def __call__(self, ctx, tool, args, next):
        start = time.monotonic()
        result = await next(ctx, tool, args)
        duration = time.monotonic() - start
        self.totals[tool.name] = self.totals.get(tool.name, 0.0) + duration
        return result

The ToolMiddleware Protocol is runtime_checkable, so isinstance(my_middleware, ToolMiddleware) works for tests and introspection.

Chain semantics

The list is composed outer-to-inner: the first middleware runs first (outermost), the last runs last (innermost). Each middleware can call next(...) exactly once or zero times.

[A, B, C]  →  A.pre → B.pre → C.pre → tool → C.post → B.post → A.post

Mutating args

A middleware can mutate the args dict in-place (Python passes dicts by reference) before calling next. The tool sees the mutated args:

class InjectRequestID:
    async def __call__(self, ctx, tool, args, next):
        args["request_id"] = ctx.context.get("request_id", "anonymous")
        return await next(ctx, tool, args)

Short-circuiting

A middleware that does not call next short-circuits the chain. Useful for circuit-breaker / cache patterns:

from augments.adk.types.output.function_tool_call_result import FunctionToolCallResult
 
 
class CacheMiddleware:
    def __init__(self):
        self.cache: dict[tuple, str] = {}
 
    async def __call__(self, ctx, tool, args, next):
        key = (tool.name, tuple(sorted(args.items())))
        if key in self.cache:
            return FunctionToolCallResult(
                type="function_call_output",
                call_id=ctx.tool_call_id or "",
                output=self.cache[key],
            )
        result = await next(ctx, tool, args)
        self.cache[key] = result.output
        return result

For deeper chains where you want to skip multiple middleware layers, raise ToolMiddlewareTermination with the desired result — the executor unwinds directly without invoking outer middleware:

from augments.adk.tools import ToolMiddlewareTermination
 
 
class CircuitBreaker:
    async def __call__(self, ctx, tool, args, next):
        if self.tripped(tool.name):
            raise ToolMiddlewareTermination(
                FunctionToolCallResult(
                    type="function_call_output",
                    call_id=ctx.tool_call_id or "",
                    output=f"[breaker tripped for {tool.name}]",
                )
            )
        return await next(ctx, tool, args)

Toolset-scoped middleware

Wrap a toolset with WrapperToolset to apply middleware only to that toolset's tools:

from augments.adk import Middleware
from augments.adk.tools import WrapperToolset, FunctionToolset
 
db_tools = FunctionToolset(tools=[query, insert]).prefixed("db")
db_with_audit = WrapperToolset(
    wrapped=db_tools,
    middleware=[AuditMiddleware()],   # toolset-scoped (a list[ToolMiddleware])
)
 
agent = Agent(
    name="Service",
    system_prompt="...",
    tools=[ad_hoc_tool, db_with_audit],
    middleware=Middleware(tools=[ToolLoggingMiddleware()]),   # agent-global
)

Composition order: agent-global middleware wraps the toolset-scoped middleware. So a call to db_query runs through:

ToolLoggingMiddleware.pre → AuditMiddleware.pre → query → AuditMiddleware.post → ToolLoggingMiddleware.post

ad_hoc_tool only runs through ToolLoggingMiddleware (no toolset wrapping).

Standard middleware shipped

Both are reference implementations meant to be copied or extended. The framework deliberately does not ship a full observability stack — choosing the metrics backend is a deployment concern.

When NOT to write a middleware

Middleware is plumbing. Anything that decides whether a call should proceed, whether content is acceptable, or what verdict to apply on detection belongs elsewhere. Common confusions and their right-place-to-go:

Allowed plumbing patterns — these all wrap the call without encoding a verdict:

The decision rule:

If the answer to "should this call proceed?" or "is this content acceptable?" is involved, write a guardrail. If only the surrounding plumbing (logging, metrics, retries, arg injection, caching) changes, write a middleware.

The contract is stated in two places: the ToolMiddleware Protocol docstring and this section.

Scope: three middleware layers

The framework ships three sibling middleware Protocols, one per execution layer. All three share the same value-return shape and the same plumbing-only contract.

Composition order during one turn is

Agent.middleware.agents (outermost)
  Agent.middleware.llms
    LLM.acomplete()
  return
  Agent.middleware.tools (per tool call)
    tool.on_invoke()
  return
return

The plumbing-only / no-verdicts contract applies unchanged at every layer — the forbidden-vs-allowed table in this page is normative for AgentMiddleware and LLMMiddleware too. Only the typed verdict surfaces (AgentInputGuardrail, AgentOutputGuardrail, ToolInputGuardrail, ToolOutputGuardrail, RunConfig.usage_limits, FunctionTool.requires_approval, FunctionTool.rate_limit, FunctionTool.schema_enforcement) may carry policy decisions.

Streaming support

All three middleware Protocols apply on both streaming and non-streaming runs:

• ToolMiddleware (this page) — drain of streaming tools' AsyncIterator[ToolStreamEvent] happens at the innermost middleware terminal so middleware sees the final accumulated value, not chunks.
• AgentMiddleware — composes inside both run_agent_loop and run_agent_loop_streamed.
• LLMMiddleware (non-streaming) and LLMStreamMiddleware (streaming) — sibling Protocols on Middleware.llms and Middleware.stream_llms slots respectively. Use make_logging_middlewares() to register a paired logger across both LLM paths.

See also

• docs/agents/agent_middleware.md — AgentMiddleware contract.
• docs/llms/llm_middleware.md — LLMMiddleware contract.
• examples/tools/middleware/ — runnable examples (including a three-layer demo).
• src/augments/adk/tools/tool_middleware.py — Protocol definition, shipped middleware, and the chain composition helpers.
• docs/tools/toolsets.md — toolset composition (where WrapperToolset.middleware lives).