Skip to content

Agent Middleware

Middleware lets you intercept and transform the inputs and outputs of an Agent's run without subclassing Agent or wrapping LLMClient. Use it for message trimming, budget enforcement, PII redaction, response-rewriting, caching, and any other cross-cutting concern that should compose cleanly.

Middleware vs Guardrails

Both plug into the Agent pipeline, but they serve different purposes.

Middleware Guardrail
Purpose Transform messages or responses Assert on content
Failure mode Can raise StopAgent to end cooperatively Raises GuardrailBlockedError (blocking) or logs (non-blocking)
Composition Onion — each middleware wraps the next Flat list — each guardrail runs independently
When to use "Trim history", "redact PII", "cap tool calls", "rewrite refusals" "Is this PII?", "Is the JSON valid?", "Is this toxic?"

Input guardrails run before middleware's before_model; output guardrails run after middleware's after_model on the final iteration.

The Three Hooks

from fastaiagent import AgentMiddleware, MiddlewareContext
from fastaiagent.llm.client import LLMResponse
from fastaiagent.llm.message import Message
from fastaiagent.tool.base import Tool, ToolResult


class MyMiddleware(AgentMiddleware):
    name = "my_middleware"

    async def before_model(
        self, ctx: MiddlewareContext, messages: list[Message]
    ) -> list[Message]:
        """Transform messages just before the LLM call."""
        return messages

    async def after_model(
        self, ctx: MiddlewareContext, response: LLMResponse
    ) -> LLMResponse:
        """Inspect or rewrite the LLM response before tool dispatch."""
        return response

    async def wrap_tool(
        self,
        ctx: MiddlewareContext,
        tool: Tool,
        args: dict,
        call_next,
    ) -> ToolResult:
        """Wrap each tool invocation. MUST ``await call_next(tool, args)``
        to proceed, or return a ToolResult to short-circuit."""
        return await call_next(tool, args)

All three hooks are optional — override only what you need.

Hook Ordering

  • before_model — declaration order (first middleware first)
  • after_model — reverse declaration order (last middleware first)
  • wrap_tool — onion (first middleware is outermost, calls into inner middleware via call_next)
Agent.arun()
 ├─ input guardrails
 ├─ middleware[0].before_model
 ├─ middleware[1].before_model
 ├─ middleware[2].before_model
 ├─ LLM.acomplete()
 ├─ middleware[2].after_model
 ├─ middleware[1].after_model
 ├─ middleware[0].after_model
 ├─ for each tool call:
 │   ├─ middleware[0].wrap_tool  (pre)
 │   │   └─ middleware[1].wrap_tool  (pre)
 │   │       └─ middleware[2].wrap_tool  (pre)
 │   │           └─ real tool
 │   │       ← middleware[2].wrap_tool  (post)
 │   │   ← middleware[1].wrap_tool  (post)
 │   ← middleware[0].wrap_tool  (post)
 ├─ (repeat LLM + tool calls per iteration)
 └─ output guardrails

The MiddlewareContext

Each run creates one MiddlewareContext that every hook sees:

Attribute Description
run_context The user-supplied RunContext, if any
turn Zero-indexed LLM iteration within the run
tool_call_index Zero-indexed tool call within the current turn
scratch Mutable dict shared across all hooks — use to pass data between before_model and after_model, or to accumulate state across tool calls
agent_name Name of the agent running

Scratch is per-run; a fresh dict on every agent.arun().

Cooperative Stop

To end a run from inside middleware, raise StopAgent:

from fastaiagent import StopAgent

class BudgetMiddleware(AgentMiddleware):
    async def before_model(self, ctx, messages):
        if ctx.turn >= 3:
            raise StopAgent("Turn budget exhausted.")
        return messages

The agent returns an AgentResult whose output is the StopAgent message. No GuardrailBlockedError, no unwinding through the caller — it's a cooperative signal.

Use StopAgent for budgets, completion signals, or feature flags. Use Guardrail for policy rejections that should surface as errors.

What the history looks like after a stop

StopAgent raised from before_model or after_model ends the turn on a boundary: nothing is half-done.

StopAgent raised from wrap_tool does not. The assistant message at the top of that turn already declares every tool call the model asked for, and the stop lands partway through dispatching them. Both OpenAI and Anthropic reject a history where an assistant tool_calls message is not answered call-for-call, so the agent closes the turn before returning (since 1.67.0):

  • the call that was in flight is answered with whatever it actually produced — a stopper that fires before call_next (ToolBudget does) produces nothing;
  • every call the stop skipped is answered with a short note saying it was not run.

The skipped calls are not dispatched. Their side effects never happen, and the note says so in terms the model can act on, rather than leaving it to assume a silent success. If your middleware needs the remaining calls to run, return a ToolResult from wrap_tool instead of raising — a short-circuit ends one call, StopAgent ends the run.

Built-in Middleware

TrimLongMessages(keep_last=20)

Keeps only the most recent keep_last messages plus any leading SystemMessage. Cheap alternative to summarization for long-running agents.

from fastaiagent import Agent, LLMClient, TrimLongMessages

agent = Agent(
    name="chatty",
    llm=LLMClient(provider="openai", model="gpt-4.1"),
    middleware=[TrimLongMessages(keep_last=30)],
)

