pydantic-ai-
middleware

# Install
pip install pydantic-ai-middleware

# Or with uv
uv add pydantic-ai-middleware

class BeforeRunLogger(AgentMiddleware[None]):

    async def before_run(self, prompt, deps, ctx):
        print(f"[before_run] {prompt!r}")
        return prompt

agent = Agent("openai:gpt-4.1")
mw_agent = MiddlewareAgent(
    agent=agent,
    middleware=[BeforeRunLogger()],
)

result = await mw_agent.run(
    "What is the capital of France?"
)

class TimingMiddleware(AgentMiddleware[None]):

    def __init__(self):
        self._start = 0

    async def before_run(self, prompt, deps, ctx):
        self._start = time.perf_counter()
        return prompt

    async def after_run(self, prompt, output, deps, ctx):
        elapsed = time.perf_counter() - self._start
        print(f"[after_run] {output!r}")
        print(f"[after_run] Elapsed: {elapsed:.3f}s")
        return output

class BeforeToolCallLogger(AgentMiddleware[None]):

    async def before_tool_call(
        self, tool_name, tool_args, deps, ctx
    ):
        print(f"[before] {tool_name}({tool_args})")
        return tool_args

toolset = FunctionToolset()

@toolset.tool
def add(a: int, b: int) -> str:
    """Add two numbers."""
    return str(a + b)

async def on_tool_error(
    self, tool_name, tool_args, error, deps, ctx
):
    print(f"[on_tool_error] {tool_name} failed!")
    print(f"  Error: {error}")
    # Replace with user-friendly error
    return RuntimeError(
        f"'{tool_name}' is unavailable"
    )

# Provides tool_name + tool_args context
# that on_error does NOT have

from pydantic_ai_middleware import (
    ToolDecision, ToolPermissionResult
)

async def before_tool_call(
    self, tool_name, tool_args, deps, ctx
):
    path = tool_args.get("path", "")
    if ".env" in path:
        return ToolPermissionResult(
            decision=ToolDecision.DENY,
            reason="Access to .env blocked",
        )
    return ToolPermissionResult(
        decision=ToolDecision.ALLOW
    )

from pydantic_ai_middleware.decorators import (
    before_run, after_run,
    before_tool_call, after_tool_call,
    before_model_request,
    on_error, on_tool_error,
)

@before_run
async def log_prompt(prompt, deps, ctx):
    print(f"Prompt: {prompt}")
    return prompt

@after_tool_call
async def log_result(name, args, result, deps, ctx):
    print(f"{name} → {result}")
    return result

# Use as regular middleware:
middleware=[log_prompt, log_result]

from pydantic_ai_middleware import MiddlewareContext
from pydantic_ai_middleware.context import HookType

async def before_run(self, prompt, deps, ctx):
    ctx.set("start_time", time.perf_counter())
    return prompt

async def after_run(self, prompt, output, deps, ctx):
    start = ctx.get_from(
        HookType.BEFORE_RUN, "start_time"
    )
    print(f"Elapsed: {time.perf_counter()-start:.3f}s")
    return output

# Access control: each hook can only
# read from earlier hooks

class FullLifecycleLogger(AgentMiddleware[None]):
    async def before_run(self, prompt, deps, ctx):
        self._log("before_run")
        return prompt
    async def before_model_request(self, msgs, deps, ctx):
        self._log("before_model_request")
        return msgs
    async def before_tool_call(self, name, args, deps, ctx):
        self._log(f"before_tool_call({name})")
        return args
    async def after_tool_call(self, name, args, res, deps, ctx):
        self._log(f"after_tool_call({name})")
        return res
    async def after_run(self, prompt, output, deps, ctx):
        self._log("after_run")
        return output

# + on_error, on_tool_error

from pydantic_deep import create_deep_agent

agent = create_deep_agent(
    model="openai:gpt-4.1",

    # Cost tracking (enabled by default)
    cost_tracking=True,
    cost_budget_usd=5.00,
    on_cost_update=my_cost_callback,

    # Context manager (enabled by default)
    context_manager=True,
    context_manager_max_tokens=200_000,
    on_context_update=my_ctx_callback,

    # Custom middleware
    middleware=[audit_mw, permission_mw],

    # Claude Code-style hooks
    hooks=[safety_gate, audit_logger],
)

BLOCKED_PATTERNS = [
    r"/etc/passwd", r"\.env$",
    r"/root/", r"\.ssh/",
    r"id_rsa", r"/proc/",
]
FILE_TOOLS = {"read_file", "write_file", "edit_file"}

class PermissionMiddleware(AgentMiddleware[DeepAgentDeps]):

    async def before_tool_call(self, tool_name, tool_args, deps, ctx):
        if tool_name not in FILE_TOOLS:
            return tool_args

        path = tool_args.get("path", "")
        for pattern in BLOCKED_PATTERNS:
            if re.search(pattern, path):
                return ToolPermissionResult(
                    decision=ToolDecision.DENY,
                    reason=f"Blocked: {pattern}",
                )
        return tool_args

async def safety_gate_handler(hook_input: HookInput) -> HookResult:
    """PRE_TOOL_USE: blocks dangerous commands."""
    command = hook_input.tool_input.get("command", "")

    dangerous = [
        r"rm\s+-rf\s+/",
        r"mkfs\.",
        r"dd\s+if=.*of=/dev/",
        r":\(\)\{",  # fork bomb
    ]

    for pattern in dangerous:
        if re.search(pattern, command):
            return HookResult(
                allow=False,
                reason=f"BLOCKED: {command}",
            )

    return HookResult(allow=True)