gpt-edit/src/server/mcp_server.py

"""MCP Server implementation for GPTEdit"""

import logging
import sys
from typing import Dict, Any, List, Union, Optional

from mcp.server import Server
from mcp.types import (
    Tool, 
    TextContent, 
    ImageContent,
    Prompt,
    PromptMessage,
    Resource
)

from ..connector import Config
from .models import MCPToolDefinitions
from .handlers import ToolHandlers, sanitize_args_for_logging

# Set up logger to use stderr to avoid interfering with stdout
logger = logging.getLogger(__name__)


class GPTEditMCPServer:
    """GPTEdit MCP server class"""
    
    def __init__(self, config: Config):
        """Initialize server"""
        self.config = config
        self.server = Server("gpt-edit")  # Simplified server name
        self.handlers = ToolHandlers(config)
        
        # Register handlers
        self._register_handlers()
        
        logger.info("GPT-Edit MCP Server initialized")
        logger.info(f"Model: {Config.MODEL}")
        logger.info(f"Max image size: {config.max_image_size_mb}MB")
        logger.info(f"Server name: gpt-edit")
    
    def _register_handlers(self) -> None:
        """Register MCP handlers"""
        
        @self.server.list_tools()
        async def handle_list_tools() -> List[Tool]:
            """Return list of available tools"""
            logger.debug("list_tools called")
            tools = MCPToolDefinitions.get_all_tools()
            logger.info(f"Returning {len(tools)} tools")
            return tools
        
        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[Union[TextContent, ImageContent]]:
            """Handle tool calls"""
            # Log tool call safely without exposing sensitive data
            safe_args = sanitize_args_for_logging(arguments)
            logger.info(f"Tool called: {name} with arguments: {safe_args}")
            
            try:
                if name == "gpt_batch_edit":
                    return await self.handlers.handle_gpt_batch_edit(arguments)
                elif name == "validate_image":
                    return await self.handlers.handle_validate_image(arguments)
                elif name == "create_mask_from_alpha":
                    return await self.handlers.handle_create_mask_from_alpha(arguments)
                elif name == "gpt_edit_image":
                    return await self.handlers.handle_gpt_edit_image(arguments)
                elif name == "gpt_edit_image_with_mask":
                    return await self.handlers.handle_gpt_edit_image_with_mask(arguments)
                else:
                    error_msg = f"Unknown tool: {name}"
                    logger.error(error_msg)
                    return [TextContent(type="text", text=f"[ERROR] {error_msg}")]
                    
            except Exception as e:
                logger.error(f"Error handling tool {name}: {e}", exc_info=True)
                return [TextContent(
                    type="text",
                    text=f"[ERROR] Error processing {name}: {str(e)}"
                )]
        
        @self.server.list_prompts()
        async def handle_list_prompts() -> List[Prompt]:
            """Return list of available prompts"""
            logger.debug("list_prompts called")
            prompts = [
                Prompt(
                    name="gpt_edit_image_with_mask",
                    description="Edit an image with a mask",
                    arguments=[
                        {
                            "name": "image_path",
                            "description": "Path to the image to edit",
                            "required": True
                        },
                        {
                            "name": "mask_path",
                            "description": "Path to the mask image",
                            "required": True
                        },
                        {
                            "name": "edit_description",
                            "description": "Description of how to edit the masked areas",
                            "required": True
                        }
                    ]
                ),
                Prompt(
                    name="gpt_batch_edit",
                    description="Edit multiple images with batch processing (for automation/Claude Code)",
                    arguments=[
                        {
                            "name": "image_paths",
                            "description": "Comma-separated list of image paths to edit",
                            "required": True
                        },
                        {
                            "name": "edit_description",
                            "description": "Description of how to edit the images",
                            "required": True
                        }
                    ]
                )
            ]
            logger.info(f"Returning {len(prompts)} prompts")
            return prompts
        
        @self.server.get_prompt()
        async def handle_get_prompt(name: str, arguments: Optional[Dict[str, Any]] = None) -> List[PromptMessage]:
            """Get a specific prompt"""
            logger.debug(f"get_prompt called for: {name}")
            
            if name == "gpt_edit_image_with_mask":
                if not arguments or "image_path" not in arguments or "mask_path" not in arguments or "edit_description" not in arguments:
                    return [PromptMessage(
                        role="user",
                        content=TextContent(
                            type="text",
                            text="Please provide image_path, mask_path, and edit_description"
                        )
                    )]
                
                return [PromptMessage(
                    role="user",
                    content=TextContent(
                        type="text",
                        text=f"Edit the image at '{arguments['image_path']}' using mask '{arguments['mask_path']}' with: {arguments['edit_description']}"
                    )
                )]
            
            elif name == "gpt_batch_edit":
                if not arguments or "image_paths" not in arguments or "edit_description" not in arguments:
                    return [PromptMessage(
                        role="user",
                        content=TextContent(
                            type="text",
                            text="Please provide image_paths and edit_description"
                        )
                    )]
                
                return [PromptMessage(
                    role="user",
                    content=TextContent(
                        type="text",
                        text=f"Batch edit the images at '{arguments['image_paths']}' with: {arguments['edit_description']}"
                    )
                )]
            
            else:
                return [PromptMessage(
                    role="user",
                    content=TextContent(
                        type="text",
                        text=f"Unknown prompt: {name}"
                    )
                )]
        
        @self.server.list_resources()
        async def handle_list_resources() -> List[Resource]:
            """Return list of available resources"""
            logger.debug("list_resources called")
            # GPTEdit doesn't expose file resources directly, return empty list
            resources = []
            logger.info(f"Returning {len(resources)} resources")
            return resources
        
        # Note: read_resource is not implemented as GPTEdit doesn't expose resources
        # The MCP server will handle the "not implemented" response automatically
        
        # Log all registered handlers
        logger.info("Registered MCP handlers:")
        logger.info("  - list_tools")
        logger.info("  - call_tool")
        logger.info("  - list_prompts")
        logger.info("  - get_prompt")
        logger.info("  - list_resources")
    
    def get_server(self) -> Server:
        """Return MCP server instance"""
        return self.server