pydantic-ai-middleware — Hooks & Guardrails

Vstorm Open Source

# pydantic-ai-
middleware

Hooks, guardrails, and lifecycle middleware for PydanticAI agents.

#### Version

0.2.x

#### Hooks

7 lifecycle events

#### Stack

PydanticAI + Python 3.10+

## What is pydantic-ai-middleware?

ASGI-style middleware for AI agents — intercept, modify, or block at every stage.

- 7 lifecycle hooks — before/after for run, model request, and tool calls, plus error handlers
- Composable pipelines — stack multiple middleware, each one wraps the next
- Tool permissions — ALLOW / DENY / ASK decisions with structured results
- Context sharing — scoped, access-controlled data between hooks
- Cost tracking — built-in token & USD monitoring with budgets
- Composable — chains, conditionals, decorator syntax

Repository

    # Install
    pip install pydantic-ai-middleware

    # Or with uv
    uv add pydantic-ai-middleware

    from pydantic_ai import Agent
    from pydantic_ai_middleware import (
        AgentMiddleware, MiddlewareAgent
    )

    class MyMiddleware(AgentMiddleware[None]):
        async def before_run(self, prompt, deps, ctx):
            print(f"Prompt: {prompt}")
            return prompt

    agent = Agent("openai:gpt-4.1")
    mw_agent = MiddlewareAgent(
        agent=agent,
        middleware=[MyMiddleware()],
    )
    result = await mw_agent.run("Hello!")

## Lifecycle Pipeline

Every agent run passes through 7 hooks in a defined order.

1. before_run → 2. before_model_request → 3. before_tool_call → 4. on_tool_error → 5. after_tool_call → 6. after_run → 7. on_error

### before_* hooks

Intercept before the agent acts.
Can modify input or block execution.

### after_* hooks

Intercept after the agent acts.
Can modify output or log results.

## Hook: before_run

First hook — intercept the user prompt before the agent starts.

Entry Point

    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?"
    )

### Console Output

$ python 01_before_run.py
[before_run] Received prompt: 'What is the capital of France?'
[result] The capital of France is Paris.
Use Cases

- Log incoming prompts
- Sanitize user input
- Block harmful content (`raise InputBlocked`)

## Hook: after_run

Last hook — inspect or modify the final output. Runs in REVERSE order.

Exit Point

    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

### Console Output

$ python 02_after_run.py
[before_run] Timer started
[after_run] Output: '2 + 2 equals 4.'
[after_run] Elapsed: 0.674s
[result] 2 + 2 equals 4.
Use Cases

- Measure response time
- Filter or transform output
- Send analytics events

## Hooks: before & after tool_call

Intercept every tool invocation — log, modify args, measure timing.

Tool Layer

    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)

    class AfterToolCallLogger(AgentMiddleware[None]):

        async def after_tool_call(
            self, tool_name, tool_args, result, deps, ctx
        ):
            print(f"[after] {tool_name} → {result!r}")
            return result

    toolset = FunctionToolset()

    @toolset.tool
    def multiply(a: int, b: int) -> str:
        """Multiply two numbers."""
        return str(a * b)

### Console Output

$ python 04_before_tool_call.py
[before_tool_call] Tool: add
[before_tool_call] Args: {'a': 17, 'b': 25}
[tool:add] Computing 17 + 25 = 42
[result] 17 + 25 = 42
$ python 05_after_tool_call.py
[tool:multiply] 12 * 8 = 96
[after_tool_call] Tool: multiply
[after_tool_call] Result: '96'
[after_tool_call] Took: 0.0007s
[result] 12 times 8 is 96.

## Error Hooks: on_tool_error & on_error

Handle failures gracefully — tool-specific context or global catch-all.

Safety

    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

    async def on_error(
        self, error, deps, ctx
    ):
        print(f"[on_error] {type(error).__name__}")
        # Return None to re-raise original
        # Return new Exception to replace
        return None

    # Catches ANY exception during the run
    # Global fallback — no tool context

### Console Output

$ python 06_on_tool_error.py
[result] Division by zero is undefined in mathematics, so 10 divided by 0 does not have a valid answer.
$ python 07_on_error.py
[on_error] Caught: InputBlocked: Prompt contains blocked content
[main] Blocked: Prompt contains blocked content
[result] Python is a high-level, interpreted programming language known for its simplicity, readability, and versatility.

