hanzo-mcp 0.1.21__py3-none-any.whl

hanzo_mcp/tools/common/context.py

@@ -0,0 +1,444 @@
+"""Enhanced Context for Hanzo Dev MCP tools.
+
+This module provides an enhanced Context class that wraps the MCP Context
+and adds additional functionality specific to Claude Code tools.
+"""
+
+import json
+import os
+import time
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any, ClassVar, final
+
+from mcp.server.fastmcp import Context as MCPContext
+from mcp.server.lowlevel.helper_types import ReadResourceContents
+
+
+@final
+class ToolContext:
+    """Enhanced context for Hanzo Dev MCP tools.
+
+    This class wraps the MCP Context and adds additional functionality
+    for tracking tool execution, progress reporting, and resource access.
+    """
+
+    # Track all active contexts for debugging
+    _active_contexts: ClassVar[set["ToolContext"]] = set()
+
+    def __init__(self, mcp_context: MCPContext) -> None:
+        """Initialize the tool context.
+
+        Args:
+            mcp_context: The underlying MCP Context
+        """
+        self._mcp_context: MCPContext = mcp_context
+        self._tool_name: str | None = None
+        self._execution_id: str | None = None
+
+        # Add to active contexts
+        ToolContext._active_contexts.add(self)
+
+    def __del__(self) -> None:
+        """Clean up when the context is destroyed."""
+        # Remove from active contexts
+        ToolContext._active_contexts.discard(self)
+
+    @property
+    def mcp_context(self) -> MCPContext:
+        """Get the underlying MCP Context.
+
+        Returns:
+            The MCP Context
+        """
+        return self._mcp_context
+
+    @property
+    def request_id(self) -> str:
+        """Get the request ID from the MCP context.
+
+        Returns:
+            The request ID
+        """
+        return self._mcp_context.request_id
+
+    @property
+    def client_id(self) -> str | None:
+        """Get the client ID from the MCP context.
+
+        Returns:
+            The client ID
+        """
+        return self._mcp_context.client_id
+
+    def set_tool_info(self, tool_name: str, execution_id: str | None = None) -> None:
+        """Set information about the currently executing tool.
+
+        Args:
+            tool_name: The name of the tool being executed
+            execution_id: Optional unique execution ID
+        """
+        self._tool_name = tool_name
+        self._execution_id = execution_id
+
+    async def info(self, message: str) -> None:
+        """Log an informational message.
+
+        Args:
+            message: The message to log
+        """
+        await self._mcp_context.info(self._format_message(message))
+
+    async def debug(self, message: str) -> None:
+        """Log a debug message.
+
+        Args:
+            message: The message to log
+        """
+        await self._mcp_context.debug(self._format_message(message))
+
+    async def warning(self, message: str) -> None:
+        """Log a warning message.
+
+        Args:
+            message: The message to log
+        """
+        await self._mcp_context.warning(self._format_message(message))
+
+    async def error(self, message: str) -> None:
+        """Log an error message.
+
+        Args:
+            message: The message to log
+        """
+        await self._mcp_context.error(self._format_message(message))
+
+    def _format_message(self, message: str) -> str:
+        """Format a message with tool information if available.
+
+        Args:
+            message: The original message
+
+        Returns:
+            The formatted message
+        """
+        if self._tool_name:
+            if self._execution_id:
+                return f"[{self._tool_name}:{self._execution_id}] {message}"
+            return f"[{self._tool_name}] {message}"
+        return message
+
+    async def report_progress(self, current: int, total: int) -> None:
+        """Report progress to the client.
+
+        Args:
+            current: Current progress value
+            total: Total progress value
+        """
+        await self._mcp_context.report_progress(current, total)
+
+    async def read_resource(self, uri: str) -> Iterable[ReadResourceContents]:
+        """Read a resource via the MCP protocol.
+
+        Args:
+            uri: The resource URI
+
+        Returns:
+            A tuple of (content, mime_type)
+        """
+        return await self._mcp_context.read_resource(uri)
+
+
+# Factory function to create a ToolContext from an MCP Context
+def create_tool_context(mcp_context: MCPContext) -> ToolContext:
+    """Create a ToolContext from an MCP Context.
+
+    Args:
+        mcp_context: The MCP Context
+
+    Returns:
+        A new ToolContext
+    """
+    return ToolContext(mcp_context)
+
+
+@final
+class DocumentContext:
+    """Manages document context and codebase understanding."""
+
+    def __init__(self) -> None:
+        """Initialize the document context."""
+        self.documents: dict[str, str] = {}
+        self.document_metadata: dict[str, dict[str, Any]] = {}
+        self.modified_times: dict[str, float] = {}
+        self.allowed_paths: set[Path] = set()
+
+    def add_allowed_path(self, path: str) -> None:
+        """Add a path to the allowed paths.
+
+        Args:
+            path: The path to allow
+        """
+        resolved_path: Path = Path(path).resolve()
+        self.allowed_paths.add(resolved_path)
+
+    def is_path_allowed(self, path: str) -> bool:
+        """Check if a path is allowed.
+
+        Args:
+            path: The path to check
+
+        Returns:
+            True if the path is allowed, False otherwise
+        """
+        resolved_path: Path = Path(path).resolve()
+
+        # Check if the path is within any allowed path
+        for allowed_path in self.allowed_paths:
+            try:
+                _ = resolved_path.relative_to(allowed_path)
+                return True
+            except ValueError:
+                continue
+
+        return False
+
+    def add_document(
+        self, path: str, content: str, metadata: dict[str, Any] | None = None
+    ) -> None:
+        """Add a document to the context.
+
+        Args:
+            path: The path of the document
+            content: The content of the document
+            metadata: Optional metadata about the document
+        """
+        self.documents[path] = content
+        self.modified_times[path] = time.time()
+
+        if metadata:
+            self.document_metadata[path] = metadata
+        else:
+            # Try to infer metadata
+            self.document_metadata[path] = self._infer_metadata(path, content)
+
+    def get_document(self, path: str) -> str | None:
+        """Get a document from the context.
+
+        Args:
+            path: The path of the document
+
+        Returns:
+            The document content, or None if not found
+        """
+        return self.documents.get(path)
+
+    def get_document_metadata(self, path: str) -> dict[str, Any] | None:
+        """Get document metadata.
+
+        Args:
+            path: The path of the document
+
+        Returns:
+            The document metadata, or None if not found
+        """
+        return self.document_metadata.get(path)
+
+    def update_document(self, path: str, content: str) -> None:
+        """Update a document in the context.
+
+        Args:
+            path: The path of the document
+            content: The new content of the document
+        """
+        self.documents[path] = content
+        self.modified_times[path] = time.time()
+
+        # Update metadata
+        self.document_metadata[path] = self._infer_metadata(path, content)
+
+    def remove_document(self, path: str) -> None:
+        """Remove a document from the context.
+
+        Args:
+            path: The path of the document
+        """
+        if path in self.documents:
+            del self.documents[path]
+
+        if path in self.document_metadata:
+            del self.document_metadata[path]
+
+        if path in self.modified_times:
+            del self.modified_times[path]
+
+    def _infer_metadata(self, path: str, content: str) -> dict[str, Any]:
+        """Infer metadata about a document.
+
+        Args:
+            path: The path of the document
+            content: The content of the document
+
+        Returns:
+            Inferred metadata
+        """
+        extension: str = Path(path).suffix.lower()
+
+        metadata: dict[str, Any] = {
+            "extension": extension,
+            "size": len(content),
+            "line_count": content.count("\n") + 1,
+        }
+
+        # Infer language based on extension
+        language_map: dict[str, list[str]] = {
+            "python": [".py"],
+            "javascript": [".js", ".jsx"],
+            "typescript": [".ts", ".tsx"],
+            "java": [".java"],
+            "c++": [".c", ".cpp", ".h", ".hpp"],
+            "go": [".go"],
+            "rust": [".rs"],
+            "ruby": [".rb"],
+            "php": [".php"],
+            "html": [".html", ".htm"],
+            "css": [".css"],
+            "markdown": [".md"],
+            "json": [".json"],
+            "yaml": [".yaml", ".yml"],
+            "xml": [".xml"],
+            "sql": [".sql"],
+            "shell": [".sh", ".bash"],
+        }
+
+        # Find matching language
+        for language, extensions in language_map.items():
+            if extension in extensions:
+                metadata["language"] = language
+                break
+        else:
+            metadata["language"] = "text"
+
+        return metadata
+
+    def load_directory(
+        self,
+        directory: str,
+        recursive: bool = True,
+        exclude_patterns: list[str] | None = None,
+    ) -> None:
+        """Load all files in a directory into the context.
+
+        Args:
+            directory: The directory to load
+            recursive: Whether to load subdirectories
+            exclude_patterns: Patterns to exclude
+        """
+        if not self.is_path_allowed(directory):
+            raise ValueError(f"Directory not allowed: {directory}")
+
+        dir_path: Path = Path(directory)
+
+        if not dir_path.exists() or not dir_path.is_dir():
+            raise ValueError(f"Not a valid directory: {directory}")
+
+        if exclude_patterns is None:
+            exclude_patterns = []
+
+        # Common directories and files to exclude
+        default_excludes: list[str] = [
+            "__pycache__",
+            ".git",
+            ".github",
+            ".ssh",
+            ".gnupg",
+            ".config",
+            "node_modules",
+            "__pycache__",
+            ".venv",
+            "venv",
+            "env",
+            ".idea",
+            ".vscode",
+            ".DS_Store",
+        ]
+
+        exclude_patterns.extend(default_excludes)
+
+        def should_exclude(path: Path) -> bool:
+            """Check if a path should be excluded.
+
+            Args:
+                path: The path to check
+
+            Returns:
+                True if the path should be excluded, False otherwise
+            """
+            for pattern in exclude_patterns:
+                if pattern.startswith("*"):
+                    if path.name.endswith(pattern[1:]):
+                        return True
+                elif pattern in str(path):
+                    return True
+            return False
+
+        # Walk the directory
+        for root, dirs, files in os.walk(dir_path):
+            # Skip excluded directories
+            dirs[:] = [d for d in dirs if not should_exclude(Path(root) / d)]
+
+            # Process files
+            for file in files:
+                file_path: Path = Path(root) / file
+
+                if should_exclude(file_path):
+                    continue
+
+                try:
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        content: str = f.read()
+
+                    # Add to context
+                    self.add_document(str(file_path), content)
+                except UnicodeDecodeError:
+                    # Skip binary files
+                    continue
+
+            # Stop if not recursive
+            if not recursive:
+                break
+
+    def to_json(self) -> str:
+        """Convert the context to a JSON string.
+
+        Returns:
+            A JSON string representation of the context
+        """
+        data: dict[str, Any] = {
+            "documents": self.documents,
+            "metadata": self.document_metadata,
+            "modified_times": self.modified_times,
+            "allowed_paths": [str(p) for p in self.allowed_paths],
+        }
+
+        return json.dumps(data)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "DocumentContext":
+        """Create a context from a JSON string.
+
+        Args:
+            json_str: The JSON string
+
+        Returns:
+            A new DocumentContext instance
+        """
+        data: dict[str, Any] = json.loads(json_str)
+
+        context = cls()
+        context.documents = data.get("documents", {})
+        context.document_metadata = data.get("metadata", {})
+        context.modified_times = data.get("modified_times", {})
+        context.allowed_paths = set(Path(p) for p in data.get("allowed_paths", []))
+
+        return context

