Diagrams MCP Server

Instructions

Implementation Reference

• src/diagrams_mcp/tools/core.py:125-142 (registration)

  Registration of the 'create_diagram' tool using the @mcp.tool decorator, including name, detailed description, input schema via Annotated parameters, and execution hints.

  @mcp.tool(
      name="create_diagram",
      description="""Generate infrastructure diagrams with 15+ providers (AWS, Azure, GCP, K8s, etc.).
  
  Examples:
  AWS: nodes=[{"id":"r53","provider":"aws","category":"network","type":"Route53",...}]
  K8s: nodes=[{"id":"ing","provider":"k8s","category":"network","type":"Ingress",...}]
  Clusters: clusters=[{"name":"VPC","node_ids":["elb","ec2"],"graph_attr":{"bgcolor":"#E5F5FD"}}]
  
  ⚠️ CRITICAL: Node types must exist in diagrams library or diagram fails silently (no arrows).
  ALWAYS verify first: list_available_nodes(provider="aws", category="compute")
  For brands (Stripe, Vercel), use create_diagram_with_custom_icons instead.""",
      annotations={
          "readOnlyHint": False,
          "destructiveHint": False,
          "idempotentHint": True,
      },
  )

• src/diagrams_mcp/tools/core.py:143-361 (handler)

  Core handler function for 'create_diagram' tool. Validates inputs, dynamically imports node classes from diagrams library, constructs Diagram with nodes, connections, clusters, generates multi-format outputs, computes metadata, optionally embeds base64 images, and formats response.

  async def create_diagram(
      name: Annotated[str, Field(description="Diagram title")],
      nodes: Annotated[List[NodeDef], Field(description="List of nodes to include")],
      connections: Annotated[
          List[ConnectionDef], Field(description="List of connections between nodes")
      ],
      clusters: Annotated[
          Optional[List[ClusterDef]],
          Field(description="Optional clusters for grouping nodes"),
      ] = None,
      direction: Annotated[
          Literal["LR", "RL", "TB", "BT"],
          Field(description="Diagram direction (LR=left-right, TB=top-bottom)"),
      ] = "LR",
      curvestyle: Annotated[
          Literal["ortho", "curved"], Field(description="Edge curve style")
      ] = "ortho",
      output_format: Annotated[
          str | List[str],
          Field(description="Output format(s): png, pdf, jpg, dot"),
      ] = "png",
      output_dir: Annotated[
          Optional[str],
          Field(
              description="Output directory (default: current directory). Auto-created if missing."
          ),
      ] = None,
      graph_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Graphviz graph attributes")
      ] = None,
      node_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Default node attributes")
      ] = None,
      edge_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Default edge attributes")
      ] = None,
      autolabel: Annotated[bool, Field(description="Auto-prefix nodes with class names")] = False,
      return_base64: Annotated[bool, Field(description="Include base64-encoded images")] = False,
  ) -> str:
      """Generate infrastructure/architecture diagram."""
      start_time = time.time()
  
      try:
          # Validate all node references
          node_ids = {node.id for node in nodes}
          for node in nodes:
              validate_node_reference(node.provider, node.category, node.type)
  
          # Validate connections reference existing nodes
          for conn in connections:
              if conn.from_node not in node_ids:
                  raise ValueError(
                      f"Connection references unknown node '{conn.from_node}'. "
                      f"Available nodes: {', '.join(sorted(node_ids))}"
                  )
  
              # Handle single target or list of targets
              targets = [conn.to_node] if isinstance(conn.to_node, str) else conn.to_node
              for target in targets:
                  if target not in node_ids:
                      raise ValueError(
                          f"Connection references unknown node '{target}'. "
                          f"Available nodes: {', '.join(sorted(node_ids))}"
                      )
  
          # Validate cluster node references
          if clusters:
              for cluster in clusters:
                  for node_id in cluster.node_ids:
                      if node_id not in node_ids:
                          raise ValueError(
                              f"Cluster '{cluster.name}' references unknown node '{node_id}'. "
                              f"Available nodes: {', '.join(sorted(node_ids))}"
                          )
  
          # Prepare output formats
          formats = [output_format] if isinstance(output_format, str) else output_format
  
          # Reject SVG - it's buggy and unsupported
          if any("svg" in fmt.lower() for fmt in formats):
              raise ValueError("SVG output is not supported. Use png, pdf, jpg, or dot instead.")
  
          # Change to output directory if specified
          original_dir = os.getcwd()
          if output_dir:
              os.makedirs(output_dir, exist_ok=True)
              os.chdir(output_dir)
  
          try:
              # Create diagram with Diagram context manager
              # Hide diagram title by setting label to empty string
              default_graph_attr = {"label": ""}
              merged_graph_attr = {**default_graph_attr, **(graph_attr or {})}
  
              with Diagram(
                  name=name,
                  show=False,  # Never auto-open in MCP server
                  direction=direction,
                  curvestyle=curvestyle,
                  outformat=formats,
                  autolabel=autolabel,
                  graph_attr=merged_graph_attr,
                  node_attr=node_attr or {},
                  edge_attr=edge_attr or {},
              ) as _:
                  # Create all nodes
                  node_objects = {}
                  for node in nodes:
                      # Import the node class
                      NodeClass = import_node_class(node.provider, node.category, node.type)
  
                      # Create node instance
                      node_obj = NodeClass(node.label)
                      node_objects[node.id] = node_obj
  
                  # Handle clusters
                  cluster_objects = {}
                  if clusters:
                      # Build cluster hierarchy (simple approach: process in order)
                      for cluster in clusters:
                          # Create cluster context
                          cluster_ctx = Cluster(
                              cluster.name,
                              graph_attr=cluster.graph_attr or {},
                          )
  
                          # Store for later reference
                          cluster_objects[cluster.name] = {
                              "context": cluster_ctx,
                              "node_ids": cluster.node_ids,
                              "parent": cluster.parent_cluster,
                          }
  
                      # Apply clusters (nodes will be added inside cluster contexts)
                      # For now, we'll use a simplified approach without nesting
                      # Full nesting support would require more complex logic
  
                  # Create all connections
                  edge_count = 0
                  for conn in connections:
                      from_obj = node_objects[conn.from_node]
  
                      # Handle single target or list of targets
                      targets = [conn.to_node] if isinstance(conn.to_node, str) else conn.to_node
  
                      for target in targets:
                          to_obj = node_objects[target]
  
                          # Create edge with optional styling
                          if conn.label or conn.color or conn.style:
                              edge = Edge(
                                  label=conn.label or "",
                                  color=conn.color or "black",
                                  style=conn.style or "solid",
                              )
  
                              if conn.direction == "forward":
                                  _ = from_obj >> edge >> to_obj
                              elif conn.direction == "reverse":
                                  _ = from_obj << edge << to_obj
                              else:  # bidirectional
                                  _ = from_obj - edge - to_obj
                          else:
                              # Simple connection without styling
                              if conn.direction == "forward":
                                  _ = from_obj >> to_obj
                              elif conn.direction == "reverse":
                                  _ = from_obj << to_obj
                              else:  # bidirectional
                                  _ = from_obj - to_obj
  
                          edge_count += 1
  
              # Get generated file paths
              # diagrams library generates files with snake_case names
              diagram_filename = name.replace(" ", "_").replace("-", "_").lower()
              file_paths = []
              for fmt in formats:
                  file_path = f"{diagram_filename}.{fmt}"
                  if output_dir:
                      file_path = os.path.join(output_dir, file_path)
                  file_paths.append(os.path.abspath(file_path))
  
              # Build metadata
              generation_time_ms = (time.time() - start_time) * 1000
              metadata = build_diagram_metadata(
                  file_paths,
                  node_count=len(nodes),
                  edge_count=edge_count,
                  cluster_count=len(clusters) if clusters else 0,
                  generation_time_ms=generation_time_ms,
              )
  
              # Optionally encode images as base64
              base64_images = None
              if return_base64:
                  base64_images = {}
                  for path in file_paths:
                      ext = Path(path).suffix[1:]
                      if ext != "dot":  # Don't encode dot files
                          try:
                              base64_images[ext] = encode_file_base64(path)
                          except Exception:
                              # Non-fatal error, just skip
                              pass
  
              return format_diagram_result(file_paths, metadata, base64_images)
  
          finally:
              # Restore original directory
              if output_dir:
                  os.chdir(original_dir)
  
      except Exception as e:
          return format_error(
              f"Failed to generate diagram: {str(e)}",
              suggestion="Check node types with list_available_nodes tool",
          )

