feat: implement Diagrams MCP server with Docker support and dynamic node discovery

mcp-server/Dockerfile

@@ -0,0 +1,25 @@
+# Use a slim Python image
+FROM python:3.11-slim
+
+# Install system dependencies (Graphviz is required for diagrams)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    graphviz \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy the current directory (which contains the diagrams library source and our server code)
+COPY . /app
+
+# Install the diagrams library from source (current dir) and other requirements
+# We install 'mcp' explicitly as it might not be in the local repo's requirements
+RUN pip install --no-cache-dir . mcp
+
+# Create an output directory for persistence if volume is mounted
+RUN mkdir -p /app/output
+
+# Run the MCP server using absolute path
+CMD ["python", "/app/src/server.py"]

mcp-server/GEMINI.md

@@ -0,0 +1,112 @@
+# Project: Diagrams MCP Server
+
+## Overview
+This project aims to implement a Model Context Protocol (MCP) server that exposes the capabilities of the [diagrams](https://diagrams.mingrammer.com/) Python library. The server will allow AI agents to dynamically discover available diagram nodes (AWS, Azure, Kubernetes, etc.) and generate architectural diagrams from Python code.
+
+## Architecture
+The solution will be containerized to ensure isolation and consistent dependencies (specifically Graphviz).
+
+-   **Runtime**: Python 3.9+
+-   **Container**: Docker (Debian-based to support Graphviz)
+-   **Communication**: Standard Input/Output (stdio) via the MCP protocol.
+-   **Libraries**:
+    -   `diagrams`: For generating diagrams.
+    -   `mcp`: Official Python SDK for the Model Context Protocol.
+    -   `graphviz`: System dependency required by the `diagrams` library.
+
+## File Structure
+```text
+.
+├── Dockerfile
+├── requirements.txt
+├── server.py
+└── src/
+    └── inspection.py  # Helper for dynamic node discovery
+```
+
+## Implementation Details
+
+### 1. Docker Environment (`Dockerfile`)
+The environment must include Graphviz, which is a system-level dependency required for rendering.
+
+*   **Base Image**: `python:3.11-slim`
+*   **System Dependencies**: `graphviz` (via `apt-get install -y graphviz`)
+*   **Python Dependencies**: `diagrams`, `mcp`
+
+### 2. MCP Server (`server.py`)
+The server will define three main tools. It should use the `mcp.server.fastmcp` or `mcp.server` standard library to define the server.
+
+#### Tool 1: `list_icons`
+**Purpose**: Dynamically discovers all available diagram nodes across all providers (AWS, Azure, GCP, SaaS, etc.) so the AI knows what classes are available to import.
+
+*   **Logic**:
+    1.  Recursively walk the `diagrams` package directory.
+    2.  Import modules dynamically.
+    3.  Inspect classes in each module.
+    4.  Filter classes that inherit from `diagrams.Node` but are not the base `Node` class itself.
+    5.  Organize into a hierarchy: `Provider -> Service -> Node`.
+    6.  **Optimization**: Cache this result at startup as it won't change.
+
+*   **Parameters**:
+    -   `provider_filter` (string, optional): If provided (e.g., "aws"), only return nodes for that provider.
+
+*   **Returns**: JSON structure:
+    ```json
+    {
+      "aws": {
+        "compute": ["EC2", "Lambda", ...],
+        "database": ["RDS", "DynamoDB", ...]
+      },
+      "k8s": { ... }
+    }
+    ```
+
+#### Tool 2: `generate_diagram`
+**Purpose**: Executes Python code to generate a diagram image.
+
+*   **Logic**:
+    1.  Accepts a string of Python code (DSL).
+    2.  **Security**: The code is executed via `exec()`. Since this runs inside a Docker container, it provides a layer of isolation.
+    3.  **Execution**:
+        -   Set up a temporary directory.
+        -   Change the working directory to this temp location.
+        -   Execute the code.
+        -   Find the generated output file (usually `.png`).
+    4.  **Result**: Return the path to the generated image or the base64 encoded content (depending on client capability, but path is preferred if sharing volume). *For this implementation, return the path inside the container and ensure the container mounts a shared volume if persistence is needed.*
+
+*   **Parameters**:
+    -   `code` (string, required): The Python code using `diagrams` DSL.
+    -   `filename` (string, optional): Desired output filename.
+
+*   **Returns**: 
+    -   `status`: "success" or "error"
+    -   `message`: Path to file or error message.
+
+#### Tool 3: `get_diagram_examples`
+**Purpose**: Provides example code snippets to help the AI understand the syntax.
+
+*   **Logic**: Return a dictionary of static examples for common patterns (Basic, Clustered, Cloud-specific).
+
+*   **Parameters**:
+    -   `provider` (string, optional): specific provider example (e.g., "aws").
+
+### 3. Dynamic Inspection Helper (`src/inspection.py`)
+This module is crucial for `list_icons`. It must robustly handle imports without crashing the server if a specific provider has missing optional dependencies.
+
+*   Use `pkgutil.walk_packages` to iterate over `diagrams`.
+*   Use `importlib.import_module` to load found modules.
+*   Use `inspect.getmembers` to find classes.
+*   Check `issubclass(obj, diagrams.Node)`.
+
+## Execution & Testing
+To run the server:
+```bash
+# Build
+docker build -t diagrams-mcp .
+
+# Run (connected to stdin/stdout for MCP)
+docker run -i --rm -v $(pwd)/output:/app/output diagrams-mcp
+```
+
+## Security Considerations
+*   **Arbitrary Code Execution**: The `generate_diagram` tool executes arbitrary Python code. This is by design but dangerous. The Docker container MUST be treated as untrusted and ephemeral. Do not mount sensitive host directories into the container.

mcp-server/README.md

@@ -0,0 +1,136 @@
+# Diagrams MCP Server
+
+This is a Model Context Protocol (MCP) server that exposes the capabilities of the [Diagrams](https://diagrams.mingrammer.com/) Python library. It allows AI agents to discover available diagram nodes (AWS, Azure, K8s, etc.) and generate architectural diagrams from Python code.
+
+## Architecture
+
+```mermaid
+graph TD
+    subgraph Host ["Host System"]
+        Client[MCP Client]
+        HostFS[Workspace / Output]
+    end
+
+    subgraph Container ["Docker Container"]
+        MCPServer["MCP Server (server.py)"]
+        NodeRegistry["(Node Registry)"]
+        
+        subgraph Logic ["Core Logic"]
+            Inspector["src/inspection.py"]
+            Executor["generate_diagram"]
+        end
+        
+        subgraph Libs ["Dependencies"]
+            DiagramsLib["diagrams package"]
+            Graphviz["Graphviz Binary"]
+        end
+    end
+
+    %% Startup Flow
+    MCPServer -- "Startup" --> Inspector
+    Inspector -- "Scans" --> DiagramsLib
+    Inspector -- "Populates" --> NodeRegistry
+
+    %% Tool Flows
+    Client -- "list_icons()" --> MCPServer
+    MCPServer -- "Query" --> NodeRegistry
+    
+    Client -- "generate_diagram(code)" --> MCPServer
+    MCPServer -- "Pass Code" --> Executor
+    Executor -- "exec()" --> DiagramsLib
+    DiagramsLib -- "Render" --> Graphviz
+    
+    %% Output
+    Graphviz -- "Generates PNG" --> Executor
+    Executor -- "Writes File (Volume Mount)" --> HostFS
+```
+
+## Features
+
+- **Dynamic Icon Discovery**: `list_icons` tool scans the `diagrams` library to find all available nodes (e.g., `EC2`, `Pod`, `BlobStorage`) organized by provider and service.
+- **Diagram Generation**: `generate_diagram` tool accepts Python code (DSL) and renders it into an image (PNG).
+- **Examples**: `get_diagram_examples` tool provides ready-to-use snippets for common patterns.
+- **Sandboxed Execution**: Runs inside a Docker container to ensure isolation and consistent dependencies (Graphviz).
+
+## Prerequisites
+
+- **Docker**: This server is designed to run as a Docker container to manage system dependencies like Graphviz.
+
+## Build
+
+Build the Docker image from the `mcp-server` directory:
+
+```bash
+cd mcp-server
+docker build -t diagrams-mcp:latest .
+```
+
+## Configuration
+
+To use this server with an MCP client (like Gemini CLI or Claude Desktop), add the following configuration.
+
+This configuration mounts the current project directory into the container, allowing the server to save the generated images directly to your workspace.
+
+```json
+{
+  "mcpServers": {
+    "diagrams": {
+      "command": "sh",
+      "args": [
+        "-c",
+        "export PROJECT_PATH=\"$\(PROJECT_PATH:-$(pwd)\")\"; docker run -i --rm -v \"$PROJECT_PATH:$PROJECT_PATH\" -w \"$PROJECT_PATH\" diagrams-mcp:latest"
+      ],
+      "env": {
+        "FASTMCP_LOG_LEVEL": "ERROR"
+      }
+    }
+  }
+}
+```
+
+### Explanation of the Command
+
+- **`docker run -i --rm`**: Runs the container interactively (for stdin/stdout communication) and removes it after exit.
+- **`-v "$PROJECT_PATH:$PROJECT_PATH"`**: Mounts the project root (where you invoke the agent) to the same path inside the container. This is crucial for the `generate_diagram` tool to write the output image file back to your host filesystem.
+- **`-w "$PROJECT_PATH"`**: Sets the working directory inside the container to match the host, ensuring relative paths work as expected.
+- **`diagrams-mcp:latest`**: The name of the image you built.
+
+## Tools
+
+### `list_icons`
+Lists available icons/nodes from the diagrams package.
+- **Inputs**: `provider_filter` (optional), `service_filter` (optional).
+- **Example**: List all AWS compute nodes.
+
+### `generate_diagram`
+Generates a diagram from Python code.
+- **Inputs**: `code` (Python DSL), `filename` (optional), `timeout` (default: 90s).
+- **Example Code**:
+  ```python
+  from diagrams import Diagram
+  from diagrams.aws.compute import EC2
+  
+  with Diagram("Simple", show=False):
+      EC2("web")
+  ```
+
+### `get_diagram_examples`
+Returns example code snippets.
+- **Inputs**: `diagram_type` (e.g., "aws", "k8s").
+
+## Development Structure
+
+```text
+mcp-server/
+├── Dockerfile          # Container definition (Python + Graphviz)
+├── requirements.txt    # Python deps
+├── README.md           # This file
+└── src/
+    ├── server.py       # Main MCP server entrypoint
+    └── inspection.py   # Helper for dynamic node discovery
+```
+
+## Autor
+
+- **Autor**: Carlos Barbero
+- **User**: carlosrgomes

mcp-server/requirements.txt

@@ -0,0 +1,2 @@
+mcp
+graphviz

mcp-server/src/inspection.py

@@ -0,0 +1,73 @@
+import pkgutil
+import importlib
+import inspect
+import sys
+from collections import defaultdict
+import diagrams
+from diagrams import Node
+
+def get_all_nodes():
+    """
+    Dynamically inspects the diagrams package and returns a dictionary of all available Nodes.
+    
+    Returns:
+        dict: A nested dictionary structure:
+              {
+                  "provider": {
+                      "service": ["NodeName1", "NodeName2", ...]
+                  }
+              }
+    """
+    # Initialize the structure
+    icons = defaultdict(lambda: defaultdict(list))
+    
+    # We also keep a flat map for the execution context: Name -> Class
+    # This handles potential name collisions by favoring the last seen or explicit logic if needed.
+    node_registry = {}
+
+    # Iterate through all subpackages in diagrams (e.g., aws, azure, k8s)
+    # We look at the path of the diagrams package
+    path = diagrams.__path__
+    prefix = diagrams.__name__ + "."
+
+    for _, provider_name, ispkg in pkgutil.iter_modules(path, prefix):
+        if not ispkg:
+            continue
+            
+        # e.g., provider_name = "diagrams.aws"
+        short_provider = provider_name.split(".")[-1]
+        
+        # Skip internal modules if any (base, etc are actually useful, but we focus on providers)
+        if short_provider in ['base', 'custom']: 
+             # 'custom' and 'base' might be treated differently, but for now we scan them
+             pass
+
+        try:
+            provider_module = importlib.import_module(provider_name)
+        except ImportError:
+            # Skip providers that might have missing system deps or issues
+            continue
+
+        # Now iterate modules within the provider (e.g., diagrams.aws.compute)
+        if hasattr(provider_module, "__path__"):
+            for _, service_name, _ in pkgutil.iter_modules(provider_module.__path__, provider_name + "."):
+                try:
+                    service_module = importlib.import_module(service_name)
+                    short_service = service_name.split(".")[-1]
+
+                    # Inspect classes in this service module
+                    for name, obj in inspect.getmembers(service_module, inspect.isclass):
+                        # Must inherit from Node
+                        if issubclass(obj, Node) and obj is not Node:
+                            # Verify it belongs to this module (to avoid re-export noise)
+                            # or at least is defined in the diagrams package
+                            if obj.__module__.startswith("diagrams"):
+                                icons[short_provider][short_service].append(name)
+                                node_registry[name] = obj
+                                
+                except ImportError:
+                    continue
+                except Exception:
+                    continue
+
+    return icons, node_registry

mcp-server/src/server.py

@@ -0,0 +1,176 @@
+import os
+import sys
+import tempfile
+import contextlib
+import base64
+from pathlib import Path
+from mcp.server.fastmcp import FastMCP
+from diagrams import Diagram, Cluster, Edge, Node
+
+# Import our helper
+from inspection import get_all_nodes
+
+# Initialize FastMCP
+mcp = FastMCP("diagrams-mcp")
+
+# Pre-load nodes for quick access and for the execution context
+print("Loading diagram nodes...", file=sys.stderr)
+ALL_ICONS, NODE_REGISTRY = get_all_nodes()
+print(f"Loaded {len(NODE_REGISTRY)} nodes.", file=sys.stderr)
+
+@mcp.tool()
+def list_icons(provider_filter: str = None, service_filter: str = None):
+    """
+    List available icons from the diagrams package, with optional filtering.
+    
+    Args:
+        provider_filter: Filter icons by provider name (e.g., "aws", "gcp", "k8s")
+        service_filter: Filter icons by service name (e.g., "compute", "database")
+    """
+    if not provider_filter:
+        # Return list of providers
+        return {"providers": list(ALL_ICONS.keys())}
+    
+    if provider_filter not in ALL_ICONS:
+        return {"error": f"Provider '{provider_filter}' not found. Available: {list(ALL_ICONS.keys())}"}
+    
+    provider_data = ALL_ICONS[provider_filter]
+    
+    if not service_filter:
+        # Return all services for this provider
+        return provider_data
+    
+    if service_filter not in provider_data:
+         return {"error": f"Service '{service_filter}' not found in '{provider_filter}'. Available: {list(provider_data.keys())}"}
+         
+    return {service_filter: provider_data[service_filter]}
+
+@mcp.tool()
+def get_diagram_examples(diagram_type: str = "all"):
+    """
+    Get example code for different types of diagrams.
+    
+    Args:
+        diagram_type: Type of diagram example to return (aws, k8s, flow, etc. or 'all')
+    """
+    examples = {
+        "aws": """
+from diagrams import Diagram
+from diagrams.aws.compute import EC2
+from diagrams.aws.database import RDS
+from diagrams.aws.network import ELB
+
+with Diagram("Web Service", show=False):
+    ELB("lb") >> EC2("web") >> RDS("userdb")
+""",
+        "k8s": """
+from diagrams import Diagram, Cluster
+from diagrams.k8s.compute import Pod
+from diagrams.k8s.network import Ingress, Service
+
+with Diagram("K8s Cluster", show=False):
+    ingress = Ingress("domain.com")
+    
+    with Cluster("App"):
+        svc = Service("svc")
+        pods = [Pod("pod1"), Pod("pod2")]
+        
+    ingress >> svc >> pods
+""",
+        "custom": """
+from diagrams import Diagram
+from diagrams.custom import Custom
+
+with Diagram("Custom", show=False):
+    # Ensure you have the icon file locally if using Custom
+    Custom("Label", "./my-icon.png")
+"""
+    }
+    
+    if diagram_type == "all":
+        return examples
+    
+    return {diagram_type: examples.get(diagram_type, "No example found for this type.")}
+
+@mcp.tool()
+def generate_diagram(code: str, filename: str = None, timeout: int = 90):
+    """
+    Generate a diagram from Python code using the diagrams package.
+    
+    Args:
+        code: Python code using the diagrams package DSL.
+        filename: Optional filename to save the diagram to.
+        timeout: Execution timeout in seconds.
+    """
+    
+    # Create a temporary directory for execution
+    with tempfile.TemporaryDirectory() as temp_dir:
+        original_cwd = os.getcwd()
+        os.chdir(temp_dir)
+        
+        try:
+            # Prepare the execution context
+            # We inject Diagram, Cluster, Edge, and ALL discovered nodes (EC2, Pod, etc.)
+            # This allows the user to write code without heavy imports if they choose,
+            # though explicit imports are still better for clarity.
+            exec_globals = {
+                "Diagram": Diagram,
+                "Cluster": Cluster,
+                "Edge": Edge,
+                "Node": Node,
+                **NODE_REGISTRY
+            }
+            
+            # Execute the code
+            # We wrap it in a try/except block within the exec to catch runtime errors
+            try:
+                exec(code, exec_globals)
+            except Exception as e:
+                return {"status": "error", "message": f"Runtime error: {str(e)}"}
+                
+            # Find the generated file
+            # Diagrams generates files based on the name passed to Diagram() class
+            # We look for any .png file created in the temp dir
+            generated_files = list(Path(".").glob("*.png"))
+            
+            if not generated_files:
+                 return {"status": "error", "message": "No diagram image was generated. Did you call with Diagram(..., show=False)?"}
+            
+            # Use the most recently modified file or the first one
+            generated_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
+            output_file = generated_files[0]
+            
+            # If a filename was requested, we might want to rename it?
+            # For now, we return the path. 
+            # In a real MCP setup, we might copy this to a mounted volume.
+            
+            # Copy the generated file back to the original working directory
+            # This ensures that if the user mounted their project to the working directory,
+            # the file appears in their project.
+            import shutil
+            
+            target_dir = Path(original_cwd)
+            target_filename = filename if filename else output_file.name
+            target_path = target_dir / target_filename
+            
+            # Ensure extension
+            if not target_path.suffix:
+                target_path = target_path.with_suffix(".png")
+                
+            shutil.copy2(output_file, target_path)
+            final_path = str(target_path)
+                
+            return {
+                "status": "success", 
+                "path": final_path,
+                "filename": target_path.name
+            }
+
+        except Exception as e:
+            return {"status": "error", "message": f"System error: {str(e)}"}
+            
+        finally:
+            os.chdir(original_cwd)
+
+if __name__ == "__main__":
+    mcp.run()