## Tool Permissions: ALLOW / DENY / ASK

Structured permission decisions — block sensitive tools or defer to a handler.

    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
        )

    async def before_tool_call(
        self, tool_name, tool_args, deps, ctx
    ):
        if tool_name == "delete_file":
            return ToolPermissionResult(
                decision=ToolDecision.ASK,
                reason="Destructive — needs approval",
            )
        return tool_args

    # ASK defers to permission_handler callback:
    mw_agent = MiddlewareAgent(
        agent=agent,
        middleware=[guard],
        permission_handler=my_handler, # async (name, args, reason) → bool
    )

### Console Output

$ python 08_tool_blocking.py
--- Safe path ---
[PathGuard] ALLOWED: read_file('/home/user/notes.txt')
[tool:read_file] Reading /home/user/notes.txt
[result] Contents of /home/user/notes.txt: [mock data]
--- Blocked path ---
[PathGuard] DENIED: read_file('/etc/passwd')
[result] Access to /etc/passwd was denied.
ALLOW ✓
DENY ✗
ASK ?

## Decorator Syntax & tool_names Filtering

Lightweight alternatives — no class needed. Filter by tool name.

    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]

    class EmailGuard(AgentMiddleware[None]):
        # Only fires for these tools:
        tool_names = {"send_email", "draft_email"}

        async def before_tool_call(
            self, tool_name, tool_args, deps, ctx
        ):
            if not tool_args.get("to"):
                raise ToolBlocked(
                    tool_name, "Recipient required!"
                )
            return tool_args

    # get_weather → guard skipped
    # send_email  → guard fires

### Console Output

$ python 10_tool_name_filtering.py
--- Weather (no guard) ---
[tool:get_weather] City: Warsaw
[result] The weather in Warsaw is 22°C and sunny.
--- Email (guard active) ---
[EmailOnlyGuard] Checking send_email: {'to': 'jan@example.com', 'subject': 'Hello', 'body': 'Hi!'}
[tool:send_email] To: jan@example.com, Subject: Hello
[result] The email has been sent to jan@example.com.

## Context Sharing & Cost Tracking

Share data between hooks. Track tokens and USD costs with budgets.

    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

    from pydantic_ai_middleware.cost_tracking import (
        CostTrackingMiddleware, CostInfo
    )

    cost_mw = CostTrackingMiddleware(
        model_name="openai:gpt-4.1",
        budget_limit_usd=1.00,
        on_cost_update=lambda info: print(
            f"Run #{info.run_count}: ${info.run_cost_usd:.6f}"
        ),
    )

    # After runs:
    print(cost_mw.total_cost)  # $0.001234
    print(cost_mw.total_request_tokens)  # 1523
    # Raises BudgetExceededError if over budget

### Console Output

$ python 14_cost_tracking.py
=== Run 1 ===
[cost] Run #1
[cost] Tokens: 20 in / 32 out
[cost] Run cost: $0.000296
[cost] Total cost: $0.000296
[result] Python is a high-level, interpreted programming language...
=== Run 2 ===
[cost] Run #2
[cost] Tokens: 20 in / 25 out
[cost] Run cost: $0.000240
[cost] Total cost: $0.000536
[result] Rust is a systems programming language...
=== Cumulative Stats ===
Total runs: 2
Total tokens: 40 in / 57 out
Total cost: $0.000536

## Full Lifecycle Trace

All hooks in one middleware — see the complete execution flow.

    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

### Console Output

$ python 16_full_lifecycle.py
=== Full Lifecycle Trace ===
1. before_run(prompt='What is 7 * 8 + 2?')
2. before_model_request(messages=1)
3. before_tool_call(tool=calculate, args={'expression': '7 * 8 + 2'})
4. after_tool_call(tool=calculate, result='58')
5. before_model_request(messages=3)
6. after_run(output='7 * 8 + 2 = 58.')
[result] 7 * 8 + 2 = 58.
=== Total events: 6 ===

## Real-World: pydantic-deep

How middleware powers a production AI agent framework — cost tracking, context management, permissions, and hooks.

Production

    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],
    )

    # Middleware assembly order:
    # 1. User middleware (audit_mw, permission_mw)
    # 2. HooksMiddleware (from hooks=[])
    # 3. CheckpointMiddleware (if enabled)
    # 4. ContextManagerMiddleware (default ON)
    # 5. CostTrackingMiddleware (default ON)

    # Auto-wraps in MiddlewareAgent when
    # any middleware is active:
    if middleware or cost_tracking or hooks:
        return MiddlewareAgent(
            agent=agent,
            middleware=all_middleware,
            context=middleware_context,
            permission_handler=permission_handler,
        )

