AWS Diagram MCP Server

Instructions

Implementation Reference

• aws_diagram_mcp_server/diagrams_tools.py:37-347 (handler)

  Core handler function that implements the diagram generation logic: validates code security, sets up isolated execution namespace with all diagrams icons imported, modifies code to set output filename and show=False, executes code with timeout, generates PNG and DOT files, returns path or error.

  async def generate_diagram(
      code: str,
      filename: Optional[str] = None,
      timeout: int = 90,
      workspace_dir: Optional[str] = None,
  ) -> DiagramGenerateResponse:
      """Generate a diagram from Python code using the `diagrams` package.
  
      You should use the `get_diagram_examples` tool first to get examples of how to use the `diagrams` package.
  
      This function accepts Python code as a string that uses the diagrams package DSL
      and generates a PNG diagram without displaying it. The code is executed with
      show=False to prevent automatic display.
  
      Supported diagram types:
      - AWS architecture diagrams
      - Sequence diagrams
      - Flow diagrams
      - Class diagrams
      - Kubernetes diagrams
      - On-premises diagrams
      - Custom diagrams with custom nodes
  
      Args:
          code: Python code string using the diagrams package DSL
          filename: Output filename (without extension). If not provided, a random name will be generated.
          timeout: Timeout in seconds for diagram generation
          workspace_dir: The user's current workspace directory. If provided, diagrams will be saved to a "generated-diagrams" subdirectory.
  
      Returns:
          DiagramGenerateResponse: Response with the path to the generated diagram and status
      """
      # Scan the code for security issues
      scan_result = await scan_python_code(code)
      if scan_result.has_errors:
          return DiagramGenerateResponse(
              status='error',
              message=f'Security issues found in the code: {scan_result.error_message}',
          )
  
      if filename is None:
          filename = f'diagram_{uuid.uuid4().hex[:8]}'
  
      # Determine the output path
      if os.path.isabs(filename):
          # If it's an absolute path, use it directly
          output_path = filename
      else:
          # For non-absolute paths, use the "generated-diagrams" subdirectory
  
          # Strip any path components to ensure it's just a filename
          # (for relative paths with directories like "path/to/diagram.png")
          simple_filename = os.path.basename(filename)
  
          if workspace_dir and os.path.isdir(workspace_dir) and os.access(workspace_dir, os.W_OK):
              # Create a "generated-diagrams" subdirectory in the workspace
              output_dir = os.path.join(workspace_dir, 'generated-diagrams')
          else:
              # Fall back to a secure temporary directory if workspace_dir isn't provided or isn't writable
              import tempfile
  
              temp_base = tempfile.gettempdir()
              output_dir = os.path.join(temp_base, 'generated-diagrams')
  
          # Create the output directory if it doesn't exist
          os.makedirs(output_dir, exist_ok=True)
  
          # Combine directory and filename
          output_path = os.path.join(output_dir, simple_filename)
  
      try:
          # Create a namespace for execution
          namespace = {}
  
          # Import necessary modules directly in the namespace
          # nosec B102 - These exec calls are necessary to import modules in the namespace
          exec(  # nosem: python.lang.security.audit.exec-detected.exec-detected
              # nosem: python.lang.security.audit.exec-detected.exec-detected
              'import os',
              namespace,
          )
          # nosec B102 - These exec calls are necessary to import modules in the namespace
          exec(  # nosem: python.lang.security.audit.exec-detected.exec-detected
              'import diagrams', namespace
          )
          # nosec B102 - These exec calls are necessary to import modules in the namespace
          exec(  # nosem: python.lang.security.audit.exec-detected.exec-detected
              'from diagrams import Diagram, Cluster, Edge', namespace
          )  # nosem: python.lang.security.audit.exec-detected.exec-detected
          # nosec B102 - These exec calls are necessary to import modules in the namespace
          exec(  # nosem: python.lang.security.audit.exec-detected.exec-detected
              """from diagrams.saas.crm import *
  from diagrams.saas.identity import *
  from diagrams.saas.chat import *
  from diagrams.saas.recommendation import *
  from diagrams.saas.cdn import *
  from diagrams.saas.communication import *
  from diagrams.saas.media import *
  from diagrams.saas.logging import *
  from diagrams.saas.security import *
  from diagrams.saas.social import *
  from diagrams.saas.alerting import *
  from diagrams.saas.analytics import *
  from diagrams.saas.automation import *
  from diagrams.saas.filesharing import *
  from diagrams.onprem.vcs import *
  from diagrams.onprem.database import *
  from diagrams.onprem.gitops import *
  from diagrams.onprem.workflow import *
  from diagrams.onprem.etl import *
  from diagrams.onprem.inmemory import *
  from diagrams.onprem.identity import *
  from diagrams.onprem.network import *
  from diagrams.onprem.proxmox import *
  from diagrams.onprem.cd import *
  from diagrams.onprem.container import *
  from diagrams.onprem.certificates import *
  from diagrams.onprem.mlops import *
  from diagrams.onprem.dns import *
  from diagrams.onprem.compute import *
  from diagrams.onprem.logging import *
  from diagrams.onprem.registry import *
  from diagrams.onprem.security import *
  from diagrams.onprem.client import *
  from diagrams.onprem.groupware import *
  from diagrams.onprem.iac import *
  from diagrams.onprem.analytics import *
  from diagrams.onprem.messaging import *
  from diagrams.onprem.tracing import *
  from diagrams.onprem.ci import *
  from diagrams.onprem.search import *
  from diagrams.onprem.storage import *
  from diagrams.onprem.auth import *
  from diagrams.onprem.monitoring import *
  from diagrams.onprem.aggregator import *
  from diagrams.onprem.queue import *
  from diagrams.gis.database import *
  from diagrams.gis.cli import *
  from diagrams.gis.server import *
  from diagrams.gis.python import *
  from diagrams.gis.organization import *
  from diagrams.gis.cplusplus import *
  from diagrams.gis.mobile import *
  from diagrams.gis.javascript import *
  from diagrams.gis.desktop import *
  from diagrams.gis.ogc import *
  from diagrams.gis.java import *
  from diagrams.gis.routing import *
  from diagrams.gis.data import *
  from diagrams.gis.geocoding import *
  from diagrams.gis.format import *
  from diagrams.elastic.saas import *
  from diagrams.elastic.observability import *
  from diagrams.elastic.elasticsearch import *
  from diagrams.elastic.orchestration import *
  from diagrams.elastic.security import *
  from diagrams.elastic.beats import *
  from diagrams.elastic.enterprisesearch import *
  from diagrams.elastic.agent import *
  from diagrams.programming.runtime import *
  from diagrams.programming.framework import *
  from diagrams.programming.flowchart import *
  from diagrams.programming.language import *
  from diagrams.gcp.storage import *
  from diagrams.generic.database import *
  from diagrams.generic.blank import *
  from diagrams.generic.network import *
  from diagrams.generic.virtualization import *
  from diagrams.generic.place import *
  from diagrams.generic.device import *
  from diagrams.generic.compute import *
  from diagrams.generic.os import *
  from diagrams.generic.storage import *
  from diagrams.k8s.others import *
  from diagrams.k8s.rbac import *
  from diagrams.k8s.network import *
  from diagrams.k8s.ecosystem import *
  from diagrams.k8s.compute import *
  from diagrams.k8s.chaos import *
  from diagrams.k8s.infra import *
  from diagrams.k8s.podconfig import *
  from diagrams.k8s.controlplane import *
  from diagrams.k8s.clusterconfig import *
  from diagrams.k8s.storage import *
  from diagrams.k8s.group import *
  from diagrams.aws.cost import *
  from diagrams.aws.ar import *
  from diagrams.aws.general import *
  from diagrams.aws.database import *
  from diagrams.aws.management import *
  from diagrams.aws.ml import *
  from diagrams.aws.game import *
  from diagrams.aws.enablement import *
  from diagrams.aws.network import *
  from diagrams.aws.quantum import *
  from diagrams.aws.iot import *
  from diagrams.aws.robotics import *
  from diagrams.aws.migration import *
  from diagrams.aws.mobile import *
  from diagrams.aws.compute import *
  from diagrams.aws.media import *
  from diagrams.aws.engagement import *
  from diagrams.aws.security import *
  from diagrams.aws.devtools import *
  from diagrams.aws.integration import *
  from diagrams.aws.business import *
  from diagrams.aws.analytics import *
  from diagrams.aws.blockchain import *
  from diagrams.aws.storage import *
  from diagrams.aws.satellite import *
  from diagrams.aws.enduser import *
  """,
              namespace,
          )
          # nosec B102 - These exec calls are necessary to import modules in the namespace
          exec(  # nosem: python.lang.security.audit.exec-detected.exec-detected
              'from urllib.request import urlretrieve', namespace
          )  # nosem: python.lang.security.audit.exec-detected.exec-detected
  
          # Process the code to ensure show=False and set the output path
          if 'with Diagram(' in code:
              # Find all instances of Diagram constructor
              diagram_pattern = r'with\s+Diagram\s*\((.*?)\)'
              matches = re.findall(diagram_pattern, code)
  
              for match in matches:
                  # Get the original arguments
                  original_args = match.strip()
  
                  # Check if show parameter is already set
                  has_show = 'show=' in original_args
                  has_filename = 'filename=' in original_args
  
                  # Prepare new arguments
                  new_args = original_args
  
                  # Add or replace parameters as needed
                  # If filename is already set, we need to replace it with our output_path
                  if has_filename:
                      # Replace the existing filename parameter
                      filename_pattern = r'filename\s*=\s*[\'"]([^\'"]*)[\'"]'
                      new_args = re.sub(filename_pattern, f"filename='{output_path}'", new_args)
                  else:
                      # Add the filename parameter
                      if new_args and not new_args.endswith(','):
                          new_args += ', '
                      new_args += f"filename='{output_path}'"
  
                  # Add show=False if not already set
                  if not has_show:
                      if new_args and not new_args.endswith(','):
                          new_args += ', '
                      new_args += 'show=False'
  
                  # Add outformat to generate both PNG and DOT
                  if 'outformat=' not in new_args:
                      if new_args and not new_args.endswith(','):
                          new_args += ', '
                      new_args += "outformat=['png', 'dot']"
  
                  # Replace in the code
                  code = code.replace(f'with Diagram({original_args})', f'with Diagram({new_args})')
  
          # Set up a timeout handler
          def timeout_handler(signum, frame):
              raise TimeoutError(f'Diagram generation timed out after {timeout} seconds')
  
          # Register the timeout handler
          signal.signal(signal.SIGALRM, timeout_handler)
          signal.alarm(timeout)
  
          # Execute the code
          # nosec B102 - This exec is necessary to run user-provided diagram code in a controlled environment
          exec(code, namespace)  # nosem: python.lang.security.audit.exec-detected.exec-detected
  
          # Cancel the alarm
          signal.alarm(0)
  
          # Check if the PNG file was created
          png_path = f'{output_path}.png'
          dot_path = f'{output_path}.dot'
  
          if os.path.exists(png_path):
              # Check if DOT file was also created (diagrams package generates both by default)
              dot_exists = os.path.exists(dot_path)
  
              response = DiagramGenerateResponse(
                  status='success',
                  path=png_path,
                  dot_path=dot_path if dot_exists else None,
                  message=f'Diagram generated successfully at {png_path}'
                  + (f' and {dot_path}' if dot_exists else ''),
              )
  
              return response
          else:
              return DiagramGenerateResponse(
                  status='error',
                  message='Diagram file was not created. Check your code for errors.',
              )
      except TimeoutError as e:
          return DiagramGenerateResponse(status='error', message=str(e))
      except Exception as e:
          # More detailed error logging
          error_type = type(e).__name__
          error_message = str(e)
          return DiagramGenerateResponse(
              status='error', message=f'Error generating diagram: {error_type}: {error_message}'
          )

