Source code for edri.integrations.pydantic_ai.mcp

"""Auditing the one MCP operation Pydantic AI exposes a public seam for.

Of the MCP protocol operations an agent performs, exactly one can be observed
without reaching into private internals on the supported dependency floor:
``tools/call``, through ``MCPToolset(process_tool_call=...)``. ``initialize``,
``tools/list``, ``resources/*`` and ``prompts/*`` round-trip correctly but offer
no public observation point, so this module makes no claim about them -- and
``MCPToolset.direct_call_tool()`` bypasses the seam entirely, so an application
calling a toolset directly is outside the observed set.

Installation is explicit, in the agent's own ``build_agent()``::

    from pydantic_ai.mcp import MCPToolset
    from edri.integrations.pydantic_ai import audited_tool_calls

    class WeatherAgent(AgentBase[AgentDeps, str]):
        def build_agent(self):
            toolset = MCPToolset(
                client, id="weather",
                process_tool_call=audited_tool_calls("weather"),
            )
            return Agent(self.model_name, deps_type=AgentDeps, toolsets=[toolset])

Explicit because the toolset is the application's object: EDRI does not reach
into an agent's toolsets to instrument them, and an integration that silently
replaced a user's own ``process_tool_call`` would be worse than one that asks
for a line of wiring.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any
from uuid import uuid4

from edri.integrations.pydantic_ai.lifecycle import (
    _current_run,
    _current_tool_action,
    _observed,
)
from edri.integrations.pydantic_ai.observability import (
    AgentActionNamespace,
    AgentActionOperation,
    _Capture,
)

if TYPE_CHECKING:  # Type-only: importing it at runtime would require the extra.
    from pydantic_ai.mcp import ProcessToolCallback

__all__ = ["audited_tool_calls"]




def audited_tool_calls(server: str | None = None) -> ProcessToolCallback:
    """A ``process_tool_call`` callback that publishes one action per call.

    ``server`` labels the server in the action stream. The seam is handed only
    the run context, the transport callable, a tool name and its arguments --
    never the connection -- so the label cannot be discovered here; pass the
    same name the toolset carries as its ``id``.

    The returned callback is transparent. It awaits the real call, returns its
    result unchanged and re-raises its exceptions unchanged, including the
    ``ModelRetry`` that a server-side failure arrives as. A call made while no
    observed run is in progress -- a shared connection's own traffic -- is
    passed straight through: its attribution is unknown, and EDRI publishes
    nothing rather than guessing a run for it.

    The return type is the real ``ProcessToolCallback``, imported only under
    ``TYPE_CHECKING``: ``pydantic_ai.mcp`` raises ``ImportError`` unless the
    ``mcp`` extra is installed, and this package must import without it.
    """

    async def audited(
        ctx: Any, call_tool: Any, name: str, args: dict[str, Any]
    ) -> Any:
        run = _current_run()
        if run is None:
            return await call_tool(name, args)

        policy = run.policy
        return await _observed(
            run.emitter,
            lambda: call_tool(name, args),
            shared={
                "action_id": uuid4(),
                "namespace": AgentActionNamespace.MCP,
                "operation": AgentActionOperation.MCP_TOOLS_CALL,
                # The tool execution that made this call, when there is one: an
                # MCP call and the agent tool containing it are distinct
                # operations, and an inspector needs to see both.
                "parent_action_id": (
                    _current_tool_action(run.state.root_action_id)
                    or run.state.root_action_id
                ),
                "root_action_id": run.state.root_action_id,
                "run_id": run.state.run_id,
                "conversation_id": run.state.conversation_id,
                "name": name,
                "metadata": {"server": server},
            },
            request=_Capture(args, policy.include_mcp_requests),
            terminal=lambda result: {
                "response_data": _Capture(result, policy.include_mcp_responses)
            },
        )

    return audited