### Features Enabled by Default

Cost CostTrackingMiddleware

Token & USD tracking, budget limits, per-run callbacks

Context ContextManagerMiddleware

Auto-compression at 90% of 200K token budget

Hooks HooksMiddleware

Claude Code-style shell/Python hooks with matchers

Custom Your AgentMiddleware

Audit, permissions, logging — compose freely

## Real-World: Path Permissions

Production middleware from pydantic-deep — regex-based path blocking for file tools.

Security

    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

### How it works

# Agent tries to read /etc/passwd:
[PermissionMiddleware] before_tool_call(read_file)
[PermissionMiddleware] Path matches: /etc/passwd
[PermissionMiddleware] Decision: DENY
# Agent gets ToolBlocked and must explain:
[result] I can't read that file — access denied.
Three decisions

- ALLOW — let the tool run (optionally with modified args)
- DENY — block with a reason (agent sees ToolBlocked)
- ASK — defer to a permission_handler callback

## Real-World: Claude Code-Style Hooks

Python handler hooks — block dangerous commands, log tool calls in the background.

Hooks

    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)

    async def audit_logger_handler(
        hook_input: HookInput,
    ) -> HookResult:
        """POST_TOOL_USE: logs all tool calls.
        Runs in background (fire-and-forget)."""

        args = str(hook_input.tool_input)[:200]
        logger.info(
            f"AUDIT: {hook_input.tool_name}({args})"
        )

        return HookResult(allow=True)

    from pydantic_deep import Hook, HookEvent

    HOOKS = [
        # Background audit — fires after every tool
        Hook(
            event=HookEvent.POST_TOOL_USE,
            handler=audit_logger_handler,
            background=True,
        ),
        # Safety gate — blocks dangerous 'execute' calls
        Hook(
            event=HookEvent.PRE_TOOL_USE,
            handler=safety_gate_handler,
            matcher="execute",  # only 'execute' tool
            timeout=5,
        ),
    ]

    agent = create_deep_agent(hooks=HOOKS)

### Hook Features

PRE_TOOL_USE

Runs before tool execution. Return `allow=False` to block. Supports `matcher` regex to target specific tools.

POST_TOOL_USE

Runs after tool completes. Can modify result. Set `background=True` for fire-and-forget logging.

POST_TOOL_USE_FAILURE

Runs when a tool raises an exception. Access error details via `hook_input.tool_error`.

Hooks can also be shell commands (exit 0 = allow, exit 2 = deny), following Claude Code conventions.

## Examples Reference

16 standalone examples covering every feature.

# | File | Hook / Feature | Description
--- | --- | --- | ---
01 | 01_before_run.py | before_run | Log incoming prompts
02 | 02_after_run.py | after_run | Measure wall-clock time
03 | 03_before_model_request.py | before_model_request | Count LLM API calls
04 | 04_before_tool_call.py | before_tool_call | Log tool invocations
05 | 05_after_tool_call.py | after_tool_call | Log results + timing
06 | 06_on_tool_error.py | on_tool_error | Handle tool failures
07 | 07_on_error.py | on_error | Global error handler
08 | 08_tool_blocking.py | ToolPermissionResult | ALLOW / DENY decisions
09 | 09_permission_handler.py | ToolDecision.ASK | Interactive approval flow
10 | 10_tool_name_filtering.py | tool_names | Per-tool middleware
11 | 11_middleware_chain.py | MiddlewareChain | Compose middleware
12 | 12_decorator_middleware.py | @before_run | Decorator syntax
13 | 13_context_sharing.py | MiddlewareContext | Share data between hooks
14 | 14_cost_tracking.py | CostTrackingMiddleware | Token & USD tracking
15 | 15_multiple_middleware.py | Execution order | Multiple middleware pipeline
16 | 16_full_lifecycle.py | All hooks | Complete lifecycle trace

## Get Started

pip install pydantic-ai-middleware

Guardrails

Block harmful inputs

Tool Control

ALLOW / DENY / ASK

Cost Tracking

Tokens & USD budgets

Composable

Chains & conditionals

vstorm-co/pydantic-ai-middleware
Vstorm Open Source 2026