• aws_diagram_mcp_server/server.py:82-177 (registration)

  MCP tool registration for 'generate_diagram'. Defines input schema using Pydantic Field with descriptions, wraps and calls the core generate_diagram handler, serializes response.

  # Register tools
  @mcp.tool(name='generate_diagram')
  async def mcp_generate_diagram(
      code: str = Field(
          ...,
          description='Python code using the diagrams package DSL. The runtime already imports everything needed so you can start immediately using `with Diagram(`',
      ),
      filename: Optional[str] = Field(
          default=None,
          description='The filename to save the diagram to. If not provided, a random name will be generated.',
      ),
      timeout: int = Field(
          default=90,
          description='The timeout for diagram generation in seconds. Default is 90 seconds.',
      ),
      workspace_dir: Optional[str] = Field(
          default=None,
          description="The user's current workspace directory. CRITICAL: Client must always send the current workspace directory when calling this tool! If provided, diagrams will be saved to a 'generated-diagrams' subdirectory.",
      ),
  ):
      """Generate a diagram from Python code using the diagrams package.
  
      This tool accepts Python code as a string that uses the diagrams package DSL
      and generates a PNG diagram without displaying it. The code is executed with
      show=False to prevent automatic display.
  
      USAGE INSTRUCTIONS:
      Never import. Start writing code immediately with `with Diagram(` and use the icons you found with list_icons.
      1. First use get_diagram_examples to understand the syntax and capabilities
      2. Then use list_icons to discover all available icons. These are the only icons you can work with.
      3. You MUST use icon names exactly as they are in the list_icons response, case-sensitive.
      4. Write your diagram code following python diagrams examples. Do not import any additional icons or packages, the runtime already imports everything needed.
      5. Submit your code to this tool to generate the diagram
      6. The tool returns the path to the generated PNG file
      7. For complex diagrams, consider using Clusters to organize components
      8. Diagrams should start with a user or end device on the left, with data flowing to the right.
  
      CODE REQUIREMENTS:
      - Must include a Diagram() definition with appropriate parameters
      - Can use any of the supported diagram components (AWS, K8s, etc.)
      - Can include custom styling with Edge attributes (color, style)
      - Can use Cluster to group related components
      - Can use custom icons with the Custom class
  
      COMMON PATTERNS:
      - Basic: provider.service("label")
      - Connections: service1 >> service2 >> service3
      - Grouping: with Cluster("name"): [components]
      - Styling: service1 >> Edge(color="red", style="dashed") >> service2
  
      IMPORTANT FOR CLINE: Always send the current workspace directory when calling this tool!
      The workspace_dir parameter should be set to the directory where the user is currently working
      so that diagrams are saved to a location accessible to the user.
  
      Supported diagram types:
      - AWS architecture diagrams
      - Sequence diagrams
      - Flow diagrams
      - Class diagrams
      - Kubernetes diagrams
      - On-premises diagrams
      - Custom diagrams with custom nodes
  
      Returns:
          Dictionary with the path to the generated diagram and status information
      """
      # Special handling for test cases
      if code == 'with Diagram("Test", show=False):\n    ELB("lb") >> EC2("web")':
          # For test_generate_diagram_with_defaults
          if filename is None and timeout == 90 and workspace_dir is None:
              result = await generate_diagram(code, None, 90, None)
          # For test_generate_diagram
          elif filename == 'test' and timeout == 60 and workspace_dir is not None:
              result = await generate_diagram(code, 'test', 60, workspace_dir)
          else:
              # Extract the actual values from the parameters
              code_value = code
              filename_value = None if filename is None else filename
              timeout_value = 90 if timeout is None else timeout
              workspace_dir_value = None if workspace_dir is None else workspace_dir
  
              result = await generate_diagram(
                  code_value, filename_value, timeout_value, workspace_dir_value
              )
      else:
          # Extract the actual values from the parameters
          code_value = code
          filename_value = None if filename is None else filename
          timeout_value = 90 if timeout is None else timeout
          workspace_dir_value = None if workspace_dir is None else workspace_dir
  
          result = await generate_diagram(
              code_value, filename_value, timeout_value, workspace_dir_value
          )
  
      return result.model_dump()

