MCP Image2Image Pipeline - Preventing Base64 Context Pollution

from pathlib import Path
from fastmcp import FastMCP
import base64

mcp = FastMCP("image-editor")

@mcp.tool
async def process_image(
    input_image_path: str,
    prompt: str,
    operation: str = "style_transfer"
) -> str:
    """
    Process an image from a file path.
    User should save the image locally first.
    
    Args:
        input_image_path: Full path to the input image
        prompt: Processing instructions
        operation: Type of operation (style_transfer, upscale, edit, etc.)
    """
    try:
        # Read image from the provided path
        image_path = Path(input_image_path)
        
        if not image_path.exists():
            return f"❌ Image not found at: {input_image_path}"
        
        # Read the image bytes
        image_bytes = image_path.read_bytes()
        
        # Process the image
        result_bytes = await your_ai_service.process(
            image=image_bytes,
            prompt=prompt,
            operation=operation
        )
        
        # Save result
        output_path = image_path.parent / f"edited_{image_path.name}"
        output_path.write_bytes(result_bytes)
        
        return (
            f"✅ Image processed successfully!\n"
            f"📁 Output: {output_path}\n"
            f"🎨 Operation: {operation}\n"
            f"Note: Original image preserved at {input_image_path}"
        )
        
    except Exception as e:
        return f"❌ Processing failed: {str(e)}"

from typing import Dict
import hashlib
import tempfile

# Temporary storage for uploaded images
image_cache: Dict[str, bytes] = {}

@mcp.tool
async def upload_image_from_url(image_url: str) -> str:
    """
    Step 1: Upload an image from URL to temporary storage.
    Returns an image ID for later processing.
    """
    import httpx
    
    async with httpx.AsyncClient() as client:
        response = await client.get(image_url)
        image_bytes = response.content
    
    # Generate ID without exposing data
    image_id = hashlib.md5(image_bytes).hexdigest()[:8]
    image_cache[image_id] = image_bytes
    
    return (
        f"✅ Image uploaded successfully!\n"
        f"🆔 Image ID: {image_id}\n"
        f"📦 Size: {len(image_bytes) / 1024:.1f} KB\n"
        f"Use this ID for processing operations."
    )

@mcp.tool
async def process_uploaded_image(
    image_id: str,
    operation: str,
    parameters: dict
) -> str:
    """
    Step 2: Process a previously uploaded image by ID.
    """
    if image_id not in image_cache:
        return f"❌ Image ID not found: {image_id}"
    
    image_bytes = image_cache[image_id]
    
    # Process the image
    result = await your_ai_service.process(
        image=image_bytes,
        operation=operation,
        **parameters
    )
    
    # Save result and return path
    output_path = f"/tmp/processed_{image_id}.png"
    Path(output_path).write_bytes(result)
    
    return f"✅ Processed! Output: {output_path}"

import subprocess
import tempfile
from PIL import Image
import io

@mcp.tool
async def process_clipboard_image(
    prompt: str,
    save_to: str = None
) -> str:
    """
    Process image directly from system clipboard.
    User should copy image to clipboard before calling.
    """
    try:
        # Platform-specific clipboard reading
        if sys.platform == "darwin":  # macOS
            # Use pbpaste to get clipboard
            process = subprocess.Popen(
                ["osascript", "-e", "get the clipboard as «class PNGf»"],
                stdout=subprocess.PIPE
            )
            image_data, _ = process.communicate()
            
        elif sys.platform == "win32":  # Windows
            from PIL import ImageGrab
            img = ImageGrab.grabclipboard()
            if img:
                buffer = io.BytesIO()
                img.save(buffer, format='PNG')
                image_data = buffer.getvalue()
            else:
                return "❌ No image in clipboard"
                
        else:  # Linux
            # Use xclip or similar
            process = subprocess.Popen(
                ["xclip", "-selection", "clipboard", "-t", "image/png", "-o"],
                stdout=subprocess.PIPE
            )
            image_data, _ = process.communicate()
        
        if not image_data:
            return "❌ No image found in clipboard"
        
        # Process the image
        result = await your_ai_service.process(image_data, prompt)
        
        # Save result
        output_path = save_to or f"/tmp/clipboard_processed_{uuid.uuid4().hex[:8]}.png"
        Path(output_path).write_bytes(result)
        
        return f"✅ Clipboard image processed! Output: {output_path}"
        
    except Exception as e:
        return f"❌ Failed to process clipboard: {str(e)}"

