MCP/kicad-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor server to use type-safe lifespan context management

Browse Source 
Implement proper context management in the KiCad MCP server:

Add dedicated context.py with typed KiCadAppContext class
Convert tools to access context instead of parameters
Implement caching for thumbnails
Add proper cleanup of resources on shutdown
Improve error handling with cancellation support
This commit is contained in:
Lama 2025-03-20 09:58:19 -04:00
parent 5007d11579
commit 9a114bce7b
3 changed files with 379 additions and 230 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
kicad_mcp/context.py Normal file
View File

@ -0,0 +1,89 @@
"""
Lifespan context management for KiCad MCP Server.
"""
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import AsyncIterator, Optional, Dict, Any
from mcp.server.fastmcp import FastMCP
from kicad_mcp.utils.logger import Logger
from kicad_mcp.utils.python_path import setup_kicad_python_path
# Create logger for this module
logger = Logger()
@dataclass
class KiCadAppContext:
"""Type-safe context for KiCad MCP server."""
kicad_modules_available: bool
# Optional cache for expensive operations
cache: Dict[str, Any]
@asynccontextmanager
async def kicad_lifespan(server: FastMCP) -> AsyncIterator[KiCadAppContext]:
"""Manage KiCad MCP server lifecycle with type-safe context.
This function handles:
1. Initializing shared resources when the server starts
2. Providing a typed context object to all request handlers
3. Properly cleaning up resources when the server shuts down
Args:
server: The FastMCP server instance
Yields:
KiCadAppContext: A typed context object shared across all handlers
"""
logger.info("Starting KiCad MCP server initialization")
# Initialize resources on startup
logger.info("Setting up KiCad Python modules")
kicad_modules_available = setup_kicad_python_path()
logger.info(f"KiCad Python modules available: {kicad_modules_available}")
# Create in-memory cache for expensive operations
cache: Dict[str, Any] = {}
# Initialize any other resources that need cleanup later
created_temp_dirs = []
try:
# Import any KiCad modules that should be preloaded
if kicad_modules_available:
try:
logger.info("Preloading KiCad Python modules")
# Core PCB module used in multiple tools
import pcbnew
logger.info(f"Successfully preloaded pcbnew module: {getattr(pcbnew, 'GetBuildVersion', lambda: 'unknown')()}")
cache["pcbnew_version"] = getattr(pcbnew, "GetBuildVersion", lambda: "unknown")()
except ImportError as e:
logger.warning(f"Failed to preload some KiCad modules: {str(e)}")
# Yield the context to the server - server runs during this time
logger.info("KiCad MCP server initialization complete")
yield KiCadAppContext(
kicad_modules_available=kicad_modules_available,
cache=cache
)
finally:
# Clean up resources when server shuts down
logger.info("Shutting down KiCad MCP server")
# Clear the cache
if cache:
logger.debug(f"Clearing cache with {len(cache)} entries")
cache.clear()
# Clean up any temporary directories
import shutil
for temp_dir in created_temp_dirs:
try:
logger.debug(f"Removing temporary directory: {temp_dir}")
shutil.rmtree(temp_dir, ignore_errors=True)
except Exception as e:
logger.error(f"Error cleaning up temporary directory {temp_dir}: {str(e)}")
logger.info("KiCad MCP server shutdown complete")

8
kicad_mcp/server.py
View File

@ -38,8 +38,8 @@ def create_server() -> FastMCP:
logger.warning("KiCad Python modules not available - some features will be disabled")
# Initialize FastMCP server
mcp = FastMCP("KiCad")
logger.info("Created FastMCP server instance")
mcp = FastMCP("KiCad", lifespan=kicad_lifespan)
logger.info("Created FastMCP server instance with lifespan management")
# Register resources
logger.debug("Registering resources...")
@ -51,8 +51,8 @@ def create_server() -> FastMCP:
logger.debug("Registering tools...")
register_project_tools(mcp)
register_analysis_tools(mcp)
register_export_tools(mcp, kicad_modules_available)
register_drc_tools(mcp, kicad_modules_available)
register_export_tools(mcp)
register_drc_tools(mcp)
# Register prompts
logger.debug("Registering prompts...")

74
kicad_mcp/tools/export_tools.py
View File

