Add scoped payment instrumentation

.changelog/agent-instrument.md

@@ -0,0 +1,5 @@
+---
+pympp: minor
+---
+
+Added scoped instrumentation for payment-aware sync/async httpx and MCP client calls.

src/mpp/agent.py

@@ -0,0 +1,287 @@
+"""Scoped instrumentation for payment-aware HTTP and MCP calls."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from contextvars import ContextVar
+from dataclasses import dataclass
+from types import MethodType
+from typing import Any, Literal
+
+import httpx
+
+from mpp.runtime import (
+    PaymentRuntime,
+    mcp_payment_flow_active,
+    payment_flow_active,
+    payment_flow_active_in_process,
+)
+
+
+@dataclass(eq=False, slots=True)
+class _Binding:
+    runtime: PaymentRuntime
+    httpx: bool
+    mcp: bool
+    active: bool = True
+
+
+_bindings: ContextVar[tuple[_Binding, ...] | None] = ContextVar(
+    "mpp_instrumentation_bindings",
+    default=None,
+)
+_httpx_active: ContextVar[bool] = ContextVar("mpp_httpx_instrumentation_active", default=False)
+_mcp_active: ContextVar[bool] = ContextVar("mpp_mcp_instrumentation_active", default=False)
+
+
+@dataclass(slots=True)
+class InstrumentationHandle:
+    """Handle returned by :func:`instrument`."""
+
+    runtime: PaymentRuntime
+    _binding: _Binding
+
+    def disable(self) -> None:
+        """Disable this binding and restore unused process patches safely."""
+        binding = self._binding
+        with _state.lock:
+            if not binding.active:
+                return
+            binding.active = False
+            _state.bindings = [item for item in _state.bindings if item is not binding]
+            _restore_unused_patches()
+
+        local = _bindings.get()
+        if local is not None:
+            _bindings.set(tuple(item for item in local if item is not binding))
+
+    def __enter__(self) -> InstrumentationHandle:
+        return self
+
+    def __exit__(self, *_args: Any) -> None:
+        self.disable()
+
+
+def instrument(
+    runtime: PaymentRuntime,
+    *,
+    httpx: bool = True,
+    mcp: Literal["auto"] | bool = "auto",
+) -> InstrumentationHandle:
+    """Make common Python HTTP and MCP client boundaries payment-aware.
+
+    Selection is context-local when instrumentation is installed in an async
+    task or request context. A bare thread uses the process fallback only when
+    exactly one runtime is active, which supports harness worker threads without
+    choosing between multiple wallets.
+    """
+    client_session = _resolve_mcp_client(required=mcp is True) if mcp is not False else None
+    binding = _Binding(runtime=runtime, httpx=httpx, mcp=client_session is not None)
+
+    with _state.lock:
+        try:
+            if httpx:
+                _install_httpx_patches()
+            if client_session is not None:
+                _install_mcp_patch(client_session)
+        except BaseException:
+            _restore_unused_patches()
+            raise
+        _state.bindings.append(binding)
+
+    local = _bindings.get()
+    _bindings.set((*(() if local is None else local), binding))
+    return InstrumentationHandle(runtime=runtime, _binding=binding)
+
+
+class _InstrumentationState:
+    def __init__(self) -> None:
+        self.lock = threading.RLock()
+        self.bindings: list[_Binding] = []
+        self.original_sync_send: Any | None = None
+        self.sync_send_patch: Any | None = None
+        self.original_async_send: Any | None = None
+        self.async_send_patch: Any | None = None
+        self.original_mcp_call_tool: Any | None = None
+        self.mcp_call_tool_patch: Any | None = None
+        self.mcp_client_session: Any | None = None
+
+
+_state = _InstrumentationState()
+
+
+def _select_runtime(protocol: Literal["httpx", "mcp"]) -> PaymentRuntime | None:
+    local = _bindings.get()
+    if local is not None:
+        for binding in reversed(local):
+            if binding.active and getattr(binding, protocol):
+                return binding.runtime
+        return None
+
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        pass
+    else:
+        return None
+
+    if payment_flow_active_in_process():
+        return None
+
+    with _state.lock:
+        runtimes: list[PaymentRuntime] = []
+        for binding in _state.bindings:
+            if not binding.active or not getattr(binding, protocol):
+                continue
+            if all(runtime is not binding.runtime for runtime in runtimes):
+                runtimes.append(binding.runtime)
+        return runtimes[0] if len(runtimes) == 1 else None
+
+
+def _install_httpx_patches() -> None:
+    if _state.original_sync_send is None:
+        original_sync_send = httpx.Client.send
+
+        def sync_send(
+            self: httpx.Client,
+            request: httpx.Request,
+            *args: Any,
+            **kwargs: Any,
+        ) -> httpx.Response:
+            if (
+                getattr(self, "_mpp_payment_wrapped", False)
+                or payment_flow_active()
+                or _httpx_active.get()
+            ):
+                return original_sync_send(self, request, *args, **kwargs)
+            runtime = getattr(self, "_mpp_payment_runtime", None) or _select_runtime("httpx")
+            if runtime is None:
+                return original_sync_send(self, request, *args, **kwargs)
+            token = _httpx_active.set(True)
+            try:
+                return runtime.send_httpx_sync(
+                    MethodType(original_sync_send, self),
+                    request,
+                    *args,
+                    **kwargs,
+                )
+            finally:
+                _httpx_active.reset(token)
+
+        _state.original_sync_send = original_sync_send
+        _state.sync_send_patch = sync_send
+        httpx.Client.send = sync_send  # type: ignore[method-assign]
+
+    if _state.original_async_send is None:
+        original_async_send = httpx.AsyncClient.send
+
+        async def async_send(
+            self: httpx.AsyncClient,
+            request: httpx.Request,
+            *args: Any,
+            **kwargs: Any,
+        ) -> httpx.Response:
+            if (
+                getattr(self, "_mpp_payment_wrapped", False)
+                or payment_flow_active()
+                or _httpx_active.get()
+            ):
+                return await original_async_send(self, request, *args, **kwargs)
+            runtime = getattr(self, "_mpp_payment_runtime", None) or _select_runtime("httpx")
+            if runtime is None:
+                return await original_async_send(self, request, *args, **kwargs)
+            token = _httpx_active.set(True)
+            try:
+                return await runtime.send_httpx(
+                    MethodType(original_async_send, self),
+                    request,
+                    *args,
+                    **kwargs,
+                )
+            finally:
+                _httpx_active.reset(token)
+
+        _state.original_async_send = original_async_send
+        _state.async_send_patch = async_send
+        httpx.AsyncClient.send = async_send  # type: ignore[method-assign]
+
+
+def _resolve_mcp_client(*, required: bool) -> Any | None:
+    try:
+        from mcp import ClientSession
+    except ImportError as error:
+        if required:
+            raise ImportError(
+                'Cannot instrument MCP calls. Install the "mcp" extra: pip install "pympp[mcp]"'
+            ) from error
+        return None
+    _ = ClientSession.call_tool
+    return ClientSession
+
+
+def _install_mcp_patch(client_session: Any) -> None:
+    if _state.original_mcp_call_tool is not None:
+        return
+    original_call_tool = client_session.call_tool
+
+    async def call_tool(
+        self: Any,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        if mcp_payment_flow_active() or _mcp_active.get():
+            return await original_call_tool(self, name, arguments, *args, **kwargs)
+        runtime = _select_runtime("mcp")
+        if runtime is None:
+            return await original_call_tool(self, name, arguments, *args, **kwargs)
+        token = _mcp_active.set(True)
+        try:
+            return await runtime.call_mcp_tool(
+                MethodType(original_call_tool, self),
+                name,
+                arguments,
+                *args,
+                **kwargs,
+            )
+        finally:
+            _mcp_active.reset(token)
+
+    client_session.call_tool = call_tool
+    _state.original_mcp_call_tool = original_call_tool
+    _state.mcp_call_tool_patch = call_tool
+    _state.mcp_client_session = client_session
+
+
+def _restore_unused_patches() -> None:
+    if not any(binding.active and binding.httpx for binding in _state.bindings):
+        if (
+            _state.sync_send_patch is not None
+            and httpx.Client.send is _state.sync_send_patch
+            and _state.original_sync_send is not None
+        ):
+            httpx.Client.send = _state.original_sync_send  # type: ignore[method-assign]
+        if (
+            _state.async_send_patch is not None
+            and httpx.AsyncClient.send is _state.async_send_patch
+            and _state.original_async_send is not None
+        ):
+            httpx.AsyncClient.send = _state.original_async_send  # type: ignore[method-assign]
+        _state.original_sync_send = None
+        _state.sync_send_patch = None
+        _state.original_async_send = None
+        _state.async_send_patch = None
+
+    if not any(binding.active and binding.mcp for binding in _state.bindings):
+        if (
+            _state.mcp_client_session is not None
+            and _state.mcp_call_tool_patch is not None
+            and _state.mcp_client_session.call_tool is _state.mcp_call_tool_patch
+            and _state.original_mcp_call_tool is not None
+        ):
+            _state.mcp_client_session.call_tool = _state.original_mcp_call_tool
+        _state.original_mcp_call_tool = None
+        _state.mcp_call_tool_patch = None
+        _state.mcp_client_session = None

src/mpp/runtime.py

@@ -27,13 +27,27 @@
 
 _T = TypeVar("_T")
 _PAYMENT_FLOW_ACTIVE: ContextVar[bool] = ContextVar("mpp_payment_flow_active", default=False)
+_MCP_FLOW_ACTIVE: ContextVar[bool] = ContextVar("mpp_mcp_flow_active", default=False)
+_payment_flow_count = 0
+_payment_flow_lock = threading.Lock()
 
 
 def payment_flow_active() -> bool:
     """Return whether the current context is creating a payment credential."""
     return _PAYMENT_FLOW_ACTIVE.get()
 
 
+def payment_flow_active_in_process() -> bool:
+    """Return whether any context is creating a payment credential."""
+    with _payment_flow_lock:
+        return _payment_flow_count > 0
+
+
+def mcp_payment_flow_active() -> bool:
+    """Return whether the current context is inside an MCP payment adapter."""
+    return _MCP_FLOW_ACTIVE.get()
+
+
 @runtime_checkable
 class Method(Protocol):
     """Payment method interface for client-side credential creation."""
@@ -299,6 +313,20 @@ async def call_mcp_tool(
         **kwargs: Any,
     ) -> Any:
         """Call an MCP tool with automatic payment handling, preserving result type."""
+        token = _MCP_FLOW_ACTIVE.set(True)
+        try:
+            return await self._call_mcp_tool(call_tool, name, arguments, *args, **kwargs)
+        finally:
+            _MCP_FLOW_ACTIVE.reset(token)
+
+    async def _call_mcp_tool(
+        self,
+        call_tool: Any,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
         from mpp.extensions.mcp.client import (
             PaymentOutcomeUnknownError,
             _extract_challenges,
@@ -471,7 +499,11 @@ async def _create_credential(
         *,
         event_payload: dict[str, Any] | None = None,
     ) -> Credential:
+        global _payment_flow_count
+
         token = _PAYMENT_FLOW_ACTIVE.set(True)
+        with _payment_flow_lock:
+            _payment_flow_count += 1
         try:
             payload = {
                 "challenge": challenge,
@@ -495,6 +527,8 @@ async def _create_credential(
             )
             return credential
         finally:
+            with _payment_flow_lock:
+                _payment_flow_count -= 1
             _PAYMENT_FLOW_ACTIVE.reset(token)
 
     async def emit_event(self, name: str, payload: EventPayload) -> Any: