diff --git a/.changelog/agent-instrument.md b/.changelog/agent-instrument.md new file mode 100644 index 0000000..828cc74 --- /dev/null +++ b/.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. diff --git a/src/mpp/agent.py b/src/mpp/agent.py new file mode 100644 index 0000000..1b5993d --- /dev/null +++ b/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 diff --git a/tests/test_agent.py b/tests/test_agent.py index 3876f6f..1932cda 100644 --- a/tests/test_agent.py +++ b/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)