Trimming preserves tool-call pairing. keep_last is a message count, so the window can land between an assistant message carrying tool_calls and the tool results that answer it. Since 1.67.0 the tail is re-balanced before it is returned:

  • a tool result whose parent call was trimmed away is dropped;
  • a tool call whose results were trimmed away is answered with a note saying the result is no longer in the window (the tool did run — only its result fell out of the window).

That can leave slightly more than keep_last messages. A provider 400 is the more expensive of the two.

ToolBudget(max_calls=10, message="...")

Raises StopAgent once max_calls tool invocations have occurred in a single run.

from fastaiagent import Agent, LLMClient, ToolBudget

agent = Agent(
    name="budgeted",
    llm=LLMClient(provider="openai", model="gpt-4.1"),
    middleware=[ToolBudget(max_calls=5)],
)

The budget is checked before call_next, so the call that trips it never runs — see What the history looks like after a stop.

With output_type. The stop message becomes the run's output, and it will not parse as your schema, so the agent spends its output_retries re-asking the model with the same history. That is worth knowing for two reasons: the re-ask costs turns the budget was meant to save, and it is the request that re-sends the stopped turn's history. Before 1.67.0 that second request is what returned

400 An assistant message with 'tool_calls' must be followed by tool messages
    responding to each 'tool_call_id'

Set AgentConfig(output_retries=0) if you would rather a budgeted run simply return the stop message unparsed.

RedactPII(patterns=..., placeholder="[REDACTED]", entities=...)

Redacts PII from outbound messages and inbound LLM responses.

from fastaiagent import Agent, LLMClient, RedactPII

agent = Agent(
    name="safe",
    llm=LLMClient(provider="openai", model="gpt-4.1"),
    middleware=[RedactPII()],
)

By default it delegates to the same detector behind no_pii(), the pii guardrail type and the PIILeakage scorer — email, US phone, SSN and credit cards, with cards Luhn-validated. Narrow it with entities=("email", "ssn").

Behaviour change in 1.63.0 — it redacts less, and that is the fix

Until 1.63.0 this middleware carried its own private copy of the PII regexes, and the copy had drifted: its card pattern was a bare \b(?:\d[ \-]?){13,19}\b with no Luhn check. Any 13–19 digit run — an order number, an invoice id, an IMEI — was redacted as a credit card.

That corruption was what the model saw and what was replayed in a guardrail re-ask. If you were relying on long digit runs being masked, pass an explicit patterns= list.

What it redacts, and what it doesn't. RedactPII redacts what the model is sent and what it returns. In conversation memory:

  • The user's messages are stored as they said them. Since 1.81.0 before_model redacts copies, so history already stored is never rewritten. (Before, it edited the stored messages in place, and turn 1 turned into [REDACTED] in memory once turn 2 ran.)
  • The model's reply is stored redacted. Memory records the agent's answer after all middleware — the same text run returns — with run and astream alike. With astream the caller has already received the raw deltas as they streamed; memory still gets the redacted reply. (Before 1.82.0, astream stored the raw text.)

To keep PII out of memory entirely, redact the input before it reaches the agent.

patterns= is unchanged: your regexes, applied verbatim, with no Luhn opinion and without consulting the shared detector. Non-string content (a multimodal message's list[ContentPart]) is now passed through untouched rather than raising.

Writing Your Own Middleware

A TokenCounter middleware that tallies tokens across a run:

from fastaiagent import Agent, AgentMiddleware, LLMClient

class TokenCounter(AgentMiddleware):
    name = "token_counter"

    async def after_model(self, ctx, response):
        used = response.usage.get("total_tokens", 0)
        ctx.scratch["tokens_total"] = ctx.scratch.get("tokens_total", 0) + used
        return response

agent = Agent(
    name="counter",
    llm=LLMClient(provider="openai", model="gpt-4.1"),
    middleware=[TokenCounter()],
)
result = await agent.arun("hello")
# Access via a second middleware, or log from after_model.

A PromptCache that returns a cached response when the message list has been seen before:

import hashlib
import json
from fastaiagent import AgentMiddleware
from fastaiagent.llm.client import LLMResponse


class PromptCache(AgentMiddleware):
    name = "prompt_cache"

    def __init__(self):
        self._cache: dict[str, LLMResponse] = {}

    def _key(self, messages):
        payload = [m.to_openai_format() for m in messages]
        return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()

    async def before_model(self, ctx, messages):
        ctx.scratch["cache_key"] = self._key(messages)
        return messages

    async def after_model(self, ctx, response):
        self._cache[ctx.scratch["cache_key"]] = response
        return response

Short-Circuiting Tools

A middleware can return a ToolResult without calling call_next:

from fastaiagent.tool.base import ToolResult

class CacheTool(AgentMiddleware):
    async def wrap_tool(self, ctx, tool, args, call_next):
        cache_key = (tool.name, tuple(sorted(args.items())))
        if cache_key in ctx.scratch.get("tool_cache", {}):
            return ToolResult(output=ctx.scratch["tool_cache"][cache_key])
        result = await call_next(tool, args)
        ctx.scratch.setdefault("tool_cache", {})[cache_key] = result.output
        return result

Interaction with Tracing

Middleware hooks are called inside the agent's root span. Custom spans created inside a middleware hook nest under it — see Tracing for the trace_context helper.

Next Steps