"""Asset viewing tools for ComfyUI MCP Server
Thin MCP wrapper around AssetOrchestrator for asset viewing and retrieval.
NO business logic - delegates to orchestrator.
"""
from fastmcp import FastMCP
from src.orchestrators.asset import AssetOrchestrator
from src.utils import get_global_logger
from src.utils.asset_processor import (
encode_preview_for_mcp,
estimate_response_chars,
get_cache_key,
)
# Try to import Image from fastmcp if available (v3 feature)
try:
from fastmcp.media import Image as FastMCPImage
except (ImportError, ModuleNotFoundError):
FastMCPImage = None # Use fallback dict format
logger = get_global_logger("MCP_Server.tools.asset")
def register_asset_tools(mcp: FastMCP, asset_orchestrator: AssetOrchestrator):
"""Register asset viewing tools with the MCP server
Args:
mcp: FastMCP server instance
asset_orchestrator: Orchestrator for asset operations
"""
@mcp.tool()
def view_image(
asset_id: str,
mode: str = "thumb",
max_dim: int | None = None,
max_b64_chars: int | None = None,
) -> dict:
"""View a generated image inline in chat (thumbnail preview only).
This tool allows the AI agent to view generated images inline in the chat interface,
enabling closed-loop iteration: generate → view → adjust → regenerate.
Only supports image assets (PNG, JPEG, WebP, GIF). For audio/video assets, use the
asset_url directly or implement separate viewing tools.
Args:
asset_id: Asset ID returned from generation tools (e.g., generate_image)
mode: Display mode - "thumb" (thumbnail preview, default) or "metadata" (info only)
max_dim: Maximum dimension in pixels (default: 512, hard cap)
max_b64_chars: Maximum base64 character count (default: 100000, ~100KB)
Returns:
MCP ImageContent structure for inline display, or metadata dict if mode="metadata"
or if image exceeds budget (refuse-inline branch).
"""
# Cleanup expired assets periodically
asset_orchestrator.cleanup_expired()
# Validate asset_id exists in registry (security: only our assets)
asset_record = asset_orchestrator.get_asset_record(asset_id)
if not asset_record:
return {
"error": f"Asset {asset_id} not found (registry is in-memory and resets on restart). "
"Generate a new asset to regenerate."
}
# Get asset URL (computed from stable identity)
asset_url = asset_orchestrator.get_asset_url(asset_record)
# If metadata mode, return info only
if mode == "metadata":
return {
"asset_id": asset_record.asset_id,
"asset_url": asset_url,
"filename": asset_record.filename,
"subfolder": asset_record.subfolder,
"folder_type": asset_record.folder_type,
"mime_type": asset_record.mime_type,
"width": asset_record.width,
"height": asset_record.height,
"bytes_size": asset_record.bytes_size,
"workflow_id": asset_record.workflow_id,
"prompt_id": asset_record.prompt_id,
"created_at": asset_record.created_at.isoformat(),
"expires_at": asset_record.expires_at.isoformat()
if asset_record.expires_at
else None,
}
# Enforce: only "thumb" mode for scoped version
if mode != "thumb":
return {
"error": f"Mode '{mode}' not supported in scoped version. Use 'thumb' or 'metadata'."
}
# Validate content type (only images supported for inline viewing)
supported_types = ("image/png", "image/jpeg", "image/jpg", "image/webp", "image/gif")
if asset_record.mime_type not in supported_types:
return {
"error": f"Asset type '{asset_record.mime_type}' not supported for inline viewing. "
f"Supported types: {', '.join(supported_types)}"
}
# Set conservative defaults
if max_dim is None:
max_dim = 512 # Hard cap for scoped version
if max_b64_chars is None:
max_b64_chars = 100_000 # 100KB base64 payload (conservative to prevent hangs)
# Process image for inline viewing
try:
# Fetch image bytes using orchestrator
image_bytes = asset_orchestrator.fetch_asset_bytes(asset_record)
# Encode with new function (accepts bytes directly)
cache_key = get_cache_key(asset_id, max_dim, 70) # Use quality=70 for cache key
encoded = encode_preview_for_mcp(
image_bytes,
max_dim=max_dim,
max_b64_chars=max_b64_chars,
quality=70,
cache_key=cache_key,
)
# Log telemetry
logger.info(
f"view_image success: asset_id={asset_id} "
f"src={asset_record.bytes_size}B src_dims={asset_record.width}x{asset_record.height} "
f"preview_dims={encoded.size_px[0]}x{encoded.size_px[1]} format=webp "
f"encoded={encoded.bytes_len}B b64_chars={encoded.b64_chars} "
f"response_est={estimate_response_chars(encoded.b64_chars)}chars"
)
# Use FastMCP.Image for inline display (not dict)
# FastMCP.Image takes raw bytes and format string
if FastMCPImage is not None:
return FastMCPImage(data=encoded.raw_bytes, format="webp")
else:
# Fallback if Image not available
return {
"content": [
{"type": "image", "data": encoded.b64_str, "mimeType": "image/webp"}
]
}
except ValueError as e:
# Image too large or processing failed - REFUSE-INLINE (non-lethal failure)
logger.warning(f"Refusing to inline image for {asset_id}: {e}")
return {
"content": [
{
"type": "text",
"text": (
f"Could not inline image (exceeds budget: {e}). "
f"Asset ID: {asset_id}. "
f"URL: {asset_url}. "
f"Source size: {asset_record.bytes_size} bytes. "
f"Source dimensions: {asset_record.width}x{asset_record.height}. "
f"Hint: Open URL locally or use metadata mode."
),
}
]
}
except ImportError as e:
return {
"error": f"Image processing not available: {e}. Install Pillow: pip install Pillow"
}
except Exception as e:
logger.exception(f"Failed to process asset {asset_id} for viewing")
return {"error": f"Failed to process asset: {str(e)}"}
