AWS Diagram MCP Server

Instructions

Implementation Reference

• aws_diagram_mcp_server/server.py:83-177 (handler)

  MCP handler function for the 'generate_diagram' tool. Defines input schema via Pydantic Fields, handles special test cases, calls core generate_diagram function, and returns response.
  @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/diagrams_tools.py:38-347 (handler)

  Core implementation that performs security scan, sets up safe execution namespace with diagrams imports, modifies code to ensure safe rendering, executes user diagram code with timeout, and generates PNG diagram file.
  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/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/models.py:35-56 (schema)

  Pydantic model for input validation/schema of diagram generation (note: handler uses inline Fields, but this defines equivalent schema).
  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/server.py:83-83 (registration)

  MCP tool registration decorator binding the name 'generate_diagram' to the mcp_generate_diagram handler.
  @mcp.tool(name='generate_diagram')
• aws_diagram_mcp_server/scanner.py:166-210 (helper)

  Security scanner used by generate_diagram to validate user-submitted code before execution.
  async def scan_python_code(code: str) -> CodeScanResult: """Use ast and bandit to scan the python code for security issues.""" # Get code metrics metrics = await count_code_metrics(code) # Check syntax syntax_valid, syntax_error = await validate_syntax(code) if not syntax_valid: return CodeScanResult( has_errors=True, syntax_valid=False, error_message=syntax_error, metrics=metrics ) # Check security security_issues = await check_security(code) # Check for dangerous functions explicitly dangerous_functions = check_dangerous_functions(code) if dangerous_functions: for func in dangerous_functions: security_issues.append( SecurityIssue( severity='HIGH', confidence='HIGH', line=func['line'], issue_text=f"Dangerous function '{func['function']}' detected", issue_type='DangerousFunctionDetection', ) ) # Determine if there are errors has_errors = bool(security_issues) # Generate error message if needed error_message = None if has_errors: messages = [f'{issue.issue_type}: {issue.issue_text}' for issue in security_issues] error_message = '\n'.join(messages) if messages else None return CodeScanResult( has_errors=has_errors, syntax_valid=True, security_issues=security_issues, error_message=error_message, metrics=metrics, )