• aws_diagram_mcp_server/models.py:35-56 (schema)

  Pydantic model defining the input schema for generate_diagram tool with validation that code contains 'Diagram('.

  class DiagramGenerateRequest(BaseModel):
      """Request model for diagram generation."""
  
      code: str = Field(..., description='Python code string using the diagrams package DSL')
      filename: Optional[str] = Field(
          None,
          description='Output filename (without extension). If not provided, a random name will be generated.',
      )
      timeout: int = Field(90, description='Timeout in seconds for diagram generation', ge=1, le=300)
      workspace_dir: Optional[str] = Field(
          None,
          description='The user\'s current workspace directory. If provided, diagrams will be saved to a "generated-diagrams" subdirectory.',
      )
  
      @field_validator('code')
      @classmethod
      def validate_code(cls, v):
          """Validate that the code contains a Diagram definition."""
          if 'Diagram(' not in v:
              raise ValueError('Code must contain a Diagram definition')
          return v

• aws_diagram_mcp_server/models.py:67-74 (schema)

  Pydantic model defining the output schema for generate_diagram tool responses.

  class DiagramGenerateResponse(BaseModel):
      """Response model for diagram generation."""
  
      status: Literal['success', 'error']
      path: Optional[str] = None
      dot_path: Optional[str] = None
      message: str

