decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

dt_arena/mcp_server/macos/client_test.py

@@ -0,0 +1,270 @@
+import asyncio
+import base64
+from typing import Optional, List, Dict, Any
+from contextlib import AsyncExitStack
+from pathlib import Path
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from anthropic import Anthropic
+from dotenv import load_dotenv
+
+load_dotenv()  # load environment variables from .env
+
+
+class MCPClient:
+    def __init__(self):
+        # Initialize session and client objects
+        self.session: Optional[ClientSession] = None
+        self.exit_stack = AsyncExitStack()
+        self.anthropic = Anthropic()
+
+    async def connect_to_server(self, server_script_path: str, extra_args: Optional[List[str]] = None):
+        """Connect to an MCP server
+
+        Args:
+            server_script_path: Path to the server script (.py or .js)
+            extra_args: Optional additional arguments to pass to the server
+        """
+        is_python = server_script_path.endswith(".py")
+        is_js = server_script_path.endswith(".js")
+        if not (is_python or is_js):
+            raise ValueError("Server script must be a .py or .js file")
+
+        command = "python" if is_python else "node"
+        args = [server_script_path]
+        if extra_args:
+            args.extend(extra_args)
+
+        server_params = StdioServerParameters(command=command, args=args, env=None)
+
+        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
+        self.stdio, self.write = stdio_transport
+        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
+
+        await self.session.initialize()
+
+        # List available tools
+        response = await self.session.list_tools()
+        tools = response.tools
+        print("\nConnected to server with tools:", [tool.name for tool in tools])
+
+    def encode_image(self, image_path: str) -> Dict[str, Any]:
+        """Encode an image file to base64 and determine its media type
+
+        Args:
+            image_path: Path to the image file
+
+        Returns:
+            Dictionary with image data and media type for Anthropic API
+        """
+        path = Path(image_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+
+        # Determine media type based on file extension
+        extension = path.suffix.lower()
+        media_type_map = {
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".png": "image/png",
+            ".gif": "image/gif",
+            ".webp": "image/webp",
+        }
+
+        media_type = media_type_map.get(extension)
+        if not media_type:
+            raise ValueError(f"Unsupported image format: {extension}. Supported formats: {list(media_type_map.keys())}")
+
+        # Read and encode the image
+        with open(path, "rb") as image_file:
+            image_data = base64.standard_b64encode(image_file.read()).decode("utf-8")
+
+        return {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": image_data}}
+
+    def convert_mcp_content_to_anthropic(self, content: Any) -> Any:
+        """Convert MCP tool result content to Anthropic API format
+
+        Args:
+            content: Content from MCP tool result
+
+        Returns:
+            Converted content in Anthropic API format
+        """
+        # If content is a list, convert each item
+        if isinstance(content, list):
+            converted = []
+            for item in content:
+                # Check if this is an MCP image format
+                if item.type == "image":
+                    # Convert MCP image format to Anthropic format
+                    converted.append(
+                        {"type": "image", "source": {"type": "base64", "media_type": item.mimeType, "data": item.data}}
+                    )
+                elif item.type == "text":
+                    # Text content remains the same
+                    converted.append(item)
+                else:
+                    # Unknown format, keep as is
+                    converted.append(item)
+            return converted
+        else:
+            # If content is not a list, return as is
+            return content
+
+    async def process_query(self, query: str, image_paths: Optional[List[str]] = None) -> str:
+        """Process a query using Claude and available tools
+
+        Args:
+            query: Text query to process
+            image_paths: Optional list of image file paths to include in the query
+        """
+        # Build message content with text and images
+        content: List[Dict[str, Any]] = []
+
+        # Add images first if provided
+        if image_paths:
+            for image_path in image_paths:
+                try:
+                    image_content = self.encode_image(image_path)
+                    content.append(image_content)
+                    print(f"Loaded image: {image_path}")
+                except Exception as e:
+                    print(f"Warning: Failed to load image {image_path}: {str(e)}")
+
+        # Add text query
+        content.append({"type": "text", "text": query})
+
+        messages = [{"role": "user", "content": content}]
+
+        response = await self.session.list_tools()
+        available_tools = [
+            {
+                "name": tool.name,
+                "description": tool.description,
+                "input_schema": tool.inputSchema,
+            }
+            for tool in response.tools
+        ]
+
+        # Initial Claude API call
+        response = self.anthropic.messages.create(
+            model="claude-sonnet-4-5",
+            max_tokens=1000,
+            messages=messages,
+            tools=available_tools,
+        )
+
+        # Process response and handle tool calls
+        final_text = []
+
+        # Agentic loop: continue while there are tool calls
+        while response.stop_reason == "tool_use":
+            # Collect all content from assistant's response
+            assistant_content = []
+            tool_results = []
+
+            for content_block in response.content:
+                if content_block.type == "text":
+                    final_text.append(content_block.text)
+                    assistant_content.append({"type": "text", "text": content_block.text})
+                elif content_block.type == "tool_use":
+                    tool_name = content_block.name
+                    tool_args = content_block.input
+                    tool_use_id = content_block.id
+
+                    assistant_content.append(
+                        {"type": "tool_use", "id": tool_use_id, "name": tool_name, "input": tool_args}
+                    )
+
+                    # Execute tool call
+                    print(f"[Calling tool {tool_name} with args {tool_args}]")
+                    result = await self.session.call_tool(tool_name, tool_args)
+
+                    # Convert MCP content format to Anthropic format
+                    converted_content = self.convert_mcp_content_to_anthropic(result.content)
+
+                    # Add tool result
+                    tool_results.append(
+                        {"type": "tool_result", "tool_use_id": tool_use_id, "content": converted_content}
+                    )
+
+            # Add assistant message with all content (text + tool_use)
+            messages.append({"role": "assistant", "content": assistant_content})
+
+            # Add user message with tool results
+            messages.append({"role": "user", "content": tool_results})
+
+            # Get next response from Claude
+            response = self.anthropic.messages.create(
+                model="claude-sonnet-4-5",
+                max_tokens=1000,
+                messages=messages,
+                tools=available_tools,
+            )
+
+        # Process final response (no more tool calls)
+        for content_block in response.content:
+            if content_block.type == "text":
+                final_text.append(content_block.text)
+
+        return "\n".join(final_text)
+
+    async def chat_loop(self):
+        """Run an interactive chat loop"""
+        print("\nMCP Client Started!")
+        print("Type your queries or 'quit' to exit.")
+        print("To include images, use: your query --image path/to/image.jpg --image path/to/another.png")
+
+        while True:
+            try:
+                query = input("\nQuery: ").strip()
+                if query.lower() == "quit":
+                    break
+
+                # Parse image paths from query
+                image_paths = []
+                parts = query.split("--image")
+
+                if len(parts) > 1:
+                    # First part is the actual query
+                    query = parts[0].strip()
+
+                    # Remaining parts are image paths
+                    for part in parts[1:]:
+                        image_path = part.strip().split()[0] if part.strip() else ""
+                        if image_path:
+                            image_paths.append(image_path)
+
+                response = await self.process_query(query, image_paths if image_paths else None)
+                print("\n" + response)
+
+            except Exception as e:
+                print(f"\nError: {str(e)}")
+
+    async def cleanup(self):
+        """Clean up resources"""
+        await self.exit_stack.aclose()
+
+
+async def main():
+    if len(sys.argv) < 2:
+        print("Usage: python client_test.py <path_to_server_script> [additional_args...]")
+        sys.exit(1)
+
+    server_script = sys.argv[1]
+    extra_args = sys.argv[2:] if len(sys.argv) > 2 else None
+
+    client = MCPClient()
+    try:
+        await client.connect_to_server(server_script, extra_args)
+        await client.chat_loop()
+    finally:
+        await client.cleanup()
+
+
+if __name__ == "__main__":
+    import sys
+
+    asyncio.run(main())

dt_arena/mcp_server/macos/mcp_server.py

@@ -0,0 +1,285 @@
+import base64
+import os
+from textwrap import dedent
+from typing import Literal, Optional
+
+import click
+import requests
+from fastmcp import FastMCP
+from fastmcp.utilities.types import Image
+
+
+API_BASE_URL = os.environ.get("MACOS_API_URL", "http://localhost:8005")
+
+instructions = dedent("""
+macOS MCP client provides tools to interact with macOS desktop through a FastAPI backend service.
+All operations are executed via HTTP requests to the backend server.
+""")
+
+mcp = FastMCP(name="macos-mcp-client", instructions=instructions)
+
+
+def make_api_call(endpoint: str, data: dict = None) -> dict:
+    """Helper function to make API calls to the FastAPI server"""
+    try:
+        response = requests.post(f"{API_BASE_URL}{endpoint}", json=data or {})
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return {"status": "error", "result": f"API call failed: {str(e)}"}
+
+
+def _extract_text(result: dict) -> str:
+    """Extract text from API response (handles both list and string formats)."""
+    data = result.get("result", result)
+    if isinstance(data, list) and len(data) > 0:
+        return data[0].get("text", str(data))
+    return str(data)
+
+
+@mcp.tool(
+    name="launch",
+    description='Launch an application on macOS by name, path, or bundle ID (e.g., "Terminal", "Safari", "Finder").',
+)
+def launch_tool(name: str) -> str:
+    result = make_api_call("/remote_macos_open_application", {"identifier": name})
+    text = _extract_text(result)
+    return (
+        text
+        + "\nNote: The application launch has been triggered but may not be fully loaded yet. Use screenshot to verify the app is ready before interacting with it."
+    )
+
+
+@mcp.tool(
+    name="shell",
+    description="Execute shell commands on the macOS VM via SSH and return the output.",
+)
+def shell_tool(command: str) -> str:
+    result = make_api_call("/shell", {"command": command})
+    if result.get("status") == "error":
+        return result.get("result", str(result))
+    stdout = result.get("stdout", "")
+    stderr = result.get("stderr", "")
+    exit_code = result.get("exit_code", 0)
+    output = stdout
+    if stderr:
+        output += f"\nSTDERR: {stderr}"
+    return f"Response: {output}\nStatus Code: {exit_code}"
+
+
+@mcp.tool(
+    name="screenshot",
+    description="Capture a screenshot of the macOS desktop. Returns visual screenshot image and text status.",
+)
+def state_tool(use_vision: bool = False):
+    result = make_api_call("/remote_macos_get_screen")
+
+    if result.get("status") == "error":
+        return result.get("result", str(result))
+
+    data = result.get("result", [])
+    if isinstance(data, list):
+        response = []
+        for item in data:
+            if item.get("type") == "text":
+                response.append(item.get("text", ""))
+            elif item.get("type") == "image":
+                image_data = item.get("data", "")
+                image_bytes = base64.b64decode(image_data)
+                mime_type = item.get("mimeType", "image/png")
+                image_format = mime_type.split("/")[-1] if "/" in mime_type else "png"
+                response.append(Image(data=image_bytes, format=image_format))
+        return response
+    return str(data)
+
+
+@mcp.tool(
+    name="click",
+    description="Click on UI elements at specific coordinates. Supports left/right/middle mouse buttons and single/double clicks. Use coordinates from screenshot output.",
+)
+def click_tool(
+    loc: list[int],
+    button: Literal["left", "right", "middle"] = "left",
+    clicks: int = 1,
+) -> str:
+    button_map = {"left": 1, "right": 3, "middle": 2}
+    btn = button_map.get(button, 1)
+
+    if clicks >= 2:
+        endpoint = "/remote_macos_mouse_double_click"
+    else:
+        endpoint = "/remote_macos_mouse_click"
+
+    result = make_api_call(
+        endpoint,
+        {
+            "x": loc[0],
+            "y": loc[1],
+            "button": btn,
+        },
+    )
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="key",
+    description='Press keys, key combinations, or type text into the currently focused element. Use "+" to combine keys (e.g., "cmd+c" for copy, "cmd+tab" for app switching). Single special keys: "enter", "escape", "tab", "space", "backspace", "delete", arrow keys, function keys. For plain text (e.g., a filename or command), just pass the text string directly.',
+)
+def key_tool(key: str) -> str:
+    payload = {}
+    if "+" in key:
+        payload["key_combination"] = key
+    else:
+        # Check if it's a special key or regular text
+        special_keys = {
+            "enter",
+            "return",
+            "escape",
+            "tab",
+            "space",
+            "backspace",
+            "delete",
+            "up",
+            "down",
+            "left",
+            "right",
+            "f1",
+            "f2",
+            "f3",
+            "f4",
+            "f5",
+            "f6",
+            "f7",
+            "f8",
+            "f9",
+            "f10",
+            "f11",
+            "f12",
+            "home",
+            "end",
+            "pageup",
+            "pagedown",
+        }
+        if key.lower() in special_keys:
+            payload["special_key"] = key
+        else:
+            payload["text"] = key
+
+    result = make_api_call("/remote_macos_send_keys", payload)
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="type",
+    description="Type text into input fields, text areas, or focused elements. Optionally provide loc coordinates to click and focus a target element first. If loc is omitted, types into the currently focused element. Set clear=True to replace existing text.",
+)
+def type_tool(text: str, loc: list[int] = None, clear: bool = False, press_enter: bool = False) -> str:
+    # Click to focus the element only if loc is provided
+    if loc is not None:
+        make_api_call("/remote_macos_mouse_click", {"x": loc[0], "y": loc[1], "button": 1})
+    if clear:
+        make_api_call("/remote_macos_send_keys", {"key_combination": "cmd+a"})
+    # Type the text
+    result = make_api_call("/remote_macos_send_keys", {"text": text})
+    if press_enter:
+        make_api_call("/remote_macos_send_keys", {"special_key": "enter"})
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="scroll",
+    description="Scroll at specific coordinates. Use wheel_times to control scroll amount.",
+)
+def scroll_tool(
+    loc: list[int] = None,
+    direction: Literal["up", "down"] = "down",
+    wheel_times: int = 1,
+) -> str:
+    x = loc[0] if loc else 683
+    y = loc[1] if loc else 384
+    # Repeat scroll for wheel_times
+    for _ in range(wheel_times):
+        result = make_api_call(
+            "/remote_macos_mouse_scroll",
+            {
+                "x": x,
+                "y": y,
+                "direction": direction,
+            },
+        )
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="drag",
+    description="Drag and drop from source to destination coordinates.",
+)
+def drag_tool(from_loc: list[int], to_loc: list[int]) -> str:
+    result = make_api_call(
+        "/remote_macos_mouse_drag_n_drop",
+        {
+            "start_x": from_loc[0],
+            "start_y": from_loc[1],
+            "end_x": to_loc[0],
+            "end_y": to_loc[1],
+        },
+    )
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="move",
+    description="Move mouse cursor to specific coordinates without clicking.",
+)
+def move_tool(to_loc: list[int]) -> str:
+    result = make_api_call(
+        "/remote_macos_mouse_move",
+        {
+            "x": to_loc[0],
+            "y": to_loc[1],
+        },
+    )
+    return _extract_text(result)
+
+
+@mcp.tool(
+    name="wait",
+    description="Pause execution for specified duration in seconds.",
+)
+def wait_tool(duration: int) -> str:
+    import time
+
+    time.sleep(duration)
+    return f"Waited for {duration} seconds"
+
+
+@click.command()
+@click.option(
+    "--transport",
+    help="The transport layer used by the MCP server.",
+    type=click.Choice(["stdio", "sse", "streamable-http"]),
+    default="stdio",
+)
+@click.option(
+    "--host", help="Host to bind the SSE/Streamable HTTP server.", default="localhost", type=str, show_default=True
+)
+@click.option("--port", help="Port to bind the SSE/Streamable HTTP server.", default=8002, type=int, show_default=True)
+@click.option(
+    "--api-url",
+    help="URL of the FastAPI backend server. If not provided, uses MACOS_API_URL env var.",
+    default=None,
+    type=str,
+)
+def main(transport, host, port, api_url):
+    global API_BASE_URL
+    if api_url:
+        API_BASE_URL = api_url
+
+    if transport == "stdio":
+        mcp.run()
+    else:
+        mcp.run(transport=transport, host=host, port=port)
+
+
+if __name__ == "__main__":
+    main()