diff --git a/.changelog/agent-runtime.md b/.changelog/agent-runtime.md new file mode 100644 index 0000000..f314a34 --- /dev/null +++ b/.changelog/agent-runtime.md @@ -0,0 +1,5 @@ +--- +pympp: minor +--- + +Added a shared payment runtime for payment-aware HTTP requests and MCP tool calls. diff --git a/src/mpp/client/transport.py b/src/mpp/client/transport.py index 3fa83db..db545ee 100644 --- a/src/mpp/client/transport.py +++ b/src/mpp/client/transport.py @@ -11,7 +11,7 @@ import logging from datetime import UTC, datetime -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any import httpx @@ -27,6 +27,7 @@ EventHandler, Unsubscribe, ) +from mpp.runtime import Method, PaymentRuntime logger = logging.getLogger(__name__) @@ -34,17 +35,6 @@ from collections.abc import Sequence -@runtime_checkable -class Method(Protocol): - """Payment method interface for client-side credential creation.""" - - name: str - - async def create_credential(self, challenge: Challenge) -> Credential: - """Create a credential to satisfy the given challenge.""" - ... - - def _client_payment_failed_payload( *, challenge: Challenge | None, @@ -90,10 +80,11 @@ def __init__( methods: Sequence[Method], inner: httpx.AsyncBaseTransport | None = None, events: EventDispatcher | None = None, + runtime: PaymentRuntime | None = None, ) -> None: - self._methods = {m.name: m for m in methods} + self._runtime = runtime or PaymentRuntime(methods, events=events) self._inner = inner or httpx.AsyncHTTPTransport() - self._events = events or EventDispatcher() + self._events = self._runtime.events def on(self, name: str, handler: EventHandler) -> Unsubscribe: """Register a client payment event handler.""" @@ -139,13 +130,15 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: continue challenges.append(parsed) - challenge = None - matched_method = None - for parsed in challenges: - if parsed.method in self._methods: - challenge = parsed - matched_method = self._methods[parsed.method] - break + try: + challenge, matched_method = self._runtime.match_challenge( + challenges, + prefer_method_order=False, + require_intent_match=False, + ) + except ValueError: + challenge = None + matched_method = None if not challenge or not matched_method: if parse_error is not None or challenges: @@ -205,7 +198,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: credential = ( event_credential if isinstance(event_credential, Credential) - else await matched_method.create_credential(challenge) + else await self._runtime.create_credential(challenge, matched_method) ) await self._events.emit( CREDENTIAL_CREATED, diff --git a/src/mpp/extensions/mcp/client.py b/src/mpp/extensions/mcp/client.py index 7310f3c..8aa7895 100644 --- a/src/mpp/extensions/mcp/client.py +++ b/src/mpp/extensions/mcp/client.py @@ -27,13 +27,11 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from typing import Any from mpp.extensions.mcp.constants import CODE_PAYMENT_REQUIRED, META_RECEIPT -from mpp.extensions.mcp.types import MCPChallenge, MCPCredential, MCPReceipt - -if TYPE_CHECKING: - from mpp import Challenge, Credential +from mpp.extensions.mcp.types import MCPChallenge, MCPReceipt +from mpp.runtime import Method, PaymentRuntime logger = logging.getLogger(__name__) @@ -51,17 +49,6 @@ def __init__(self, challenge: MCPChallenge, cause: Exception) -> None: ) -@runtime_checkable -class Method(Protocol): - """Payment method interface for MCP client credential creation.""" - - name: str - - async def create_credential(self, challenge: Challenge) -> Credential: - """Create a credential to satisfy the given challenge.""" - ... - - def _error_code(error: Exception) -> int | None: code = getattr(error, "code", None) if code is not None: @@ -173,7 +160,7 @@ class McpClient: def __init__(self, session: Any, methods: list[Method]) -> None: self._session = session - self._methods = methods + self._runtime = PaymentRuntime(methods) def __getattr__(self, name: str) -> Any: return getattr(self._session, name) @@ -206,46 +193,20 @@ async def call_tool( PaymentOutcomeUnknownError: If the paid retry fails after sending a credential. ValueError: If no installed method matches the server's challenge. """ - from mcp.shared.exceptions import McpError - call_kwargs: dict[str, Any] = {} if timeout is not None: call_kwargs["read_timeout_seconds"] = timeout if meta is not None: call_kwargs["meta"] = meta - try: - result = await self._session.call_tool(name, arguments, **call_kwargs) - receipt = self._extract_receipt(result) - return McpToolResult(result=result, receipt=receipt) - - except McpError as e: - if not _is_payment_required_error(e): - raise - - challenges = _extract_challenges(e) - if not challenges: - raise ValueError("Server returned malformed payment challenges") from e - - challenge, method = self._match_challenge(challenges) - - core_credential = await method.create_credential(challenge.to_core()) - mcp_credential = MCPCredential.from_core(core_credential, challenge) - - retry_meta = dict(meta) if meta else {} - retry_meta.update(mcp_credential.to_meta()) - - retry_kwargs: dict[str, Any] = {"meta": retry_meta} - if timeout is not None: - retry_kwargs["read_timeout_seconds"] = timeout - - try: - retry_result = await self._session.call_tool(name, arguments, **retry_kwargs) - except Exception as exc: - raise PaymentOutcomeUnknownError(challenge, exc) from exc - - receipt = self._extract_receipt(retry_result) - return McpToolResult(result=retry_result, receipt=receipt) + result = await self._runtime.call_mcp_tool( + self._session.call_tool, + name, + arguments, + **call_kwargs, + ) + receipt = self._extract_receipt(result) + return McpToolResult(result=result, receipt=receipt) def _match_challenge(self, challenges: list[MCPChallenge]) -> tuple[MCPChallenge, Method]: """Match a challenge to an installed method. @@ -253,25 +214,7 @@ def _match_challenge(self, challenges: list[MCPChallenge]) -> tuple[MCPChallenge Iterates installed methods in order (client preference) and returns the first match by ``name`` and ``intent``. """ - for method in self._methods: - supported_intents = self._intent_names(method) - for challenge in challenges: - if challenge.method == method.name and challenge.intent in supported_intents: - return challenge, method - - available = [challenge.method for challenge in challenges] - installed = [m.name for m in self._methods] - raise ValueError( - f"No compatible payment method. Server offered: {available}, client has: {installed}" - ) - - @staticmethod - def _intent_names(method: Method) -> set[str]: - """Get intent names supported by a method.""" - intents = getattr(method, "intents", None) or getattr(method, "_intents", None) - if isinstance(intents, dict): - return set(intents.keys()) - return {"charge"} + return self._runtime.match_mcp_challenge(challenges) @staticmethod def _extract_receipt(result: Any) -> MCPReceipt | None: diff --git a/src/mpp/runtime.py b/src/mpp/runtime.py new file mode 100644 index 0000000..b4b95b1 --- /dev/null +++ b/src/mpp/runtime.py @@ -0,0 +1,235 @@ +"""Shared payment runtime for HTTP and MCP clients.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from urllib.parse import urlparse + +import httpx + +from mpp import Challenge, Credential +from mpp.events import EventDispatcher + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from mpp.client import PaymentTransport + + +@runtime_checkable +class Method(Protocol): + """Payment method interface for client-side credential creation.""" + + name: str + + async def create_credential(self, challenge: Challenge) -> Credential: + """Create a credential to satisfy the given challenge.""" + ... + + +class _BoundSendTransport(httpx.AsyncBaseTransport): + def __init__(self, send: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None: + self._send = send + self._args = args + self._kwargs = kwargs + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + return await self._send(request, *self._args, **self._kwargs) + + async def aclose(self) -> None: + return None + + +class PaymentRuntime: + """Reusable payment runtime for HTTP and MCP payment handling.""" + + def __init__( + self, + methods: Sequence[Method], + *, + events: EventDispatcher | None = None, + allowed_origins: Sequence[str] | None = None, + ) -> None: + self.methods = list(methods) + self._methods_by_name = {method.name: method for method in self.methods} + self.events = events or EventDispatcher() + self._allowed = _AllowedOrigins(allowed_origins) + + def payment_transport(self, inner: httpx.AsyncBaseTransport | None = None) -> PaymentTransport: + """Create an httpx transport using this runtime's payment methods.""" + from mpp.client import PaymentTransport + + return PaymentTransport( + methods=self.methods, + inner=inner, + events=self.events, + runtime=self, + ) + + def wrap_async_client(self, client: httpx.AsyncClient) -> httpx.AsyncClient: + """Make one existing AsyncClient payment-aware without global instrumentation.""" + client._mpp_payment_runtime = self # type: ignore[attr-defined] + if getattr(client, "_mpp_payment_wrapped", False): + return client + + original_send = client.send + + async def send(request: httpx.Request, *args: Any, **kwargs: Any) -> httpx.Response: + runtime = getattr(client, "_mpp_payment_runtime", self) + return await runtime.send_httpx(original_send, request, *args, **kwargs) + + client._mpp_payment_original_send = original_send # type: ignore[attr-defined] + client._mpp_payment_wrapped = True # type: ignore[attr-defined] + client.send = send # type: ignore[method-assign] + return client + + def httpx_client_factory( + self, + base_factory: Any = httpx.AsyncClient, + ) -> Callable[..., httpx.AsyncClient]: + """Return an AsyncClient factory suitable for MCP SDK framework hooks.""" + + def factory(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + return self.wrap_async_client(base_factory(*args, **kwargs)) + + return factory + + async def send_httpx( + self, + send: Any, + request: httpx.Request, + *args: Any, + **kwargs: Any, + ) -> httpx.Response: + """Send one httpx request with automatic 402 payment handling.""" + if not self._allowed.http_url(request.url): + return await send(request, *args, **kwargs) + + transport = _BoundSendTransport(send, args, dict(kwargs)) + return await self.payment_transport(inner=transport).handle_async_request(request) + + async def call_mcp_tool( + self, + call_tool: Any, + name: str, + arguments: dict[str, Any] | None = None, + *args: Any, + **kwargs: Any, + ) -> Any: + """Call an MCP tool with automatic payment handling, preserving result type.""" + from mcp.shared.exceptions import McpError + + from mpp.extensions.mcp.client import ( + PaymentOutcomeUnknownError, + _extract_challenges, + _is_payment_required_error, + ) + from mpp.extensions.mcp.types import MCPCredential + + try: + return await call_tool(name, arguments, *args, **kwargs) + except McpError as error: + if not _is_payment_required_error(error): + raise + + challenges = [ + challenge + for challenge in _extract_challenges(error) + if self._allowed.mcp_realm(challenge.realm) + ] + if not challenges: + raise ValueError( + "Server returned malformed payment challenges or disallowed payment origins" + ) from error + + challenge, method = self.match_challenge(challenges) + core_credential = await self.create_credential(challenge.to_core(), method) + mcp_credential = MCPCredential.from_core(core_credential, challenge) + + retry_kwargs = dict(kwargs) + retry_meta = dict(retry_kwargs.get("meta") or {}) + retry_meta.update(mcp_credential.to_meta()) + retry_kwargs["meta"] = retry_meta + + try: + return await call_tool(name, arguments, *args, **retry_kwargs) + except Exception as exc: + raise PaymentOutcomeUnknownError(challenge, exc) from exc + + def match_challenge( + self, + challenges: list[Any], + *, + prefer_method_order: bool = True, + require_intent_match: bool = True, + ) -> tuple[Any, Method]: + """Match payment challenges against configured methods.""" + if prefer_method_order: + for method in self.methods: + for challenge in challenges: + if challenge.method != method.name: + continue + if require_intent_match and challenge.intent not in _intent_names(method): + continue + return challenge, method + else: + for challenge in challenges: + method = self._methods_by_name.get(challenge.method) + if method is None: + continue + if require_intent_match and challenge.intent not in _intent_names(method): + continue + return challenge, method + + available = [challenge.method for challenge in challenges] + installed = list(self._methods_by_name) + raise ValueError( + f"No compatible payment method. Server offered: {available}, client has: {installed}" + ) + + def match_mcp_challenge(self, challenges: list[Any]) -> tuple[Any, Method]: + """Match MCP payment challenges against configured methods.""" + return self.match_challenge(challenges) + + async def create_credential(self, challenge: Challenge, method: Method) -> Credential: + """Create a credential for a matched challenge and method.""" + return await method.create_credential(challenge) + + +def _intent_names(method: Method) -> set[str]: + intents = getattr(method, "intents", None) or getattr(method, "_intents", None) + if isinstance(intents, dict): + return set(intents.keys()) + return {"charge"} + + +class _AllowedOrigins: + def __init__(self, allowed_origins: Sequence[str] | None) -> None: + self._allow_all = allowed_origins is None + self._values = set[str]() + if allowed_origins is None: + return + for value in allowed_origins: + parsed = urlparse(str(value)) + if parsed.scheme and parsed.netloc: + self._values.add(f"{parsed.scheme}://{parsed.netloc}") + self._values.add(parsed.netloc) + else: + self._values.add(str(value)) + + def http_url(self, url: httpx.URL) -> bool: + if self._allow_all: + return True + origin = f"{url.scheme}://{url.host}" + if url.port is not None: + origin = f"{origin}:{url.port}" + return origin in self._values or url.host in self._values + + def mcp_realm(self, realm: str) -> bool: + if self._allow_all: + return True + parsed = urlparse(realm) + if parsed.scheme and parsed.netloc: + origin = f"{parsed.scheme}://{parsed.netloc}" + return origin in self._values or parsed.netloc in self._values + return realm in self._values diff --git a/tests/test_agent.py b/tests/test_agent.py new file mode 100644 index 0000000..3876f6f --- /dev/null +++ b/tests/test_agent.py @@ -0,0 +1,220 @@ +"""Tests for opt-in agent runtime instrumentation.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from types import ModuleType +from typing import Any +from unittest.mock import AsyncMock + +import httpx +import pytest + +from mpp import Challenge, ChallengeEcho, Credential +from mpp.extensions.mcp import META_CREDENTIAL +from mpp.runtime import PaymentRuntime + + +class MockMethod: + name = "tempo" + _intents = {"charge": True} + + def __init__(self) -> None: + self.create_credential = AsyncMock( + return_value=Credential( + challenge=ChallengeEcho( + id="test-id", + realm="example.com", + method="tempo", + intent="charge", + request="e30", + ), + payload={"hash": "0xabc"}, + source="0x1234", + ) + ) + + +class MockTransport(httpx.AsyncBaseTransport): + def __init__(self, responses: list[httpx.Response]) -> None: + self.responses = responses + self.requests: list[httpx.Request] = [] + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + return self.responses.pop(0) + + +class FakeCallToolResult: + def __init__(self, text: str = "ok") -> None: + self.content = [{"type": "text", "text": text}] + self.isError = False + + +class FakeMcpError(Exception): + def __init__(self, code: int, data: Any = None) -> None: + super().__init__("mcp error") + self.code = code + self.data = data + + +class FakeClientSession: + def __init__(self, side_effects: list[Any]) -> None: + self.side_effects = side_effects + self.calls: list[tuple[str, dict[str, Any] | None, tuple[Any, ...], dict[str, Any]]] = [] + + async def call_tool( + self, + name: str, + arguments: dict[str, Any] | None = None, + *args: Any, + **kwargs: Any, + ) -> Any: + self.calls.append((name, arguments, args, kwargs)) + item = self.side_effects.pop(0) + if isinstance(item, Exception): + raise item + return item + + +@dataclass +class FakeMcpModules: + client_session: type[FakeClientSession] + + +def install_fake_mcp(monkeypatch: pytest.MonkeyPatch) -> FakeMcpModules: + mcp = ModuleType("mcp") + shared = ModuleType("mcp.shared") + exceptions = ModuleType("mcp.shared.exceptions") + exceptions.McpError = FakeMcpError # type: ignore[attr-defined] + shared.exceptions = exceptions # type: ignore[attr-defined] + mcp.ClientSession = FakeClientSession # type: ignore[attr-defined] + mcp.shared = shared # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "mcp", mcp) + monkeypatch.setitem(sys.modules, "mcp.shared", shared) + monkeypatch.setitem(sys.modules, "mcp.shared.exceptions", exceptions) + return FakeMcpModules(client_session=FakeClientSession) + + +def payment_challenge_header() -> str: + challenge = Challenge(id="test-id", method="tempo", intent="charge", request={}) + return challenge.to_www_authenticate("example.com") + + +def mcp_payment_error() -> FakeMcpError: + return FakeMcpError( + -32042, + data={ + "challenges": [ + { + "id": "test-id", + "realm": "example.com", + "method": "tempo", + "intent": "charge", + "request": {"amount": "1000"}, + } + ] + }, + ) + + +class TestHttpxInstrumentation: + @pytest.mark.asyncio + async def test_runtime_wrap_async_client_without_global_hook(self) -> None: + inner = MockTransport( + [ + httpx.Response(402, headers={"www-authenticate": payment_challenge_header()}), + httpx.Response(200, json={"ok": True}), + ] + ) + runtime = PaymentRuntime([MockMethod()]) + client = runtime.wrap_async_client(httpx.AsyncClient(transport=inner)) + try: + response = await client.get("https://example.com/paid") + finally: + await client.aclose() + + assert response.status_code == 200 + assert len(inner.requests) == 2 + assert inner.requests[1].headers["authorization"].startswith("Payment ") + + @pytest.mark.asyncio + async def test_disallowed_http_origin_does_not_retry_payment(self) -> None: + inner = MockTransport( + [ + httpx.Response(402, headers={"www-authenticate": payment_challenge_header()}), + ] + ) + method = MockMethod() + runtime = PaymentRuntime( + [method], + allowed_origins=["https://other.example"], + ) + client = runtime.wrap_async_client(httpx.AsyncClient(transport=inner)) + try: + response = await client.get("https://example.com/paid") + finally: + await client.aclose() + + assert response.status_code == 402 + assert len(inner.requests) == 1 + method.create_credential.assert_not_called() + + +class TestMcpInstrumentation: + @pytest.mark.asyncio + async def test_payment_runtime_call_mcp_tool_pays_and_preserves_raw_result( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = install_fake_mcp(monkeypatch) + raw_result = FakeCallToolResult("paid") + session = fake.client_session([mcp_payment_error(), raw_result]) + runtime = PaymentRuntime([MockMethod()]) + + result = await runtime.call_mcp_tool( + session.call_tool, + "premium_tool", + {"query": "test"}, + progress_callback="callback", + meta={"trace_id": "abc"}, + ) + + 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_disallowed_mcp_realm_does_not_retry_payment( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + fake = install_fake_mcp(monkeypatch) + method = MockMethod() + session = fake.client_session([mcp_payment_error()]) + runtime = PaymentRuntime( + [method], + allowed_origins=["other.example"], + ) + + with pytest.raises(ValueError, match="disallowed"): + await runtime.call_mcp_tool(session.call_tool, "premium_tool", {"query": "test"}) + + assert len(session.calls) == 1 + method.create_credential.assert_not_called() + + @pytest.mark.asyncio + async def test_payment_runtime_call_mcp_tool_preserves_raw_result( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + install_fake_mcp(monkeypatch) + raw_result = FakeCallToolResult("paid") + session = FakeClientSession([mcp_payment_error(), raw_result]) + runtime = PaymentRuntime([MockMethod()]) + + result = await runtime.call_mcp_tool(session.call_tool, "premium_tool", {}) + + assert result is raw_result + assert len(session.calls) == 2 diff --git a/tests/test_client.py b/tests/test_client.py index 2164399..8efa5d1 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -8,6 +8,7 @@ from mpp import Challenge, Credential from mpp.client import Client, PaymentTransport, get, post, request +from mpp.runtime import PaymentRuntime from tests import make_credential @@ -89,6 +90,35 @@ async def test_handles_402_with_matching_method(self) -> None: method.create_credential.assert_called_once() + @pytest.mark.asyncio + async def test_delegates_payment_matching_to_runtime(self) -> None: + """Should use an injected runtime for payment matching and credentials.""" + challenge = Challenge( + id="test-id", + method="tempo", + intent="charge", + request={"amount": "1000"}, + ) + inner = MockTransport( + [ + httpx.Response( + 402, + headers={"www-authenticate": challenge.to_www_authenticate("example.com")}, + ), + httpx.Response(200, content=b'{"data": "ok"}'), + ] + ) + + method = MockMethod() + runtime = PaymentRuntime([method]) + transport = PaymentTransport(methods=[], inner=inner, runtime=runtime) + + response = await transport.handle_async_request(httpx.Request("GET", "https://example.com")) + + assert response.status_code == 200 + assert inner.requests[1].headers["Authorization"].startswith("Payment ") + method.create_credential.assert_called_once() + @pytest.mark.asyncio async def test_emits_client_payment_events(self) -> None: """Should emit mppx-compatible client payment lifecycle events."""