@ -1,23 +1,26 @@
"""
Analysis and validation tools for KiCad projects.
Export tools for KiCad projects.
"""
import os
import tempfile
import subprocess
import shutil
import asyncio
from typing import Dict, Any, Optional
from mcp.server.fastmcp import FastMCP, Context, Image
from kicad_mcp.utils.file_utils import get_project_files
from kicad_mcp.utils.logger import Logger
from kicad_mcp.config import KICAD_APP_PATH, system
# Create logger for this module
logger = Logger()
def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False) -> None:
"""Register analysis and validation tools with the MCP server.
def register_export_tools(mcp: FastMCP) -> None:
"""Register export tools with the MCP server.
Args:
mcp: The FastMCP server instance
kicad_modules_available: Whether KiCad Python modules are available
"""
@mcp.tool()
@ -70,10 +73,16 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
Args:
project_path: Path to the KiCad project file (.kicad_pro)
ctx: Context for MCP communication
Returns:
Thumbnail image of the PCB or None if generation failed
"""
try:
# Access the context
app_context = ctx.request_context.lifespan_context
kicad_modules_available = app_context.kicad_modules_available
logger.info(f"Generating thumbnail for project: {project_path}")
if not os.path.exists(project_path):
@ -91,6 +100,12 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
pcb_file = files["pcb"]
logger.info(f"Found PCB file: {pcb_file}")
# Check cache
cache_key = f"thumbnail_{pcb_file}_{os.path.getmtime(pcb_file)}"
if hasattr(app_context, 'cache') and cache_key in app_context.cache:
logger.info(f"Using cached thumbnail for {pcb_file}")
return app_context.cache[cache_key]
await ctx.report_progress(10, 100)
ctx.info(f"Generating thumbnail for {os.path.basename(pcb_file)}")
@ -99,6 +114,9 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
try:
thumbnail = await generate_thumbnail_with_pcbnew(pcb_file, ctx)
if thumbnail:
# Cache the result if possible
if hasattr(app_context, 'cache'):
app_context.cache[cache_key] = thumbnail
return thumbnail
# If pcbnew method failed, log it but continue to try alternative method
@ -113,6 +131,9 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
try:
thumbnail = await generate_thumbnail_with_cli(pcb_file, ctx)
if thumbnail:
# Cache the result if possible
if hasattr(app_context, 'cache'):
app_context.cache[cache_key] = thumbnail
return thumbnail
except Exception as e:
logger.error(f"Error using CLI for thumbnail: {str(e)}", exc_info=True)
@ -122,9 +143,22 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
ctx.info("Could not generate thumbnail for PCB - all methods failed")
return None
except asyncio.CancelledError:
logger.info("Thumbnail generation cancelled")
raise # Re-raise to let MCP know the task was cancelled
except Exception as e:
logger.error(f"Unexpected error in thumbnail generation: {str(e)}")
ctx.info(f"Error: {str(e)}")
return None
@mcp.tool()
async def generate_project_thumbnail(project_path: str, ctx: Context) -> Optional[Image]:
"""Generate a thumbnail of a KiCad project's PCB layout."""
try:
# Access the context
app_context = ctx.request_context.lifespan_context
kicad_modules_available = app_context.kicad_modules_available
logger.info(f"Generating thumbnail for project: {project_path}")
if not os.path.exists(project_path):
@ -147,6 +181,12 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
ctx.info("KiCad Python modules are not available")
return None
# Check cache
cache_key = f"project_thumbnail_{pcb_file}_{os.path.getmtime(pcb_file)}"
if hasattr(app_context, 'cache') and cache_key in app_context.cache:
logger.info(f"Using cached project thumbnail for {pcb_file}")
return app_context.cache[cache_key]
try:
# Try to import pcbnew
import pcbnew
@ -221,7 +261,13 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
img_data = f.read()
logger.info(f"Successfully generated thumbnail, size: {len(img_data)} bytes")
return Image(data=img_data, format="png")
# Create and cache the image
thumbnail = Image(data=img_data, format="png")
if hasattr(app_context, 'cache'):
app_context.cache[cache_key] = thumbnail
return thumbnail
except ImportError as e:
logger.error(f"Failed to import pcbnew module: {str(e)}")
@ -232,6 +278,14 @@ def register_analysis_tools(mcp: FastMCP, kicad_modules_available: bool = False)
ctx.info(f"Error generating thumbnail: {str(e)}")
return None
except asyncio.CancelledError:
logger.info("Project thumbnail generation cancelled")
raise
except Exception as e:
logger.error(f"Unexpected error in project thumbnail generation: {str(e)}", exc_info=True)
ctx.info(f"Error: {str(e)}")
return None
# Helper functions for thumbnail generation
async def generate_thumbnail_with_pcbnew(pcb_file: str, ctx: Context) -> Optional[Image]:
"""Generate PCB thumbnail using the pcbnew Python module.
@ -344,8 +398,7 @@ async def generate_thumbnail_with_cli(pcb_file: str, ctx: Context) -> Optional[I
Returns:
Image object containing the PCB thumbnail or None if generation failed
"""
import subprocess
try:
logger.info("Attempting to generate thumbnail using command line tools")
await ctx.report_progress(20, 100)
@ -424,3 +477,10 @@ async def generate_thumbnail_with_cli(pcb_file: str, ctx: Context) -> Optional[I
except Exception as e:
logger.error(f"Error running CLI command: {str(e)}", exc_info=True)
return None
except asyncio.CancelledError:
logger.info("CLI thumbnail generation cancelled")
raise
except Exception as e:
logger.error(f"Unexpected error in CLI thumbnail generation: {str(e)}")
return None