|
| 1 | +"""Payment-aware MCP client wrapper. |
| 2 | +
|
| 3 | +Wraps an MCP SDK ``ClientSession`` with automatic payment handling. |
| 4 | +When a tool call returns a ``-32042`` payment required error, the wrapper |
| 5 | +creates a Credential and retries the call—mirroring the TypeScript |
| 6 | +``McpClient.wrap`` API. |
| 7 | +
|
| 8 | +Example: |
| 9 | + from mcp import ClientSession |
| 10 | + from mcp.client.sse import sse_client |
| 11 | + from mpp.extensions.mcp import McpClient |
| 12 | + from mpp.methods.tempo import tempo, TempoAccount, ChargeIntent |
| 13 | +
|
| 14 | + account = TempoAccount.from_key("0x...") |
| 15 | + method = tempo(account=account, intents={"charge": ChargeIntent()}) |
| 16 | +
|
| 17 | + async with sse_client("http://localhost:8000/sse") as streams: |
| 18 | + async with ClientSession(streams[0], streams[1]) as session: |
| 19 | + await session.initialize() |
| 20 | +
|
| 21 | + client = McpClient(session, methods=[method]) |
| 22 | + result = await client.call_tool("premium_tool", {"query": "hello"}) |
| 23 | + print(result.receipt) |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import logging |
| 29 | +from dataclasses import dataclass |
| 30 | +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable |
| 31 | + |
| 32 | +from mpp.extensions.mcp.constants import CODE_PAYMENT_REQUIRED, META_RECEIPT |
| 33 | +from mpp.extensions.mcp.types import MCPChallenge, MCPCredential, MCPReceipt |
| 34 | + |
| 35 | +if TYPE_CHECKING: |
| 36 | + from mpp import Challenge, Credential |
| 37 | + |
| 38 | +logger = logging.getLogger(__name__) |
| 39 | + |
| 40 | + |
| 41 | +@runtime_checkable |
| 42 | +class Method(Protocol): |
| 43 | + """Payment method interface for MCP client credential creation.""" |
| 44 | + |
| 45 | + name: str |
| 46 | + |
| 47 | + async def create_credential(self, challenge: Challenge) -> Credential: |
| 48 | + """Create a credential to satisfy the given challenge.""" |
| 49 | + ... |
| 50 | + |
| 51 | + |
| 52 | +def _is_payment_required_error(error: Exception) -> bool: |
| 53 | + """Check whether an MCP error is a -32042 payment required error. |
| 54 | +
|
| 55 | + Distinguishes payment errors from other uses of -32042 (such as |
| 56 | + URL elicitation) by checking for a ``challenges`` array in ``error.data``. |
| 57 | + """ |
| 58 | + code = getattr(error, "code", None) |
| 59 | + if code != CODE_PAYMENT_REQUIRED: |
| 60 | + return False |
| 61 | + data = getattr(error, "data", None) |
| 62 | + if not isinstance(data, dict): |
| 63 | + return False |
| 64 | + challenges = data.get("challenges") |
| 65 | + return isinstance(challenges, list) and len(challenges) > 0 |
| 66 | + |
| 67 | + |
| 68 | +def _extract_challenges(error: Exception) -> list[dict[str, Any]]: |
| 69 | + """Extract the challenges array from a payment required error.""" |
| 70 | + data = getattr(error, "data", {}) |
| 71 | + return data.get("challenges", []) if isinstance(data, dict) else [] |
| 72 | + |
| 73 | + |
| 74 | +@dataclass(frozen=True, slots=True) |
| 75 | +class McpToolResult: |
| 76 | + """Result of a payment-aware tool call. |
| 77 | +
|
| 78 | + Wraps the raw MCP ``CallToolResult`` and surfaces the payment receipt. |
| 79 | + """ |
| 80 | + |
| 81 | + result: Any |
| 82 | + receipt: MCPReceipt | None = None |
| 83 | + |
| 84 | + |
| 85 | +class McpClient: |
| 86 | + """Payment-aware MCP client wrapper. |
| 87 | +
|
| 88 | + Wraps an MCP SDK ``ClientSession`` and overrides ``call_tool`` with |
| 89 | + automatic payment handling. When a tool call returns ``-32042``, the |
| 90 | + wrapper matches the challenge to an installed payment method, creates |
| 91 | + a credential, and retries. |
| 92 | +
|
| 93 | + Args: |
| 94 | + session: An initialized ``mcp.ClientSession``. |
| 95 | + methods: Payment methods available for credential creation. |
| 96 | +
|
| 97 | + Example: |
| 98 | + client = McpClient(session, methods=[tempo(...)]) |
| 99 | + result = await client.call_tool("premium_tool", {"query": "hello"}) |
| 100 | + print(result.receipt) |
| 101 | + """ |
| 102 | + |
| 103 | + def __init__(self, session: Any, methods: list[Method]) -> None: |
| 104 | + self._session = session |
| 105 | + self._methods = methods |
| 106 | + |
| 107 | + async def call_tool( |
| 108 | + self, |
| 109 | + name: str, |
| 110 | + arguments: dict[str, Any] | None = None, |
| 111 | + *, |
| 112 | + timeout: float | None = None, |
| 113 | + meta: dict[str, Any] | None = None, |
| 114 | + ) -> McpToolResult: |
| 115 | + """Call an MCP tool with automatic payment handling. |
| 116 | +
|
| 117 | + On a ``-32042`` error, matches the challenge to an installed method, |
| 118 | + creates a credential, and retries the call with the credential in |
| 119 | + ``params._meta``. |
| 120 | +
|
| 121 | + Args: |
| 122 | + name: Tool name. |
| 123 | + arguments: Tool arguments. |
| 124 | + timeout: Per-call timeout override (passed as ``read_timeout_seconds``). |
| 125 | + meta: Additional ``_meta`` fields to include in the request. |
| 126 | +
|
| 127 | + Returns: |
| 128 | + An ``McpToolResult`` with the tool result and an optional receipt. |
| 129 | +
|
| 130 | + Raises: |
| 131 | + McpError: If the error is not payment-related or no method matches. |
| 132 | + ValueError: If no installed method matches the server's challenge. |
| 133 | + """ |
| 134 | + from mcp.shared.exceptions import McpError |
| 135 | + |
| 136 | + call_kwargs: dict[str, Any] = {} |
| 137 | + if timeout is not None: |
| 138 | + call_kwargs["read_timeout_seconds"] = timeout |
| 139 | + if meta is not None: |
| 140 | + call_kwargs["meta"] = meta |
| 141 | + |
| 142 | + try: |
| 143 | + result = await self._session.call_tool(name, arguments, **call_kwargs) |
| 144 | + receipt = self._extract_receipt(result) |
| 145 | + return McpToolResult(result=result, receipt=receipt) |
| 146 | + |
| 147 | + except McpError as e: |
| 148 | + if not _is_payment_required_error(e): |
| 149 | + raise |
| 150 | + |
| 151 | + challenges_data = _extract_challenges(e) |
| 152 | + challenge, method = self._match_challenge(challenges_data) |
| 153 | + |
| 154 | + core_credential = await method.create_credential(challenge.to_core()) |
| 155 | + mcp_credential = MCPCredential.from_core(core_credential, challenge) |
| 156 | + |
| 157 | + retry_meta = dict(meta) if meta else {} |
| 158 | + retry_meta.update(mcp_credential.to_meta()) |
| 159 | + |
| 160 | + retry_kwargs: dict[str, Any] = {"meta": retry_meta} |
| 161 | + if timeout is not None: |
| 162 | + retry_kwargs["read_timeout_seconds"] = timeout |
| 163 | + |
| 164 | + retry_result = await self._session.call_tool(name, arguments, **retry_kwargs) |
| 165 | + receipt = self._extract_receipt(retry_result) |
| 166 | + return McpToolResult(result=retry_result, receipt=receipt) |
| 167 | + |
| 168 | + def _match_challenge( |
| 169 | + self, challenges_data: list[dict[str, Any]] |
| 170 | + ) -> tuple[MCPChallenge, Method]: |
| 171 | + """Match a challenge to an installed method. |
| 172 | +
|
| 173 | + Iterates installed methods in order (client preference) and returns |
| 174 | + the first match by ``name`` and ``intent``. |
| 175 | + """ |
| 176 | + for method in self._methods: |
| 177 | + for cd in challenges_data: |
| 178 | + if cd.get("method") == method.name and cd.get("intent") in self._intent_names( |
| 179 | + method |
| 180 | + ): |
| 181 | + return MCPChallenge.from_dict(cd), method |
| 182 | + |
| 183 | + available = [cd.get("method") for cd in challenges_data] |
| 184 | + installed = [m.name for m in self._methods] |
| 185 | + raise ValueError( |
| 186 | + f"No compatible payment method. " |
| 187 | + f"Server offered: {available}, client has: {installed}" |
| 188 | + ) |
| 189 | + |
| 190 | + @staticmethod |
| 191 | + def _intent_names(method: Method) -> set[str]: |
| 192 | + """Get intent names supported by a method.""" |
| 193 | + intents = getattr(method, "intents", None) or getattr(method, "_intents", None) |
| 194 | + if isinstance(intents, dict): |
| 195 | + return set(intents.keys()) |
| 196 | + return {"charge"} |
| 197 | + |
| 198 | + @staticmethod |
| 199 | + def _extract_receipt(result: Any) -> MCPReceipt | None: |
| 200 | + """Extract a payment receipt from a tool result's _meta.""" |
| 201 | + meta = getattr(result, "meta", None) |
| 202 | + if not meta or not isinstance(meta, dict): |
| 203 | + return None |
| 204 | + receipt_data = meta.get(META_RECEIPT) |
| 205 | + if receipt_data is None: |
| 206 | + return None |
| 207 | + try: |
| 208 | + return MCPReceipt.from_dict(receipt_data) |
| 209 | + except (KeyError, TypeError): |
| 210 | + logger.warning("Failed to parse receipt from _meta") |
| 211 | + return None |
0 commit comments