Cleanup code on infra-mcp

scripts/infra-mcp/pyproject.toml

@@ -1,8 +1,10 @@
 [project]
 name = "infra-mcp"
 version = "0.1.0"
-description = "MCP server for task runner integration"
+description = "MCP server for task runner integration, container operations, and infrastructure management"
+readme = "README.md"
 requires-python = ">=3.13"
+license = { text = "MIT" }
 dependencies = [
     "fastmcp>=2.10.0,<3",
     "requests>=2.32.4,<3",

scripts/infra-mcp/server.py

@@ -8,6 +8,7 @@
 import io
 import logging
 import os
+import signal
 import sys
 
 from fastmcp import FastMCP
@@ -20,11 +21,17 @@
 from tools.get_container_categories import ContainerCategoryFinder
 from tools.get_container_tags import ContainerTagFinder
 from tools.get_dashboard_groups import DashboardGroupFinder
+from utils.constants import DEFAULT_CONTAINER_ARCHITECTURE, DEFAULT_SAME_HASH_LIMIT, DEFAULT_TAG_LIMIT
 from utils.git import get_git_root
+from utils.models import ContainerTagFinderArgs
 from utils.security import validate_url_for_ssrf
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+# Configure logging from environment
+LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO")
+logging.basicConfig(
+    level=getattr(logging, LOG_LEVEL),
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
 logger = logging.getLogger("infra-mcp")
 logger.info("Starting Infra MCP server")
 
@@ -36,6 +43,7 @@
 
 @mcp.custom_route("/healthz", methods=["GET"])
 async def health_check(_request: Request) -> PlainTextResponse:
+    """Health check endpoint for monitoring server status."""
     return PlainTextResponse("OK")
 
 
@@ -96,7 +104,7 @@ def get_container_categories() -> list[str]:
 
 
 @mcp.tool(name="list-container-tags")
-def list_container_tags(image: str, limit: int = 10) -> list[str]:
+def list_container_tags(image: str, limit: int = DEFAULT_TAG_LIMIT) -> list[str]:
     """
     List recent tags for a container image.
 
@@ -109,26 +117,23 @@ def list_container_tags(image: str, limit: int = 10) -> list[str]:
     """
     tag_finder = ContainerTagFinder()
     try:
-        # Create a namespace to simulate command line args
-        class Args:
-            pass
-
-        args = Args()
-        args.image = image
-        args.architecture = "linux/amd64"
-        args.limit = limit
-        args.quiet = True
-        args.registry = None
+        args = ContainerTagFinderArgs(
+            image=image,
+            architecture=DEFAULT_CONTAINER_ARCHITECTURE,
+            limit=limit,
+            quiet=True,
+            registry=None,
+        )
 
         tags, _, _, _ = tag_finder.get_image_tags(args)
         return [tag["name"] for tag in tags[:limit]] if tags else []
     except Exception:
-        logger.exception("list-container-tags failed for image=%r", image)
+        logger.exception(f"list-container-tags failed for image={image!r}")
         return []
 
 
 @mcp.tool(name="list-same-hash-container-tags")
-def list_same_hash_container_tags(image: str, tag: str | None = None, limit: int = 100) -> list[str]:
+def list_same_hash_container_tags(image: str, tag: str | None = None, limit: int = DEFAULT_SAME_HASH_LIMIT) -> list[str]:
     """
     List tags that have the same hash as a specified container tag.
 
@@ -142,28 +147,25 @@ def list_same_hash_container_tags(image: str, tag: str | None = None, limit: int
     """
     tag_finder = ContainerTagFinder()
     try:
-        # Create a namespace to simulate command line args
-        class Args:
-            pass
-
-        args = Args()
-        args.image = image
-        args.tag = tag
-        args.architecture = "linux/amd64"
-        args.limit = limit
-        args.quiet = True
-        args.registry = None
+        args = ContainerTagFinderArgs(
+            image=image,
+            tag=tag,
+            architecture=DEFAULT_CONTAINER_ARCHITECTURE,
+            limit=limit,
+            quiet=True,
+            registry=None,
+        )
 
         # Get same hash tags but don't output to stdout
         same_hash_tags = tag_finder.list_same_hash_tags(args, suppress_output=True)
         return [tag["name"] for tag in same_hash_tags] if same_hash_tags else []
     except Exception:
-        logger.exception("list-same-hash-container-tags failed for image=%r tag=%r", image, tag)
+        logger.exception(f"list-same-hash-container-tags failed for image={image!r} tag={tag!r}")
         return []
 
 
 @mcp.tool(name="get-most-specific-container-tag")
-def get_most_specific_container_tag(image: str, tag: str | None = None, limit: int = 100) -> str:
+def get_most_specific_container_tag(image: str, tag: str | None = None, limit: int = DEFAULT_SAME_HASH_LIMIT) -> str:
     """
     Find the most specific container version tag from tags with the same hash.
 
@@ -177,35 +179,49 @@ def get_most_specific_container_tag(image: str, tag: str | None = None, limit: i
     """
     tag_finder = ContainerTagFinder()
     try:
-        # Create a namespace to simulate command line args
-        class Args:
-            pass
-
-        args = Args()
-        args.image = image
-        args.tag = tag
-        args.architecture = "linux/amd64"
-        args.limit = limit
-        args.quiet = True
-        args.registry = None
+        args = ContainerTagFinderArgs(
+            image=image,
+            tag=tag,
+            architecture=DEFAULT_CONTAINER_ARCHITECTURE,
+            limit=limit,
+            quiet=True,
+            registry=None,
+        )
 
         # Suppress any prints from the finder
         with contextlib.redirect_stdout(io.StringIO()):
             most_specific = tag_finder.get_most_specific_tag(args)
             same_hash = tag_finder.list_same_hash_tags(args, suppress_output=True)
+    except Exception:
+        logger.exception(f"get-most-specific-container-tag failed for image={image!r} tag={tag!r}")
+        return tag or "latest"
+    else:
         if most_specific:
             return most_specific["name"]
-        elif same_hash:
+        if same_hash:
             return same_hash[0]["name"]
-        else:
-            return tag or "latest"
-    except Exception:
-        logger.exception("get-most-specific-container-tag failed for image=%r tag=%r", image, tag)
         return tag or "latest"
 
 
 # --- Configure the FastMCP server ---
 
+
+def handle_shutdown(signum: int, _frame: object) -> None:
+    """Handle shutdown signals gracefully.
+
+    Args:
+        signum: Signal number received
+        _frame: Current stack frame (unused)
+    """
+    signal_name = signal.Signals(signum).name
+    logger.info(f"Received {signal_name}, shutting down gracefully...")
+    sys.exit(0)
+
+
+# Register signal handlers for graceful shutdown
+signal.signal(signal.SIGINT, handle_shutdown)
+signal.signal(signal.SIGTERM, handle_shutdown)
+
 try:
     # Get the repository root path
     repository_root_path = get_git_root()
@@ -227,12 +243,22 @@ class Args:
         add_container_operation_tools(mcp, repository_root_path)
     else:
         logger.info("Container operation tools disabled by environment variable")
-except Exception:  # noqa: BLE001
-    logger.exception("Failed to initialize server")
+except FileNotFoundError:
+    logger.exception("Repository not found")
+    sys.exit(1)
+except RuntimeError:
+    logger.exception("Git operation failed")
+    sys.exit(1)
+except ImportError:
+    logger.exception("Failed to import required module")
+    sys.exit(1)
+except Exception:
+    logger.exception("Unexpected error during server initialization")
     sys.exit(1)
 
 
-def main():
+def main() -> None:
+    """Start the MCP server."""
     # Start the server
     mcp.run()
     # Or start with parameters:

scripts/infra-mcp/tools/collections/container_tools.py

@@ -9,15 +9,23 @@
 import subprocess
 import sys
 from collections.abc import Callable
+from pathlib import Path
 
 from fastmcp import FastMCP
 from fastmcp.tools import Tool
 
+# Import constants from the shared constants module
+try:
+    from ...utils.constants import TASK_COMMAND_TIMEOUT
+except ImportError:
+    # Fallback for standalone execution
+    TASK_COMMAND_TIMEOUT = 600
+
 # Configure logging
 logger = logging.getLogger("infra-mcp")
 
 
-def get_container_operations():
+def get_container_operations() -> list[dict[str, str]]:
     """
     Get the list of valid container operations
 
@@ -34,7 +42,7 @@ def get_container_operations():
     ]
 
 
-def execute_container_operation(operation: str, service_name: str, repository_root_path: str) -> str:
+def execute_container_operation(operation: str, service_name: str, repository_root_path: str | Path) -> str:
     """
     Execute one operation on the specified service and return the output
 
@@ -55,25 +63,30 @@ def execute_container_operation(operation: str, service_name: str, repository_ro
     if not re.match(r"^[a-zA-Z0-9_/-]+$", service_name):
         return f"Invalid service name format: {service_name}"
 
+    repository_root_str = str(repository_root_path)
     cmd = [
         sys.executable,  # Use the current Python interpreter
-        os.path.join(repository_root_path, "scripts", "labctl.py"),
+        os.path.join(repository_root_str, "scripts", "labctl.py"),
         "service",
         operation,
         service_name,
     ]
 
     try:
         result = subprocess.run(  # noqa: S603
-            cmd, capture_output=True, text=True, check=True
+            cmd,
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=TASK_COMMAND_TIMEOUT,
         )
     except subprocess.CalledProcessError as e:
         return f"Error running operation: {e.stderr or str(e)}"
     else:
         return result.stdout or "(No output)"
 
 
-def create_operation_function(op: str, repository_root_path: str) -> Callable[[str], str]:
+def create_operation_function(op: str, repository_root_path: str | Path) -> Callable[[str], str]:
     """
     Create a function that executes a specific container service operation.
 
@@ -100,7 +113,7 @@ def operation_fn(service_name: str) -> str:
     return operation_fn
 
 
-def add_container_operation_tools(mcp_server: FastMCP, repository_root_path: str) -> None:
+def add_container_operation_tools(mcp_server: FastMCP, repository_root_path: str | Path) -> None:
     """
     Create and add tools to MCP server for container service operations.
 

scripts/infra-mcp/tools/collections/task_tools.py

@@ -8,15 +8,23 @@
 import shutil
 import subprocess
 from collections.abc import Callable
+from pathlib import Path
 
 from fastmcp import FastMCP
 from fastmcp.tools import Tool
 
+# Import constants from the shared constants module
+try:
+    from ...utils.constants import TASK_COMMAND_TIMEOUT
+except ImportError:
+    # Fallback for standalone execution
+    TASK_COMMAND_TIMEOUT = 600
+
 # Configure logging
 logger = logging.getLogger("infra-mcp")
 
 
-def get_task_list(repository_root_path: str) -> list[dict[str, str]]:
+def get_task_list(repository_root_path: str | Path) -> list[dict[str, str]]:
     """
     Get the list of available tasks by running 'task --list-all'.
 
@@ -33,10 +41,11 @@ def get_task_list(repository_root_path: str) -> list[dict[str, str]]:
         return []
     try:
         result = subprocess.run(  # noqa: S603
-            [task_bin, "--list-all", "--dir", repository_root_path],
+            [task_bin, "--list-all", "--dir", str(repository_root_path)],
             capture_output=True,
             text=True,
             check=True,
+            timeout=TASK_COMMAND_TIMEOUT,
         )
     except subprocess.CalledProcessError:
         logger.exception("Error getting task list")
@@ -58,7 +67,7 @@ def get_task_list(repository_root_path: str) -> list[dict[str, str]]:
         return tasks
 
 
-def execute_task(task_name: str, repository_root_path: str) -> str:
+def execute_task(task_name: str, repository_root_path: str | Path) -> str:
     """
     Execute a task command and return the output.
 
@@ -76,17 +85,18 @@ def execute_task(task_name: str, repository_root_path: str) -> str:
         return f"Error executing task {task_name}: 'task' binary not found"
     try:
         return subprocess.run(  # noqa: S603
-            [task_bin, task_name, "--dir", repository_root_path],
+            [task_bin, task_name, "--dir", str(repository_root_path)],
             capture_output=True,
             text=True,
             check=True,
+            timeout=TASK_COMMAND_TIMEOUT,
         ).stdout.strip()
     except subprocess.CalledProcessError as e:
         logger.exception(f"Error executing task {task_name}")
         return f"Error executing task {task_name}: {e.stderr}"
 
 
-def create_task_function(task_name: str, repository_root_path: str) -> Callable[[], str]:
+def create_task_function(task_name: str, repository_root_path: str | Path) -> Callable[[], str]:
     """
     Create a function that executes a specific task.
 
@@ -104,7 +114,7 @@ def task_fn() -> str:
     return task_fn
 
 
-def add_task_tools(mcp_server: FastMCP, repository_root_path: str) -> None:
+def add_task_tools(mcp_server: FastMCP, repository_root_path: str | Path) -> None:
     """
     Get list of tasks, then create and add tools to MCP server for each task.
 
@@ -116,7 +126,7 @@ def add_task_tools(mcp_server: FastMCP, repository_root_path: str) -> None:
 
     for task_info in tasks:
         task_name = task_info["name"]
-        tool_name = task_name.replace(":", "--")
+        tool_name = task_name.replace(":", "-")
         description = task_info["description"]
         task_fn = create_task_function(task_name, repository_root_path)