Add agent payment instrumentation hook

.changelog/agent-instrument.md

@@ -0,0 +1,5 @@
+---
+pympp: minor
+---
+
+Added opt-in process-wide instrumentation for payment-aware `httpx.AsyncClient` requests and MCP `ClientSession.call_tool` calls.

src/mpp/agent.py

@@ -0,0 +1,162 @@
+"""Opt-in instrumentation for payment-aware HTTP and MCP agent runtimes."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Literal
+
+import httpx
+
+from mpp.events import EventDispatcher
+from mpp.runtime import Method, PaymentRuntime
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+@dataclass(slots=True)
+class InstrumentationHandle:
+    """Handle returned by instrument(), used to disable runtime instrumentation."""
+
+    runtime: PaymentRuntime
+    _httpx_enabled: bool
+    _mcp_enabled: bool
+    _disabled: bool = False
+
+    def disable(self) -> None:
+        """Disable this instrumentation handle and restore patches when possible."""
+        if self._disabled:
+            return
+        self._disabled = True
+        if self._httpx_enabled and self.runtime in _state.httpx_runtimes:
+            _state.httpx_runtimes.remove(self.runtime)
+        if self._mcp_enabled and self.runtime in _state.mcp_runtimes:
+            _state.mcp_runtimes.remove(self.runtime)
+        _restore_unused_patches()
+
+    def __enter__(self) -> InstrumentationHandle:
+        return self
+
+    def __exit__(self, *_args: Any) -> None:
+        self.disable()
+
+
+def instrument(
+    methods: Sequence[Method],
+    *,
+    httpx: bool = True,
+    mcp: Literal["auto"] | bool = "auto",
+    allowed_origins: Sequence[str] | None = None,
+    events: EventDispatcher | None = None,
+) -> InstrumentationHandle:
+    """Instrument common Python agent runtime boundaries for payments.
+
+    Patches are process-wide and opt-in. Existing and future httpx.AsyncClient
+    instances use the latest active runtime; MCP instrumentation is installed
+    only when the official Python MCP SDK is importable, unless mcp=True is
+    passed, in which case a missing SDK raises ImportError.
+    """
+    runtime = PaymentRuntime(methods, events=events, allowed_origins=allowed_origins)
+
+    httpx_enabled = False
+    if httpx:
+        _install_httpx_patch()
+        _state.httpx_runtimes.append(runtime)
+        httpx_enabled = True
+
+    mcp_enabled = False
+    if mcp is True or mcp == "auto":
+        installed = _install_mcp_patch(required=mcp is True)
+        if installed:
+            _state.mcp_runtimes.append(runtime)
+            mcp_enabled = True
+
+    return InstrumentationHandle(
+        runtime=runtime,
+        _httpx_enabled=httpx_enabled,
+        _mcp_enabled=mcp_enabled,
+    )
+
+
+class _InstrumentationState:
+    def __init__(self) -> None:
+        self.httpx_runtimes: list[PaymentRuntime] = []
+        self.mcp_runtimes: list[PaymentRuntime] = []
+        self.original_httpx_send: Any | None = None
+        self.original_mcp_call_tool: Any | None = None
+        self.mcp_client_session: Any | None = None
+
+
+_state = _InstrumentationState()
+
+
+def _install_httpx_patch() -> None:
+    if _state.original_httpx_send is not None:
+        return
+    original_send = httpx.AsyncClient.send
+    _state.original_httpx_send = original_send
+
+    async def send(
+        self: httpx.AsyncClient,
+        request: httpx.Request,
+        *args: Any,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        runtime = getattr(self, "_mpp_payment_runtime", None)
+        if runtime is None and _state.httpx_runtimes:
+            runtime = _state.httpx_runtimes[-1]
+        if runtime is None:
+            return await original_send(self, request, *args, **kwargs)
+        bound_send = MethodType(original_send, self)
+        return await runtime.send_httpx(bound_send, request, *args, **kwargs)
+
+    httpx.AsyncClient.send = send  # type: ignore[method-assign]
+
+
+def _install_mcp_patch(*, required: bool) -> bool:
+    if _state.original_mcp_call_tool is not None:
+        return True
+    try:
+        from mcp import ClientSession
+    except ImportError as exc:
+        if required:
+            raise ImportError(
+                'Cannot instrument MCP calls. Install the "mcp" extra: pip install "pympp[mcp]"'
+            ) from exc
+        return False
+
+    original_call_tool = ClientSession.call_tool
+    _state.original_mcp_call_tool = original_call_tool
+    _state.mcp_client_session = ClientSession
+
+    async def call_tool(
+        self: Any,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        runtime = _state.mcp_runtimes[-1] if _state.mcp_runtimes else None
+        if runtime is None:
+            return await original_call_tool(self, name, arguments, *args, **kwargs)
+        bound_call_tool = MethodType(original_call_tool, self)
+        return await runtime.call_mcp_tool(bound_call_tool, name, arguments, *args, **kwargs)
+
+    ClientSession.call_tool = call_tool  # type: ignore[method-assign]
+    return True
+
+
+def _restore_unused_patches() -> None:
+    if not _state.httpx_runtimes and _state.original_httpx_send is not None:
+        httpx.AsyncClient.send = _state.original_httpx_send  # type: ignore[method-assign]
+        _state.original_httpx_send = None
+
+    if (
+        not _state.mcp_runtimes
+        and _state.original_mcp_call_tool is not None
+        and _state.mcp_client_session is not None
+    ):
+        _state.mcp_client_session.call_tool = _state.original_mcp_call_tool
+        _state.original_mcp_call_tool = None
+        _state.mcp_client_session = None

tests/test_agent.py

@@ -12,6 +12,7 @@
 import pytest
 
 from mpp import Challenge, ChallengeEcho, Credential
+from mpp.agent import instrument
 from mpp.extensions.mcp import META_CREDENTIAL
 from mpp.runtime import PaymentRuntime
 
@@ -120,6 +121,26 @@ def mcp_payment_error() -> FakeMcpError:
 
 
 class TestHttpxInstrumentation:
+    @pytest.mark.asyncio
+    async def test_existing_async_client_pays_after_global_hook(self) -> None:
+        inner = MockTransport(
+            [
+                httpx.Response(402, headers={"www-authenticate": payment_challenge_header()}),
+                httpx.Response(200, json={"ok": True}),
+            ]
+        )
+        client = httpx.AsyncClient(transport=inner)
+        handle = instrument([MockMethod()], mcp=False)
+        try:
+            response = await client.get("https://example.com/paid")
+        finally:
+            await client.aclose()
+            handle.disable()
+
+        assert response.status_code == 200
+        assert len(inner.requests) == 2
+        assert inner.requests[1].headers["authorization"].startswith("Payment ")
+
     @pytest.mark.asyncio
     async def test_runtime_wrap_async_client_without_global_hook(self) -> None:
         inner = MockTransport(
@@ -161,8 +182,50 @@ async def test_disallowed_http_origin_does_not_retry_payment(self) -> None:
         assert len(inner.requests) == 1
         method.create_credential.assert_not_called()
 
+    def test_httpx_patch_is_idempotent_and_restored(self) -> None:
+        original_send = httpx.AsyncClient.send
+        first = instrument([], mcp=False)
+        second = instrument([], mcp=False)
+        try:
+            patched_send = httpx.AsyncClient.send
+            assert patched_send is not original_send
+            assert httpx.AsyncClient.send is patched_send
+            second.disable()
+            assert httpx.AsyncClient.send is patched_send
+        finally:
+            first.disable()
+            second.disable()
+
+        assert httpx.AsyncClient.send is original_send
+
 
 class TestMcpInstrumentation:
+    @pytest.mark.asyncio
+    async def test_mcp_call_tool_pays_after_global_hook(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        fake = install_fake_mcp(monkeypatch)
+        raw_result = FakeCallToolResult("paid")
+        session = fake.client_session([mcp_payment_error(), raw_result])
+
+        handle = instrument([MockMethod()], httpx=False, mcp=True)
+        try:
+            result = await session.call_tool(
+                "premium_tool",
+                {"query": "test"},
+                progress_callback="callback",
+                meta={"trace_id": "abc"},
+            )
+        finally:
+            handle.disable()
+
+        assert result is raw_result
+        assert len(session.calls) == 2
+        retry_kwargs = session.calls[1][3]
+        assert retry_kwargs["progress_callback"] == "callback"
+        assert retry_kwargs["meta"]["trace_id"] == "abc"
+        assert META_CREDENTIAL in retry_kwargs["meta"]
+
     @pytest.mark.asyncio
     async def test_payment_runtime_call_mcp_tool_pays_and_preserves_raw_result(
         self, monkeypatch: pytest.MonkeyPatch
@@ -205,6 +268,39 @@ async def test_disallowed_mcp_realm_does_not_retry_payment(
         assert len(session.calls) == 1
         method.create_credential.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_mcp_patch_is_idempotent_and_restored(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        fake = install_fake_mcp(monkeypatch)
+        original_call_tool = fake.client_session.call_tool
+        first = instrument([MockMethod()], httpx=False, mcp=True)
+        second = instrument([MockMethod()], httpx=False, mcp=True)
+        try:
+            patched_call_tool = fake.client_session.call_tool
+            assert patched_call_tool is not original_call_tool
+            assert fake.client_session.call_tool is patched_call_tool
+
+            session = fake.client_session([mcp_payment_error(), FakeCallToolResult("paid")])
+            result = await session.call_tool("premium_tool", {})
+            assert result.content[0]["text"] == "paid"
+            assert len(session.calls) == 2
+
+            second.disable()
+            assert fake.client_session.call_tool is patched_call_tool
+
+            session_after_one_disable = fake.client_session(
+                [mcp_payment_error(), FakeCallToolResult("still-paid")]
+            )
+            result_after_one_disable = await session_after_one_disable.call_tool("premium_tool", {})
+            assert result_after_one_disable.content[0]["text"] == "still-paid"
+            assert len(session_after_one_disable.calls) == 2
+        finally:
+            first.disable()
+            second.disable()
+
+        assert fake.client_session.call_tool is original_call_tool
+
     @pytest.mark.asyncio
     async def test_payment_runtime_call_mcp_tool_preserves_raw_result(
         self, monkeypatch: pytest.MonkeyPatch
@@ -218,3 +314,9 @@ async def test_payment_runtime_call_mcp_tool_preserves_raw_result(
 
         assert result is raw_result
         assert len(session.calls) == 2
+
+    def test_mcp_true_raises_when_mcp_sdk_is_missing(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setitem(sys.modules, "mcp", None)
+
+        with pytest.raises(ImportError, match="pympp\\[mcp\\]"):
+            instrument([], httpx=False, mcp=True)