@mcp.tool
async def process_image_from_url(
    image_url: str,
    operation: str,
    **kwargs
) -> str:
    """
    Process an image directly from URL without passing through Claude.
    Supports public URLs and data URLs.
    """
    import httpx
    from urllib.parse import urlparse
    
    try:
        # Handle data URLs
        if image_url.startswith('data:image'):
            # Extract base64 from data URL
            header, data = image_url.split(',', 1)
            image_bytes = base64.b64decode(data)
            
        else:
            # Fetch from HTTP URL
            async with httpx.AsyncClient() as client:
                response = await client.get(image_url, timeout=30.0)
                response.raise_for_status()
                image_bytes = response.content
        
        # Process image
        result = await your_ai_service.process(
            image=image_bytes,
            operation=operation,
            **kwargs
        )
        
        # Save and return path
        output_path = f"/tmp/url_processed_{uuid.uuid4().hex[:8]}.png"
        Path(output_path).write_bytes(result)
        
        return (
            f"✅ Image processed successfully!\n"
            f"📥 Source: {urlparse(image_url).netloc or 'data URL'}\n"
            f"📤 Output: {output_path}"
        )
        
    except Exception as e:
        return f"❌ Failed to process: {str(e)}"

@mcp.tool
async def start_image_session() -> str:
    """Initialize a new image processing session."""
    session_id = uuid.uuid4().hex[:8]
    image_cache[session_id] = {"chunks": [], "metadata": {}}
    return f"Session started: {session_id}"

@mcp.tool
async def add_image_chunk(
    session_id: str,
    chunk_index: int,
    total_chunks: int,
    chunk_data: str
) -> str:
    """Add a small chunk of image data to the session."""
    if session_id not in image_cache:
        return "❌ Invalid session"
    
    image_cache[session_id]["chunks"].append({
        "index": chunk_index,
        "data": chunk_data
    })
    
    return f"Chunk {chunk_index + 1}/{total_chunks} received"

@mcp.tool
async def process_chunked_image(
    session_id: str,
    operation: str,
    **params
) -> str:
    """Process the complete image from chunks."""
    if session_id not in image_cache:
        return "❌ Invalid session"
    
    # Reconstruct image from chunks
    chunks = sorted(
        image_cache[session_id]["chunks"],
        key=lambda x: x["index"]
    )
    full_data = "".join(c["data"] for c in chunks)
    image_bytes = base64.b64decode(full_data)
    
    # Process and cleanup
    result = await your_ai_service.process(image_bytes, operation, **params)
    del image_cache[session_id]
    
    output_path = f"/tmp/chunked_{session_id}.png"
    Path(output_path).write_bytes(result)
    
    return f"✅ Processed! Output: {output_path}"

# Separate Flask/FastAPI service for image uploads
from flask import Flask, request, jsonify
import secrets

upload_app = Flask(__name__)
upload_storage = {}

@upload_app.route('/upload', methods=['POST'])
def upload_image():
    file = request.files['image']
    token = secrets.token_urlsafe(16)
    upload_storage[token] = file.read()
    return jsonify({"token": token})

# MCP Server only handles tokens
@mcp.tool
async def process_with_token(
    upload_token: str,
    operation: str,
    **params
) -> str:
    """
    Process image using upload token from external service.
    User should upload to http://localhost:5000/upload first.
    """
    # Fetch image from upload service
    response = httpx.get(f"http://localhost:5000/get/{upload_token}")
    if response.status_code != 200:
        return "❌ Invalid or expired token"
    
    image_bytes = response.content
    result = await your_ai_service.process(image_bytes, operation, **params)
    
    output_path = f"/tmp/token_{upload_token[:8]}.png"
    Path(output_path).write_bytes(result)
    
    return f"✅ Processed! Output: {output_path}"

