hanzo-mcp 0.3.8__py3-none-any.whl → 0.5.1__py3-none-any.whl

hanzo_mcp/prompts/utils.py

@@ -0,0 +1,286 @@
+import platform
+from pathlib import Path
+
+try:
+    from git import Repo
+
+    GIT_AVAILABLE = True
+except ImportError:
+    GIT_AVAILABLE = False
+
+
+def get_os_info() -> tuple[str, str, str]:
+    """Get the operating system information.
+    Returns:
+        tuple: A tuple containing the system name, release, and version.
+    """
+    system = platform.system()  # noqa: F821
+    release = platform.release()
+    version = platform.version()
+
+    if system == "Darwin":
+        system = "MacOS"
+    elif system == "Linux":
+        try:
+            with open("/etc/os-release") as f:
+                for line in f:
+                    if line.startswith("NAME="):
+                        name = line.split("=")[1].strip().strip('"')
+                        if "Ubuntu" in name:
+                            system = "Ubuntu"
+                        elif "Debian" in name:
+                            system = "Debian"
+                        elif "Fedora" in name:
+                            system = "Fedora"
+                        elif "CentOS" in name:
+                            system = "CentOS"
+                        elif "Arch Linux" in name:
+                            system = "Arch Linux"
+                        system = name
+        except FileNotFoundError:
+            dist = platform.freedesktop_os_release()
+            if dist and "NAME" in dist:
+                name = dist["NAME"]
+                if "Ubuntu" in name:
+                    system = "Ubuntu"
+                system = name
+            system = "Linux"
+    elif system == "Java":
+        system = "Java"
+
+    return system, release, version
+
+
+def get_directory_structure(
+    path: str, max_depth: int = 3, include_filtered: bool = False
+) -> str:
+    """Get a directory structure similar to directory_tree tool.
+
+    Args:
+        path: The directory path to scan
+        max_depth: Maximum depth to traverse (0 for unlimited)
+        include_filtered: Whether to include normally filtered directories
+
+    Returns:
+        Formatted directory structure as a string
+    """
+    try:
+        dir_path = Path(path)
+
+        if not dir_path.exists() or not dir_path.is_dir():
+            return f"Error: {path} is not a valid directory"
+
+        # Define filtered directories (same as directory_tree.py)
+        FILTERED_DIRECTORIES = {
+            ".git",
+            "node_modules",
+            ".venv",
+            "venv",
+            "__pycache__",
+            ".pytest_cache",
+            ".idea",
+            ".vs",
+            ".vscode",
+            "dist",
+            "build",
+            "target",
+            ".ruff_cache",
+            ".llm-context",
+        }
+
+        def should_filter(current_path: Path) -> bool:
+            """Check if a directory should be filtered."""
+            # Don't filter if it's the explicitly requested path
+            if str(current_path.absolute()) == str(dir_path.absolute()):
+                return False
+            # Filter based on directory name if filtering is enabled
+            return current_path.name in FILTERED_DIRECTORIES and not include_filtered
+
+        def build_tree(current_path: Path, current_depth: int = 0) -> list[dict]:
+            """Build directory tree recursively."""
+            result = []
+
+            try:
+                # Sort entries: directories first, then files alphabetically
+                entries = sorted(
+                    current_path.iterdir(), key=lambda x: (not x.is_dir(), x.name)
+                )
+
+                for entry in entries:
+                    if entry.is_dir():
+                        entry_data = {"name": entry.name, "type": "directory"}
+
+                        # Check if we should filter this directory
+                        if should_filter(entry):
+                            entry_data["skipped"] = "filtered-directory"
+                            result.append(entry_data)
+                            continue
+
+                        # Check depth limit (if enabled)
+                        if max_depth > 0 and current_depth >= max_depth:
+                            entry_data["skipped"] = "depth-limit"
+                            result.append(entry_data)
+                            continue
+
+                        # Process children recursively
+                        entry_data["children"] = build_tree(entry, current_depth + 1)
+                        result.append(entry_data)
+                    else:
+                        # Add files only if within depth limit
+                        if max_depth <= 0 or current_depth < max_depth:
+                            result.append({"name": entry.name, "type": "file"})
+
+            except Exception:
+                # Skip directories we can't read
+                pass
+
+            return result
+
+        def format_tree(tree_data: list[dict], level: int = 0) -> list[str]:
+            """Format tree data as indented strings."""
+            lines = []
+
+            for item in tree_data:
+                # Indentation based on level
+                indent = "  " * level
+
+                # Format based on type
+                if item["type"] == "directory":
+                    if "skipped" in item:
+                        lines.append(
+                            f"{indent}{item['name']}/ [skipped - {item['skipped']}]"
+                        )
+                    else:
+                        lines.append(f"{indent}{item['name']}/")
+                        # Add children with increased indentation if present
+                        if "children" in item:
+                            lines.extend(format_tree(item["children"], level + 1))
+                else:
+                    # File
+                    lines.append(f"{indent}{item['name']}")
+
+            return lines
+
+        # Build and format the tree
+        tree_data = build_tree(dir_path)
+        formatted_lines = format_tree(tree_data)
+
+        # Add the root directory path as a prefix
+        result = f"- {dir_path}/"
+        if formatted_lines:
+            result += "\n" + "\n".join(f"  {line}" for line in formatted_lines)
+
+        return result
+
+    except Exception as e:
+        return f"Error generating directory structure: {str(e)}"
+
+
+def get_git_info(path: str) -> dict[str, str | None]:
+    """Get git information for a repository.
+
+    Args:
+        path: Path to the git repository
+
+    Returns:
+        Dictionary containing git information
+    """
+    if not GIT_AVAILABLE:
+        return {
+            "current_branch": None,
+            "main_branch": None,
+            "git_status": "GitPython not available",
+            "recent_commits": "GitPython not available",
+        }
+
+    try:
+        repo = Repo(path)
+
+        # Get current branch
+        try:
+            current_branch = repo.active_branch.name
+        except Exception:
+            current_branch = "HEAD (detached)"
+
+        # Try to determine main branch
+        main_branch = "main"  # default
+        try:
+            # Check if 'main' exists
+            if "origin/main" in [ref.name for ref in repo.refs]:
+                main_branch = "main"
+            elif "origin/master" in [ref.name for ref in repo.refs]:
+                main_branch = "master"
+            elif "main" in [ref.name for ref in repo.refs]:
+                main_branch = "main"
+            elif "master" in [ref.name for ref in repo.refs]:
+                main_branch = "master"
+        except Exception:
+            pass
+
+        # Get git status
+        try:
+            status_lines = []
+
+            # Check for staged changes
+            staged_files = list(repo.index.diff("HEAD"))
+            if staged_files:
+                for item in staged_files[:25]:  # Limit to first 25
+                    change_type = item.change_type
+                    status_lines.append(f"{change_type[0].upper()} {item.a_path}")
+                if len(staged_files) > 25:
+                    status_lines.append(
+                        f"... and {len(staged_files) - 25} more staged files"
+                    )
+
+            # Check for unstaged changes
+            unstaged_files = list(repo.index.diff(None))
+            if unstaged_files:
+                for item in unstaged_files[:25]:  # Limit to first 25
+                    status_lines.append(f"M {item.a_path}")
+                if len(unstaged_files) > 25:
+                    status_lines.append(
+                        f"... and {len(unstaged_files) - 25} more modified files"
+                    )
+
+            # Check for untracked files
+            untracked_files = repo.untracked_files
+            if untracked_files:
+                for file in untracked_files[:25]:  # Limit to first 25
+                    status_lines.append(f"?? {file}")
+                if len(untracked_files) > 25:
+                    status_lines.append(
+                        f"... and {len(untracked_files) - 25} more untracked files"
+                    )
+
+            git_status = (
+                "\n".join(status_lines) if status_lines else "Working tree clean"
+            )
+
+        except Exception:
+            git_status = "Unable to get git status"
+
+        # Get recent commits
+        try:
+            commits = []
+            for commit in repo.iter_commits(max_count=5):
+                short_hash = commit.hexsha[:7]
+                message = commit.message.split("\n")[0]  # First line only
+                commits.append(f"{short_hash} {message}")
+            recent_commits = "\n".join(commits)
+        except Exception:
+            recent_commits = "Unable to get recent commits"
+
+        return {
+            "current_branch": current_branch,
+            "main_branch": main_branch,
+            "git_status": git_status,
+            "recent_commits": recent_commits,
+        }
+
+    except Exception as e:
+        return {
+            "current_branch": None,
+            "main_branch": None,
+            "git_status": f"Error: {str(e)}",
+            "recent_commits": f"Error: {str(e)}",
+        }