• src/diagrams_mcp/tools/core.py:144-181 (schema)

  Input schema and type definitions for the create_diagram tool parameters, using Pydantic Field for descriptions and validation.

      name: Annotated[str, Field(description="Diagram title")],
      nodes: Annotated[List[NodeDef], Field(description="List of nodes to include")],
      connections: Annotated[
          List[ConnectionDef], Field(description="List of connections between nodes")
      ],
      clusters: Annotated[
          Optional[List[ClusterDef]],
          Field(description="Optional clusters for grouping nodes"),
      ] = None,
      direction: Annotated[
          Literal["LR", "RL", "TB", "BT"],
          Field(description="Diagram direction (LR=left-right, TB=top-bottom)"),
      ] = "LR",
      curvestyle: Annotated[
          Literal["ortho", "curved"], Field(description="Edge curve style")
      ] = "ortho",
      output_format: Annotated[
          str | List[str],
          Field(description="Output format(s): png, pdf, jpg, dot"),
      ] = "png",
      output_dir: Annotated[
          Optional[str],
          Field(
              description="Output directory (default: current directory). Auto-created if missing."
          ),
      ] = None,
      graph_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Graphviz graph attributes")
      ] = None,
      node_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Default node attributes")
      ] = None,
      edge_attr: Annotated[
          Optional[Dict[str, Any]], Field(description="Default edge attributes")
      ] = None,
      autolabel: Annotated[bool, Field(description="Auto-prefix nodes with class names")] = False,
      return_base64: Annotated[bool, Field(description="Include base64-encoded images")] = False,
  ) -> str:

• src/diagrams_mcp/tools/core.py:41-70 (helper)

  Helper function to dynamically import node classes from the 'diagrams' library based on provider, category, and type, crucial for runtime node instantiation.

  def import_node_class(provider: str, category: str, node_type: str):
      """Dynamically import a node class from the diagrams library.
  
      Args:
          provider: Provider name (e.g., "aws")
          category: Category name (e.g., "compute")
          node_type: Node type (e.g., "EC2")
  
      Returns:
          Node class
  
      Raises:
          ImportError: If the node class cannot be imported
      """
      try:
          # Construct import path: diagrams.aws.compute
          module_path = f"diagrams.{provider}.{category}"
          module = __import__(module_path, fromlist=[node_type])
  
          # Get the class
          node_class = getattr(module, node_type)
          return node_class
  
      except (ImportError, AttributeError) as e:
          raise ImportError(
              f"Failed to import node '{node_type}' from {provider}.{category}. "
              f"Error: {str(e)}. "
              f"Check that the provider, category, and type are correct."
          )