hanzo_mcp/tools/common/permissions.py

@@ -0,0 +1,253 @@
+"""Permission system for the Hanzo Dev MCP server."""
+
+import json
+import os
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import Any, TypeVar, final
+
+# Define type variables for better type annotations
+T = TypeVar("T")
+P = TypeVar("P")
+
+
+@final
+class PermissionManager:
+    """Manages permissions for file and command operations."""
+
+    def __init__(self) -> None:
+        """Initialize the permission manager."""
+        # Allowed paths
+        self.allowed_paths: set[Path] = set(
+            [Path("/tmp").resolve(), Path("/var").resolve()]
+        )
+
+        # Excluded paths
+        self.excluded_paths: set[Path] = set()
+        self.excluded_patterns: list[str] = []
+
+        # Default excluded patterns
+        self._add_default_exclusions()
+
+    def _add_default_exclusions(self) -> None:
+        """Add default exclusions for sensitive files and directories."""
+        # Sensitive directories
+        sensitive_dirs: list[str] = [
+            # ".git" is now allowed by default
+            ".ssh",
+            ".gnupg",
+            ".config",
+            "node_modules",
+            "__pycache__",
+            ".venv",
+            "venv",
+            "env",
+            ".idea",
+            ".vscode",
+            ".DS_Store",
+        ]
+        self.excluded_patterns.extend(sensitive_dirs)
+
+        # Sensitive file patterns
+        sensitive_patterns: list[str] = [
+            ".env",
+            "*.key",
+            "*.pem",
+            "*.crt",
+            "*password*",
+            "*secret*",
+            "*.sqlite",
+            "*.db",
+            "*.sqlite3",
+            "*.log",
+        ]
+        self.excluded_patterns.extend(sensitive_patterns)
+
+    def add_allowed_path(self, path: str) -> None:
+        """Add a path to the allowed paths.
+
+        Args:
+            path: The path to allow
+        """
+        resolved_path: Path = Path(path).resolve()
+        self.allowed_paths.add(resolved_path)
+
+    def remove_allowed_path(self, path: str) -> None:
+        """Remove a path from the allowed paths.
+
+        Args:
+            path: The path to remove
+        """
+        resolved_path: Path = Path(path).resolve()
+        if resolved_path in self.allowed_paths:
+            self.allowed_paths.remove(resolved_path)
+
+    def exclude_path(self, path: str) -> None:
+        """Exclude a path from allowed operations.
+
+        Args:
+            path: The path to exclude
+        """
+        resolved_path: Path = Path(path).resolve()
+        self.excluded_paths.add(resolved_path)
+
+    def add_exclusion_pattern(self, pattern: str) -> None:
+        """Add an exclusion pattern.
+
+        Args:
+            pattern: The pattern to exclude
+        """
+        self.excluded_patterns.append(pattern)
+
+    def is_path_allowed(self, path: str) -> bool:
+        """Check if a path is allowed.
+
+        Args:
+            path: The path to check
+
+        Returns:
+            True if the path is allowed, False otherwise
+        """
+        resolved_path: Path = Path(path).resolve()
+
+        # Check exclusions first
+        if self._is_path_excluded(resolved_path):
+            return False
+
+        # Check if the path is within any allowed path
+        for allowed_path in self.allowed_paths:
+            try:
+                resolved_path.relative_to(allowed_path)
+                return True
+            except ValueError:
+                continue
+
+        return False
+
+    def _is_path_excluded(self, path: Path) -> bool:
+        """Check if a path is excluded.
+
+        Args:
+            path: The path to check
+
+        Returns:
+            True if the path is excluded, False otherwise
+        """
+
+        # Check exact excluded paths
+        if path in self.excluded_paths:
+            return True
+
+        # Check excluded patterns
+        path_str: str = str(path)
+
+        # Get path parts to check for exact directory/file name matches
+        path_parts = path_str.split(os.sep)
+
+        for pattern in self.excluded_patterns:
+            # Handle wildcard patterns (e.g., "*.log")
+            if pattern.startswith("*"):
+                if path_str.endswith(pattern[1:]):
+                    return True
+            else:
+                # For non-wildcard patterns, check if any path component matches exactly
+                if pattern in path_parts:
+                    return True
+
+        return False
+
+    def to_json(self) -> str:
+        """Convert the permission manager to a JSON string.
+
+        Returns:
+            A JSON string representation of the permission manager
+        """
+        data: dict[str, Any] = {
+            "allowed_paths": [str(p) for p in self.allowed_paths],
+            "excluded_paths": [str(p) for p in self.excluded_paths],
+            "excluded_patterns": self.excluded_patterns,
+        }
+
+        return json.dumps(data)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "PermissionManager":
+        """Create a permission manager from a JSON string.
+
+        Args:
+            json_str: The JSON string
+
+        Returns:
+            A new PermissionManager instance
+        """
+        data: dict[str, Any] = json.loads(json_str)
+
+        manager = cls()
+
+        for path in data.get("allowed_paths", []):
+            manager.add_allowed_path(path)
+
+        for path in data.get("excluded_paths", []):
+            manager.exclude_path(path)
+
+        manager.excluded_patterns = data.get("excluded_patterns", [])
+
+        return manager
+
+
+class PermissibleOperation:
+    """A decorator for operations that require permission."""
+
+    def __init__(
+        self,
+        permission_manager: PermissionManager,
+        operation: str,
+        get_path_fn: Callable[[list[Any], dict[str, Any]], str] | None = None,
+    ) -> None:
+        """Initialize the permissible operation.
+
+        Args:
+            permission_manager: The permission manager
+            operation: The operation type (read, write, execute, etc.)
+            get_path_fn: Optional function to extract the path from args and kwargs
+        """
+        self.permission_manager: PermissionManager = permission_manager
+        self.operation: str = operation
+        self.get_path_fn: Callable[[list[Any], dict[str, Any]], str] | None = (
+            get_path_fn
+        )
+
+    def __call__(
+        self, func: Callable[..., Awaitable[T]]
+    ) -> Callable[..., Awaitable[T]]:
+        """Decorate the function.
+
+        Args:
+            func: The function to decorate
+
+        Returns:
+            The decorated function
+        """
+
+        async def wrapper(*args: Any, **kwargs: Any) -> T:
+            # Extract the path
+            if self.get_path_fn:
+                # Pass args as a list and kwargs as a dict to the path function
+                path = self.get_path_fn(list(args), kwargs)
+            else:
+                # Default to first argument
+                path = args[0] if args else next(iter(kwargs.values()), None)
+
+            if not isinstance(path, str):
+                raise ValueError(f"Invalid path type: {type(path)}")
+
+            # Check permission
+            if not self.permission_manager.is_path_allowed(path):
+                raise PermissionError(
+                    f"Operation '{self.operation}' not allowed for path: {path}"
+                )
+
+            # Call the function
+            return await func(*args, **kwargs)
+
+        return wrapper