hanzo_mcp/server.py

@@ -1,121 +1,189 @@
-"""MCP server implementing Hanzo capabilities.
-
-Includes improved error handling and debugging for tool execution.
-"""
+"""MCP server implementing Hanzo capabilities."""
 
+import atexit
+import signal
+import threading
+import time
 from typing import Literal, cast, final
 
-from mcp.server.fastmcp import FastMCP
+from fastmcp import FastMCP
 
+from hanzo_mcp.prompts import register_all_prompts
 from hanzo_mcp.tools import register_all_tools
-from hanzo_mcp.tools.common.context import DocumentContext
+
 from hanzo_mcp.tools.common.permissions import PermissionManager
-from hanzo_mcp.tools.project.analysis import ProjectAnalyzer, ProjectManager
-from hanzo_mcp.tools.shell.command_executor import CommandExecutor
+from hanzo_mcp.tools.shell.session_storage import SessionStorage
 
 
 @final
-class HanzoServer:
-    """MCP server implementing Hanzo capabilities.
-
-Includes improved error handling and debugging for tool execution.
-"""
+class HanzoMCPServer:
+    """MCP server implementing Hanzo capabilities."""
 
     def __init__(
         self,
-        name: str = "claude-code",
+        name: str = "hanzo",
         allowed_paths: list[str] | None = None,
+        project_paths: list[str] | None = None,
         project_dir: str | None = None,
         mcp_instance: FastMCP | None = None,
         agent_model: str | None = None,
         agent_max_tokens: int | None = None,
         agent_api_key: str | None = None,
+        agent_base_url: str | None = None,
         agent_max_iterations: int = 10,
         agent_max_tool_uses: int = 30,
         enable_agent_tool: bool = False,
+        command_timeout: float = 120.0,
         disable_write_tools: bool = False,
         disable_search_tools: bool = False,
-        host: str = "0.0.0.0",
-        port: int = 3001,
+        host: str = "127.0.0.1",
+        port: int = 3000,
+        enabled_tools: dict[str, bool] | None = None,
+        disabled_tools: list[str] | None = None,
     ):
