|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
import httpx
|
|
import websockets
|
|
from websockets.client import WebSocketClientProtocol
|
|
|
|
|
|
@dataclass
|
|
class Position:
|
|
line: int
|
|
character: int
|
|
|
|
|
|
@dataclass
|
|
class Selection:
|
|
start: Position
|
|
end: Position
|
|
|
|
|
|
@dataclass
|
|
class Context:
|
|
workspace_path: str
|
|
current_file: Optional[str] = None
|
|
selection: Optional[Selection] = None
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
|
|
|
|
class MCPError(Exception):
|
|
"""Base exception for MCP client errors."""
|
|
pass
|
|
|
|
|
|
class ConnectionError(MCPError):
|
|
"""Raised when connection to MCP server fails."""
|
|
pass
|
|
|
|
|
|
class MCPClient:
|
|
"""Client for interacting with the Model Context Protocol (MCP) server."""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "localhost",
|
|
port: int = 3000,
|
|
timeout: float = 30.0,
|
|
auto_reconnect: bool = True
|
|
):
|
|
"""Initialize the MCP client.
|
|
|
|
Args:
|
|
host: Server hostname
|
|
port: Server port
|
|
timeout: Request timeout in seconds
|
|
auto_reconnect: Whether to automatically reconnect on connection loss
|
|
"""
|
|
self.host = host
|
|
self.port = port
|
|
self.timeout = timeout
|
|
self.auto_reconnect = auto_reconnect
|
|
self.base_url = f"http://{host}:{port}"
|
|
self.ws_url = f"ws://{host}:{port}/mcp"
|
|
self._ws: Optional[WebSocketClientProtocol] = None
|
|
self._http = httpx.AsyncClient(timeout=timeout)
|
|
self._event_handlers: Dict[str, List[callable]] = {}
|
|
|
|
async def __aenter__(self):
|
|
"""Async context manager entry."""
|
|
await self.connect()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
"""Async context manager exit."""
|
|
await self.close()
|
|
|
|
async def connect(self) -> None:
|
|
"""Establish connection to the MCP server."""
|
|
try:
|
|
self._ws = await websockets.connect(self.ws_url)
|
|
await self._ws.send(json.dumps({
|
|
"type": "connect",
|
|
"client": "python-mcp-client",
|
|
"version": "1.0.0"
|
|
}))
|
|
except Exception as e:
|
|
raise ConnectionError(f"Failed to connect to MCP server: {e}")
|
|
|
|
async def close(self) -> None:
|
|
"""Close all connections."""
|
|
if self._ws:
|
|
await self._ws.close()
|
|
await self._http.aclose()
|
|
|
|
async def _request(
|
|
self,
|
|
method: str,
|
|
endpoint: str,
|
|
**kwargs
|
|
) -> Dict[str, Any]:
|
|
"""Make an HTTP request to the MCP server.
|
|
|
|
Args:
|
|
method: HTTP method
|
|
endpoint: API endpoint
|
|
**kwargs: Additional arguments for httpx
|
|
|
|
Returns:
|
|
JSON response from server
|
|
"""
|
|
url = f"{self.base_url}{endpoint}"
|
|
try:
|
|
response = await self._http.request(method, url, **kwargs)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except httpx.HTTPError as e:
|
|
raise MCPError(f"HTTP request failed: {e}")
|
|
|
|
async def send_prompt(
|
|
self,
|
|
prompt: str,
|
|
context: Optional[Union[Context, Dict[str, Any]]] = None
|
|
) -> Dict[str, Any]:
|
|
"""Send a prompt to the MCP server.
|
|
|
|
Args:
|
|
prompt: The prompt text
|
|
context: Optional context information
|
|
|
|
Returns:
|
|
Server response
|
|
"""
|
|
if isinstance(context, Context):
|
|
context = {
|
|
"workspace_path": context.workspace_path,
|
|
"current_file": context.current_file,
|
|
"selection": context.selection.__dict__ if context.selection else None,
|
|
"metadata": context.metadata
|
|
}
|
|
|
|
data = {
|
|
"prompt": prompt,
|
|
"context": context
|
|
}
|
|
|
|
return await self._request("POST", "/api/mcp/prompt", json=data)
|
|
|
|
async def list_profiles(self) -> List[Dict[str, Any]]:
|
|
"""Get list of available profiles.
|
|
|
|
Returns:
|
|
List of profile information
|
|
"""
|
|
response = await self._request("GET", "/api/profiles")
|
|
return response["profiles"]
|
|
|
|
async def save_profile(
|
|
self,
|
|
name: str,
|
|
description: Optional[str] = None
|
|
) -> Dict[str, Any]:
|
|
"""Save current settings as a new profile.
|
|
|
|
Args:
|
|
name: Profile name
|
|
description: Optional profile description
|
|
|
|
Returns:
|
|
Created profile information
|
|
"""
|
|
data = {
|
|
"name": name,
|
|
"description": description
|
|
}
|
|
return await self._request("POST", "/api/profiles", json=data)
|
|
|
|
async def load_profile(self, name: str) -> Dict[str, Any]:
|
|
"""Load a saved profile.
|
|
|
|
Args:
|
|
name: Profile name to load
|
|
|
|
Returns:
|
|
Loaded profile information
|
|
"""
|
|
return await self._request("POST", f"/api/profiles/{name}/load")
|
|
|
|
async def delete_profile(self, name: str) -> None:
|
|
"""Delete a profile.
|
|
|
|
Args:
|
|
name: Profile name to delete
|
|
"""
|
|
await self._request("DELETE", f"/api/profiles/{name}")
|
|
|
|
async def update_profile(
|
|
self,
|
|
name: str,
|
|
description: Optional[str] = None,
|
|
metadata: Optional[Dict[str, Any]] = None
|
|
) -> Dict[str, Any]:
|
|
"""Update profile metadata.
|
|
|
|
Args:
|
|
name: Profile name
|
|
description: New description
|
|
metadata: New metadata
|
|
|
|
Returns:
|
|
Updated profile information
|
|
"""
|
|
data = {
|
|
"description": description,
|
|
"metadata": metadata
|
|
}
|
|
return await self._request(
|
|
"PATCH",
|
|
f"/api/profiles/{name}",
|
|
json={k: v for k, v in data.items() if v is not None}
|
|
)
|
|
|
|
def on(self, event: str, callback: callable) -> None:
|
|
"""Register an event handler.
|
|
|
|
Args:
|
|
event: Event name
|
|
callback: Callback function
|
|
"""
|
|
if event not in self._event_handlers:
|
|
self._event_handlers[event] = []
|
|
self._event_handlers[event].append(callback)
|
|
|
|
async def _handle_events(self) -> None:
|
|
"""Handle incoming WebSocket events."""
|
|
while True:
|
|
try:
|
|
if not self._ws:
|
|
if self.auto_reconnect:
|
|
await self.connect()
|
|
else:
|
|
break
|
|
|
|
message = await self._ws.recv()
|
|
data = json.loads(message)
|
|
event_type = data.get("type")
|
|
|
|
if event_type in self._event_handlers:
|
|
for handler in self._event_handlers[event_type]:
|
|
await handler(data)
|
|
|
|
except websockets.ConnectionClosed:
|
|
if self.auto_reconnect:
|
|
await asyncio.sleep(1)
|
|
continue
|
|
break
|
|
except Exception as e:
|
|
print(f"Error handling WebSocket message: {e}", file=sys.stderr)
|
|
if not self.auto_reconnect:
|
|
break
|
|
|
|
async def start_event_loop(self) -> None:
|
|
"""Start the event handling loop."""
|
|
await self._handle_events()
|
|
|
|
|
|
async def ensure_server_running(
|
|
host: str = "localhost",
|
|
port: int = 3000,
|
|
timeout: float = 30.0
|
|
) -> None:
|
|
"""Ensure the MCP server is running.
|
|
|
|
Args:
|
|
host: Server hostname
|
|
port: Server port
|
|
timeout: Maximum time to wait for server
|
|
"""
|
|
start_time = time.time()
|
|
async with httpx.AsyncClient() as client:
|
|
while True:
|
|
try:
|
|
response = await client.get(
|
|
f"http://{host}:{port}/health",
|
|
timeout=1.0
|
|
)
|
|
if response.status_code == 200:
|
|
return
|
|
except httpx.RequestError:
|
|
if time.time() - start_time > timeout:
|
|
raise TimeoutError("MCP server failed to start")
|
|
await asyncio.sleep(0.1)
|
|
|
|
|
|
async def main():
|
|
"""Main entry point for CLI usage."""
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="MCP Client")
|
|
parser.add_argument(
|
|
"--host",
|
|
default="localhost",
|
|
help="MCP server hostname"
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=3000,
|
|
help="MCP server port"
|
|
)
|
|
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# List profiles
|
|
subparsers.add_parser("list", help="List available profiles")
|
|
|
|
# Save profile
|
|
save_parser = subparsers.add_parser("save", help="Save current settings as profile")
|
|
save_parser.add_argument("name", help="Profile name")
|
|
save_parser.add_argument("-d", "--description", help="Profile description")
|
|
|
|
# Load profile
|
|
load_parser = subparsers.add_parser("load", help="Load a profile")
|
|
load_parser.add_argument("name", help="Profile name")
|
|
|
|
# Delete profile
|
|
delete_parser = subparsers.add_parser("delete", help="Delete a profile")
|
|
delete_parser.add_argument("name", help="Profile name")
|
|
|
|
# Update profile
|
|
update_parser = subparsers.add_parser("update", help="Update profile metadata")
|
|
update_parser.add_argument("name", help="Profile name")
|
|
update_parser.add_argument("-d", "--description", help="New description")
|
|
update_parser.add_argument(
|
|
"-m",
|
|
"--metadata",
|
|
type=json.loads,
|
|
help="New metadata (as JSON)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
await ensure_server_running(args.host, args.port)
|
|
except TimeoutError:
|
|
print("Error: MCP server is not running", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
async with MCPClient(args.host, args.port) as client:
|
|
if args.command == "list":
|
|
profiles = await client.list_profiles()
|
|
print("\nAvailable Profiles:")
|
|
for profile in profiles:
|
|
print(f"\n{profile['name']}")
|
|
if profile.get("description"):
|
|
print(f" Description: {profile['description']}")
|
|
if profile.get("metadata"):
|
|
print(f" Metadata: {json.dumps(profile['metadata'], indent=2)}")
|
|
|
|
elif args.command == "save":
|
|
profile = await client.save_profile(args.name, args.description)
|
|
print(f"\nSaved profile: {profile['name']}")
|
|
|
|
elif args.command == "load":
|
|
profile = await client.load_profile(args.name)
|
|
print(f"\nLoaded profile: {profile['name']}")
|
|
|
|
elif args.command == "delete":
|
|
await client.delete_profile(args.name)
|
|
print(f"\nDeleted profile: {args.name}")
|
|
|
|
elif args.command == "update":
|
|
profile = await client.update_profile(
|
|
args.name,
|
|
args.description,
|
|
args.metadata
|
|
)
|
|
print(f"\nUpdated profile: {profile['name']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
print("\nOperation cancelled by user", file=sys.stderr)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"\nError: {e}", file=sys.stderr)
|
|
sys.exit(1) |
