MCP Integration

1
import asyncio
2
import logging
3
import sys
4
5
# Set up basic logging
6
logging.basicConfig(
7
    level=logging.INFO, 
8
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
9
)
10
logger = logging.getLogger("mcp-client")
11
12
async def connect_to_mcp():
13
    """Connect to an MCP server and use its tools"""
14
    from mcp.client.streamable_http import streamablehttp_client
15
    from mcp import ClientSession
16
17
    # Headers for authentication and other requirements
18
    headers = {
19
        "Authorization": "Bearer <auth_token>",
20
        "Mcp-Session-Id": "session-" + str(asyncio.current_task().get_name()), 
21
        "X-Chat-Id": "chat-" + str(asyncio.current_task().get_name())
22
    }
23
    
24
    # MCP server URL
25
    server_url = "https://ah-api.merge.dev/api/v1/tool-packs/<tool_pack_id>/registered-users/<registered_user_id>/mcp"
26
    
27
    logger.info(f"Connecting to MCP server: {server_url}")
28
    
29
    try:
30
        # Connect to the MCP server with streamablehttp_client
31
        async with streamablehttp_client(server_url, headers=headers) as (read_stream, write_stream, _):
32
            logger.info("Connection established, creating ClientSession")
33
            # Create a session using the client streams
34
            async with ClientSession(read_stream, write_stream) as session:
35
                # Initialize the connection with the correct protocol version
36
                logger.info("Initializing session with protocol version 2024-11-05")
37
                await session.initialize()
38
                
39
                # List available tools
40
                logger.info("Listing tools")
41
                tools_result = await session.list_tools()
42
                logger.info("\nMCP Tools:")
43
                logger.info(tools_result)
44
45
    except Exception as e:
46
        logger.exception(f"Error in MCP client: {e}")
47
48
async def main():
49
    logger.info("Starting MCP client")
50
    await connect_to_mcp()
51
52
if __name__ == "__main__":
53
    try:
54
        asyncio.run(main())
55
    except KeyboardInterrupt:
56
        logger.info("Program terminated by user")
57
        sys.exit(0)
58
    except Exception as e:
59
        logger.exception(f"Unhandled exception: {e}")
60
        sys.exit(1)

1
import json
2
import uuid
3
import asyncio
4
import aiohttp
5
6
class MCPClient:
7
    """
8
    Model Context Protocol (MCP) client for interacting with Merge Agent Handler MCP Server
9
    
10
    This client implements the JSON-RPC 2.0 protocol used by the MCP server.
11
    It supports initializing a session, listing available tools, and executing tools.
12
    Uses async methods to match the server's async implementation.
13
    """
14
    
15
    def __init__(self, base_url, auth_token=None):
16
        """
17
        Initialize MCP client
18
        
19
        Args:
20
            base_url (str): The base URL for the MCP server
21
            auth_token (str, optional): Bearer token for authentication
22
        """
23
        self.base_url = base_url
24
        self.session_id = str(uuid.uuid4())
25
        self.request_id = 1
26
        self.headers = {
27
            "Content-Type": "application/json",
28
            "Accept": "application/json, text/event-stream",
29
            "Mcp-Session-Id": self.session_id
30
        }
31
        
32
        if auth_token:
33
            self.headers["Authorization"] = f"Bearer {auth_token}"
34
        
35
        self.session = None
36
    
37
    async def _ensure_session(self):
38
        """Ensure that an aiohttp ClientSession exists"""
39
        if self.session is None or self.session.closed:
40
            self.session = aiohttp.ClientSession()
41
        return self.session
42
            
43
    async def _send_request(self, method, params=None):
44
        """Send an async JSON-RPC request to the MCP server"""
45
        payload = {
46
            "jsonrpc": "2.0",
47
            "id": self.request_id,
48
            "method": method
49
        }
50
        
51
        if params:
52
            payload["params"] = params
53
        
54
        session = await self._ensure_session()
55
        
56
        try:
57
            async with session.post(
58
                self.base_url,
59
                headers=self.headers,
60
                json=payload
61
            ) as response:
62
                if "Mcp-Session-Id" in response.headers:
63
                    self.session_id = response.headers["Mcp-Session-Id"]
64
                    self.headers["Mcp-Session-Id"] = self.session_id
65
                
66
                self.request_id += 1
67
                response.raise_for_status()
68
                
69
                try:
70
                    return await response.json()
71
                except aiohttp.ContentTypeError:
72
                    return {"text": await response.text()}
73
                    
74
        except aiohttp.ClientError as e:
75
            return {"error": f"Request failed: {str(e)}"}
76
    
77
    async def initialize(self, protocol_version="2024-11-05"):
78
        """Initialize a session with the MCP server"""
79
        params = {"protocolVersion": protocol_version}
80
        return await self._send_request("initialize", params)
81
    
82
    async def list_tools(self):
83
        """List available tools from the MCP server"""
84
        return await self._send_request("tools/list")
85
    
86
    async def call_tool(self, tool_name, arguments=None):
87
        """Execute a tool on the MCP server"""
88
        params = {
89
            "name": tool_name,
90
            "arguments": arguments or {}
91
        }
92
        return await self._send_request("tools/call", params)
93
    
94
    async def close(self):
95
        """Close the aiohttp session"""
96
        if self.session and not self.session.closed:
97
            await self.session.close()
98
99
# Example usage
100
async def main():
101
    client = MCPClient(
102
        base_url="https://ah-api.merge.dev/api/v1/tool-packs/<tool_pack_id>/registered-users/<user_id>/mcp",
103
        auth_token="<ah-production-access-key>"
104
    )
105
    
106
    try:
107
        # Initialize the session
108
        init_result = await client.initialize()
109
        print("\nMCP Initialization Response:")
110
        print(json.dumps(init_result, indent=2))
111
        
112
        # List available tools
113
        tools_result = await client.list_tools()
114
        print("\nMCP Tools List Response:")
115
        print(json.dumps(tools_result, indent=2))
116
        
117
    finally:
118
        await client.close()
119
120
if __name__ == "__main__":
121
    asyncio.run(main())