-        """Initialize the Hanzo server.
+        """Initialize the Hanzo MCP server.
 
         Args:
             name: The name of the server
             allowed_paths: list of paths that the server is allowed to access
-            project_dir: Optional project directory to use as initial working directory
+            project_paths: list of project paths to generate prompts for
+            project_dir: single project directory (added to allowed_paths and project_paths)
             mcp_instance: Optional FastMCP instance for testing
             agent_model: Optional model name for agent tool in LiteLLM format
             agent_max_tokens: Optional maximum tokens for agent responses
             agent_api_key: Optional API key for the LLM provider
+            agent_base_url: Optional base URL for the LLM provider API endpoint
             agent_max_iterations: Maximum number of iterations for agent (default: 10)
             agent_max_tool_uses: Maximum number of total tool uses for agent (default: 30)
             enable_agent_tool: Whether to enable the agent tool (default: False)
-            disable_write_tools: Whether to disable write/edit tools (default: False)
+            command_timeout: Default timeout for command execution in seconds (default: 120.0)
+            disable_write_tools: Whether to disable write tools (default: False)
             disable_search_tools: Whether to disable search tools (default: False)
-            host: Host to bind to for SSE transport (default: '0.0.0.0')
-            port: Port to use for SSE transport (default: 3001)
+            host: Host for SSE server (default: 127.0.0.1)
+            port: Port for SSE server (default: 3000)
+            enabled_tools: Dictionary of individual tool enable states (default: None)
+            disabled_tools: List of tool names to disable (default: None)
         """
         self.mcp = mcp_instance if mcp_instance is not None else FastMCP(name)
 
-        # Initialize context, permissions, and command executor
-        self.document_context = DocumentContext()
+        # Initialize permissions and command executor
         self.permission_manager = PermissionManager()
 
-        # Initialize command executor
-        self.command_executor = CommandExecutor(
-            permission_manager=self.permission_manager,
-            verbose=False,  # Set to True for debugging
-        )
-
-        # If project_dir is specified, set it as initial working directory for all sessions
+        # Handle project_dir parameter
         if project_dir:
-            initial_session_id = name  # Use server name as default session ID
-            self.command_executor.set_working_dir(initial_session_id, project_dir)
-
-        # Initialize project analyzer
-        self.project_analyzer = ProjectAnalyzer(self.command_executor)
-
-        # Initialize project manager
-        self.project_manager = ProjectManager(
-            self.document_context, self.permission_manager, self.project_analyzer
-        )
+            if allowed_paths is None:
+                allowed_paths = []
+            if project_dir not in allowed_paths:
+                allowed_paths.append(project_dir)
+            if project_paths is None:
+                project_paths = []
+            if project_dir not in project_paths:
+                project_paths.append(project_dir)
 
         # Add allowed paths
         if allowed_paths:
             for path in allowed_paths:
                 self.permission_manager.add_allowed_path(path)
-                self.document_context.add_allowed_path(path)
+
+        # Store paths and options
+        self.project_paths = project_paths
+        self.project_dir = project_dir
+        self.disable_write_tools = disable_write_tools
+        self.disable_search_tools = disable_search_tools
+        self.host = host
+        self.port = port
+        self.enabled_tools = enabled_tools or {}
+        self.disabled_tools = disabled_tools or []
 
         # Store agent options
         self.agent_model = agent_model
         self.agent_max_tokens = agent_max_tokens
         self.agent_api_key = agent_api_key
+        self.agent_base_url = agent_base_url
         self.agent_max_iterations = agent_max_iterations
         self.agent_max_tool_uses = agent_max_tool_uses
         self.enable_agent_tool = enable_agent_tool
-        self.disable_write_tools = disable_write_tools
-        self.disable_search_tools = disable_search_tools
+        self.command_timeout = command_timeout
 
-        # Store network options
-        self.host = host
-        self.port = port
+        # Initialize cleanup tracking
+        self._cleanup_thread: threading.Thread | None = None
+        self._shutdown_event = threading.Event()
+        self._cleanup_registered = False
+
+        # Apply disabled_tools to enabled_tools
+        final_enabled_tools = self.enabled_tools.copy()
+        for tool_name in self.disabled_tools:
+            final_enabled_tools[tool_name] = False
+        
+        # Store the final processed tool configuration
+        self.enabled_tools = final_enabled_tools
 
         # Register all tools
         register_all_tools(
             mcp_server=self.mcp,
-            document_context=self.document_context,
             permission_manager=self.permission_manager,
             agent_model=self.agent_model,
             agent_max_tokens=self.agent_max_tokens,
             agent_api_key=self.agent_api_key,
+            agent_base_url=self.agent_base_url,
             agent_max_iterations=self.agent_max_iterations,
             agent_max_tool_uses=self.agent_max_tool_uses,
             enable_agent_tool=self.enable_agent_tool,
             disable_write_tools=self.disable_write_tools,
             disable_search_tools=self.disable_search_tools,
+            enabled_tools=final_enabled_tools,
         )
 
+        register_all_prompts(mcp_server=self.mcp, projects=self.project_paths)
+
+    def _setup_cleanup_handlers(self) -> None:
+        """Set up signal handlers and background cleanup thread."""
+        if self._cleanup_registered:
+            return
+
+        # Register cleanup on normal exit
+        atexit.register(self._cleanup_sessions)
+
+        # Register signal handlers for graceful shutdown
+        def signal_handler(signum, frame):
+            self._cleanup_sessions()
+            self._shutdown_event.set()
+
+        signal.signal(signal.SIGTERM, signal_handler)
+        signal.signal(signal.SIGINT, signal_handler)
+
+        # Start background cleanup thread for periodic cleanup
+        self._cleanup_thread = threading.Thread(
+            target=self._background_cleanup, daemon=True
+        )
+        self._cleanup_thread.start()
+
+        self._cleanup_registered = True
+
+    def _background_cleanup(self) -> None:
+        """Background thread for periodic session cleanup."""
+        while not self._shutdown_event.is_set():
+            try:
+                # Clean up expired sessions every 2 minutes
+                # Using shorter TTL of 5 minutes (300 seconds)
+                SessionStorage.cleanup_expired_sessions(max_age_seconds=300)
+
+                # Wait for 2 minutes or until shutdown
+                self._shutdown_event.wait(timeout=120)
+            except Exception:
+                # Ignore cleanup errors and continue
+                pass
+
+    def _cleanup_sessions(self) -> None:
+        """Clean up all active sessions."""
+        try:
+            cleared_count = SessionStorage.clear_all_sessions()
+            if cleared_count > 0:
+                print(f"Cleaned up {cleared_count} tmux sessions on shutdown")
+        except Exception:
+            # Ignore cleanup errors during shutdown
+            pass
+
     def run(self, transport: str = "stdio", allowed_paths: list[str] | None = None):
         """Run the MCP server.
 