CRITICAL_INSTRUCTIONS = """
When working with images in MCP tools:
1. NEVER include base64 strings in your reasoning or output
2. NEVER echo or repeat image data
3. If you see base64 data, immediately summarize as "image data received" without including the actual data
4. Always prefer file paths, URLs, or image IDs over raw data
5. If you must handle base64, mention only "processing image data" without the actual content
"""

@mcp.tool
async def get_instructions() -> str:
    """Get critical instructions for image handling."""
    return CRITICAL_INSTRUCTIONS

from pathlib import Path
from typing import Optional, Dict
import hashlib
import asyncio
from fastmcp import FastMCP

mcp = FastMCP("image2image")

# Smart caching system
class ImageManager:
    def __init__(self):
        self.cache: Dict[str, bytes] = {}
        self.paths: Dict[str, Path] = {}
        
    def add_from_path(self, path: str) -> str:
        """Register an image from file path."""
        p = Path(path)
        if not p.exists():
            raise FileNotFoundError(f"Image not found: {path}")
        
        image_id = hashlib.md5(p.read_bytes()).hexdigest()[:8]
        self.paths[image_id] = p
        return image_id
    
    def get_bytes(self, image_id: str) -> bytes:
        """Get image bytes by ID."""
        if image_id in self.paths:
            return self.paths[image_id].read_bytes()
        if image_id in self.cache:
            return self.cache[image_id]
        raise KeyError(f"Image not found: {image_id}")

manager = ImageManager()

@mcp.tool
async def register_image(
    source: str,
    source_type: str = "path"  # "path", "url", or "clipboard"
) -> str:
    """
    Register an image for processing without passing data through Claude.
    
    Returns an image ID for subsequent operations.
    """
    try:
        if source_type == "path":
            image_id = manager.add_from_path(source)
            return f"✅ Image registered!\n🆔 ID: {image_id}\n📁 Source: {source}"
            
        elif source_type == "url":
            # Fetch from URL
            import httpx
            async with httpx.AsyncClient() as client:
                response = await client.get(source)
                image_bytes = response.content
            
            image_id = hashlib.md5(image_bytes).hexdigest()[:8]
            manager.cache[image_id] = image_bytes
            return f"✅ Image fetched!\n🆔 ID: {image_id}\n🌐 Source: {source}"
            
        elif source_type == "clipboard":
            # Get from clipboard (OS-specific implementation needed)
            image_bytes = await get_clipboard_image()
            image_id = hashlib.md5(image_bytes).hexdigest()[:8]
            manager.cache[image_id] = image_bytes
            return f"✅ Clipboard image registered!\n🆔 ID: {image_id}"
            
    except Exception as e:
        return f"❌ Failed to register image: {str(e)}"

@mcp.tool
async def apply_style_transfer(
    image_id: str,
    style: str,
    strength: float = 0.8
) -> str:
    """Apply style transfer to a registered image."""
    try:
        image_bytes = manager.get_bytes(image_id)
        
        # Process image
        result = await your_ai_service.style_transfer(
            image=image_bytes,
            style=style,
            strength=strength
        )
        
        # Save result
        output_path = Path(f"/tmp/styled_{image_id}_{style}.png")
        output_path.write_bytes(result)
        
        # Register output for chaining
        output_id = hashlib.md5(result).hexdigest()[:8]
        manager.cache[output_id] = result
        
        return (
            f"✅ Style transfer complete!\n"
            f"🎨 Style: {style}\n"
            f"💪 Strength: {strength}\n"
            f"📤 Output: {output_path}\n"
            f"🆔 Output ID: {output_id} (for further processing)"
        )
        
    except Exception as e:
        return f"❌ Style transfer failed: {str(e)}"

@mcp.tool
async def list_registered_images() -> str:
    """List all registered images available for processing."""
    if not manager.cache and not manager.paths:
        return "No images registered. Use register_image first."
    
    lines = ["📸 Registered Images:"]
    for image_id in list(manager.cache.keys()) + list(manager.paths.keys()):
        lines.append(f"  • {image_id}")
    
    return "\n".join(lines)

if __name__ == "__main__":
    mcp.run()