from collections.abc import Sequence

from typing import Any, Literal

import mcp.types as types

from mcp import Client

from mcp.client import ClaimContext, ClientExtension, ResultClaim

from mcp.server.context import CallNext, HandlerResult, ServerRequestContext

from mcp.server.extension import Extension

from mcp.server.mcpserver import MCPServer, require_client_extension

EXTENSION_ID = "com.example/receipts"

class ReceiptResult(types.Result):

"""The claimed result shape; `result_type` pins the wire tag."""

result_type: Literal["receipt"] = "receipt"

receipt_token: str

class ReceiptIssuer(Extension):

"""Server half: answers `buy` with a receipt instead of a final result."""

identifier = EXTENSION_ID

async def intercept_tool_call(

self,

params: types.CallToolRequestParams,

ctx: ServerRequestContext[Any, Any],

call_next: CallNext,

) -> HandlerResult:

if params.name != "buy":

return await call_next(ctx)

require_client_extension(ctx, EXTENSION_ID)

return {"resultType": "receipt", "receiptToken": "r-117"}

class Receipts(ClientExtension):

"""Client half: claims the `receipt` shape and supplies the code that finishes it."""

identifier = EXTENSION_ID

def claims(self) -> Sequence[ResultClaim[Any]]:

return [ResultClaim(result_type="receipt", model=ReceiptResult, resolve=self._redeem)]

async def _redeem(self, claimed: ReceiptResult, ctx: ClaimContext) -> types.CallToolResult:

return await ctx.session.call_tool("redeem", {"token": claimed.receipt_token})

mcp = MCPServer("shop", extensions=[ReceiptIssuer()])

@mcp.tool()

def buy(item: str) -> types.CallToolResult:

"""Buy an item."""

raise NotImplementedError # ReceiptIssuer answers `buy` before the tool runs

@mcp.tool()

def redeem(token: str) -> str:

"""Exchange a receipt token for the goods."""

return f"goods for {token}"

async def main() -> None:

async with Client(mcp, extensions=[Receipts()]) as client:

result = await client.call_tool("buy", {"item": "lamp"})

print(result.content)

# [TextContent(text='goods for r-117')]