@@ -127,60 +195,10 @@ Includes improved error handling and debugging for tool execution.
         allowed_paths_list = allowed_paths or []
         for path in allowed_paths_list:
             self.permission_manager.add_allowed_path(path)
-            self.document_context.add_allowed_path(path)
 
-        # If using SSE, set the port and host in the environment variables
-        if transport == "sse":
-            import os
-            # Set environment variables for FastMCP settings
-            os.environ["FASTMCP_PORT"] = str(self.port)
-            os.environ["FASTMCP_HOST"] = self.host
-            print(f"Starting SSE server on {self.host}:{self.port}")
+        # Set up cleanup handlers before running
+        self._setup_cleanup_handlers()
 
         # Run the server
         transport_type = cast(Literal["stdio", "sse"], transport)
         self.mcp.run(transport=transport_type)
-
-
-def main():
-    """Run the Hanzo MCP server."""
-    import argparse
-
-    parser = argparse.ArgumentParser(
-        description="MCP server implementing Hanzo capabilities"
-    )
-
-    _ = parser.add_argument(
-        "--name",
-        default="claude-code",
-        help="Name of the MCP server (default: claude-code)",
-    )
-
-    _ = parser.add_argument(
-        "--transport",
-        choices=["stdio", "sse"],
-        default="stdio",
-        help="Transport protocol to use (default: stdio)",
-    )
-
-    _ = parser.add_argument(
-        "--allow-path",
-        action="append",
-        dest="allowed_paths",
-        help="Add an allowed path (can be specified multiple times)",
-    )
-
-    args = parser.parse_args()
-
-    # Type annotations for args to avoid Any warnings
-    name: str = args.name
-    transport: str = args.transport
-    allowed_paths: list[str] | None = args.allowed_paths
-
-    # Create and run the server
-    server = HanzoServer(name=name, allowed_paths=allowed_paths)
-    server.run(transport=transport, allowed_paths=allowed_paths or [])
-
-
-if __name__ == "__main__":
-    main()