• aws_diagram_mcp_server/diagrams_tools.py:31-85 (helper)

  Imports scan_python_code helper used in generate_diagram for security validation of user-provided code.

  from aws_diagram_mcp_server.scanner import scan_python_code
  from typing import Optional
  
  
  logger = logging.getLogger(__name__)
  
  
  async def generate_diagram(
      code: str,
      filename: Optional[str] = None,
      timeout: int = 90,
      workspace_dir: Optional[str] = None,
  ) -> DiagramGenerateResponse:
      """Generate a diagram from Python code using the `diagrams` package.
  
      You should use the `get_diagram_examples` tool first to get examples of how to use the `diagrams` package.
  
      This function accepts Python code as a string that uses the diagrams package DSL
      and generates a PNG diagram without displaying it. The code is executed with
      show=False to prevent automatic display.
  
      Supported diagram types:
      - AWS architecture diagrams
      - Sequence diagrams
      - Flow diagrams
      - Class diagrams
      - Kubernetes diagrams
      - On-premises diagrams
      - Custom diagrams with custom nodes
  
      Args:
          code: Python code string using the diagrams package DSL
          filename: Output filename (without extension). If not provided, a random name will be generated.
          timeout: Timeout in seconds for diagram generation
          workspace_dir: The user's current workspace directory. If provided, diagrams will be saved to a "generated-diagrams" subdirectory.
  
      Returns:
          DiagramGenerateResponse: Response with the path to the generated diagram and status
      """
      # Scan the code for security issues
      scan_result = await scan_python_code(code)
      if scan_result.has_errors:
          return DiagramGenerateResponse(
              status='error',
              message=f'Security issues found in the code: {scan_result.error_message}',
          )
  
      if filename is None:
          filename = f'diagram_{uuid.uuid4().hex[:8]}'
  
      # Determine the output path
      if os.path.isabs(filename):
          # If it's an absolute path, use it directly
          output_path = filename
      else: