hanzo-mcp 0.3.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

hanzo_mcp/tools/common/batch_tool.py

@@ -0,0 +1,330 @@
+"""Batch tool implementation for Hanzo MCP.
+
+This module provides the BatchTool that allows for executing multiple tools in
+parallel or serial depending on their characteristics.
+"""
+
+import asyncio
+from typing import Annotated, Any, TypedDict, Unpack, final, override
+
+from fastmcp import Context as MCPContext
+from fastmcp import FastMCP
+from fastmcp.server.dependencies import get_context
+from pydantic import Field
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import create_tool_context
+
+
+class InvocationItem(TypedDict):
+    """A single tool invocation item.
+
+    Attributes:
+        tool_name: The name of the tool to invoke
+        input: The input to pass to the tool
+    """
+
+    tool_name: Annotated[
+        str,
+        Field(
+            description="The name of the tool to invoke",
+            min_length=1,
+        ),
+    ]
+    input: Annotated[
+        dict[str, Any],
+        Field(
+            description="The input to pass to the tool",
+        ),
+    ]
+
+
+Description = Annotated[
+    str,
+    Field(
+        description="A short (3-5 word) description of the batch operation",
+        min_length=1,
+    ),
+]
+
+Invocations = Annotated[
+    list[InvocationItem],
+    Field(
+        description="The list of tool invocations to execute (required -- you MUST provide at least one tool invocation)",
+        min_length=1,
+    ),
+]
+
+
+class BatchToolParams(TypedDict):
+    """Parameters for the BatchTool.
+
+    Attributes:
+        description: A short (3-5 word) description of the batch operation
+        invocations: The list of tool invocations to execute (required -- you MUST provide at least one tool invocation)
+    """
+
+    description: Description
+    invocations: Invocations
+
+
+@final
+class BatchTool(BaseTool):
+    """Tool for executing multiple tools in a single request.
+
+    Executes a list of tool invocations in parallel when possible, or
+    otherwise serially. Returns the collected results from all invocations.
+    """
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name.
+
+        Returns:
+            Tool name
+        """
+        return "batch"
+
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description.
+
+        Returns:
+            Tool description
+        """
+        return """Batch execution tool that runs multiple tool invocations in a single request.
+
+Tools are executed in parallel when possible, and otherwise serially.
+Takes a list of tool invocations (tool_name and input pairs).
+Returns the collected results from all invocations.
+Use this tool when you need to run multiple independent tool operations at once -- it is awesome for speeding up your workflow, reducing both context usage and latency.
+Each tool will respect its own permissions and validation rules.
+The tool's outputs are NOT shown to the user; to answer the user's query, you MUST send a message with the results after the tool call completes, otherwise the user will not see the results.
+
+<batch_example>
+When dispatching multiple agents to find necessary information.
+batch(
+  description="Update import statements across modules",
+  invocations=[
+    {tool_name: "dispatch_agent", input: {prompt: "Search for all instances of 'logger' configuration in /app/config directory"}},
+    {tool_name: "dispatch_agent", input: {prompt: "Find all test files that reference 'UserService' in /app/tests"}},
+  ]
+)
+
+Common scenarios for effective batching:
+1. Reading multiple related files in one operation
+2. Performing a series of simple mechanical changes
+3. Running multiple diagnostic commands
+4. Dispatch multiple agents to complete the task
+
+To make a batch call, provide the following:
+1. description: A short (3-5 word) description of the batch operation
+2. invocations: List of invocation [{"tool_name": "...", "input": "..."}], tool_name: The name of the tool to invoke,newText: The input to pass to the tool
+
+
+Available tools in batch call:
+Tool: dispatch_agent,read,directory_tree,grep,grep_ast,run_command,notebook_read
+Not available: think,write,edit,multi_edit,notebook_edit
+"""
+
+    def __init__(self, tools: dict[str, BaseTool]) -> None:
+        """Initialize the batch tool.
+
+        Args:
+            tools: Dictionary mapping tool names to tool instances
+        """
+        self.tools = tools
+
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[BatchToolParams],
+    ) -> str:
+        """Execute the tool with the given parameters.
+
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+
+        Returns:
+            Tool result
+        """
+        tool_ctx = create_tool_context(ctx)
+        tool_ctx.set_tool_info(self.name)
+
+        # Extract parameters
+        description = params.get("description")
+        invocations: list[dict[str, Any]] = params.get("invocations", list())
+
+        # Validate required parameters
+        if not description:
+            await tool_ctx.error(
+                "Parameter 'description' is required but was None or empty"
+            )
+            return "Error: Parameter 'description' is required but was None or empty"
+
+        if not invocations:
+            await tool_ctx.error(
+                "Parameter 'invocations' is required but was None or empty"
+            )
+            return "Error: Parameter 'invocations' is required but was None or empty"
+
+        if not isinstance(invocations, list) or len(invocations) == 0:
+            await tool_ctx.error("Parameter 'invocations' must be a non-empty list")
+            return "Error: Parameter 'invocations' must be a non-empty list"
+
+        await tool_ctx.info(
+            f"Executing batch operation: {description} ({len(invocations)} invocations)"
+        )
+
+        # Execute all tool invocations in parallel
+        tasks: list[asyncio.Future[dict[str, Any]]] = []
+        invocation_map: dict[
+            asyncio.Future[dict[str, Any]], dict[str, Any]
+        ] = {}  # Map task Future to invocation
+
+        for i, invocation in enumerate(invocations):
+            # Extract tool name and input from invocation
+            tool_name: str = invocation.get("tool_name", "")
+            tool_input: dict[str, Any] = invocation.get("input", {})
+
+            # Validate tool invocation
+            if not tool_name:
+                error_message = f"Tool name is required in invocation {i}"
+                await tool_ctx.error(error_message)
+                # Add direct result for this invocation
+                tasks.append(asyncio.Future())
+                tasks[-1].set_result(
+                    {"invocation": invocation, "result": f"Error: {error_message}"}
+                )
+                invocation_map[tasks[-1]] = invocation
+                continue
+
+            # Check if the tool exists
+            if tool_name not in self.tools:
+                error_message = f"Tool '{tool_name}' not found"
+                await tool_ctx.error(error_message)
+                # Add direct result for this invocation
+                tasks.append(asyncio.Future())
+                tasks[-1].set_result(
+                    {"invocation": invocation, "result": f"Error: {error_message}"}
+                )
+                invocation_map[tasks[-1]] = invocation
+                continue
+
+            # Create a task for this tool invocation
+            try:
+                tool = self.tools[tool_name]
+                await tool_ctx.info(f"Creating task for tool: {tool_name}")
+
+                # Create coroutine for this tool execution
+                async def execute_tool(
+                    tool_obj: BaseTool, tool_name: str, tool_input: dict[str, Any]
+                ):
+                    try:
+                        await tool_ctx.info(f"Executing tool: {tool_name}")
+                        result = await tool_obj.call(ctx, **tool_input)
+                        await tool_ctx.info(f"Tool '{tool_name}' execution completed")
+                        return {
+                            "invocation": {"tool_name": tool_name, "input": tool_input},
+                            "result": result,
+                        }
+                    except Exception as e:
+                        error_message = f"Error executing tool '{tool_name}': {str(e)}"
+                        await tool_ctx.error(error_message)
+                        return {
+                            "invocation": {"tool_name": tool_name, "input": tool_input},
+                            "result": f"Error: {error_message}",
+                        }
+
+                # Schedule the task
+                task = asyncio.create_task(execute_tool(tool, tool_name, tool_input))
+                tasks.append(task)
+                invocation_map[task] = invocation
+            except Exception as e:
+                error_message = f"Error scheduling tool '{tool_name}': {str(e)}"
+                await tool_ctx.error(error_message)
+                # Add direct result for this invocation
+                tasks.append(asyncio.Future())
+                tasks[-1].set_result(
+                    {"invocation": invocation, "result": f"Error: {error_message}"}
+                )
+                invocation_map[tasks[-1]] = invocation
+
+        # Wait for all tasks to complete
+        await tool_ctx.info(f"Waiting for {len(tasks)} tool executions to complete")
+        results: list[dict[str, Any]] = []
+
+        # As tasks complete, collect their results
+        for task in asyncio.as_completed(tasks):
+            try:
+                result = await task
+                results.append(result)
+            except Exception as e:
+                invocation = invocation_map[task]
+                tool_name: str = invocation.get("tool_name", "unknown")
+                error_message = f"Unexpected error in tool '{tool_name}': {str(e)}"
+                await tool_ctx.error(error_message)
+                results.append(
+                    {"invocation": invocation, "result": f"Error: {error_message}"}
+                )
+
+        # Format the results
+        formatted_results = self._format_results(results)
+        await tool_ctx.info(
+            f"Batch operation '{description}' completed with {len(results)} results"
+        )
+
+        return formatted_results
+
+    def _format_results(self, results: list[dict[str, dict[str, Any]]]) -> str:
+        """Format the results from multiple tool invocations.
+
+        Args:
+            results: List of tool invocation results
+
+        Returns:
+            Formatted results string
+        """
+        formatted_parts: list[str] = []
+        for i, result in enumerate(results):
+            invocation: dict[str, Any] = result["invocation"]
+            tool_name: str = invocation.get("tool_name", "unknown")
+
+            # Add the result header
+            formatted_parts.append(f"### Result {i + 1}: {tool_name}")
+            # Add the result content - use multi-line code blocks for code outputs
+            if "\n" in result["result"]:
+                formatted_parts.append(f"```\n{result['result']}\n```")
+            else:
+                formatted_parts.append(result["result"])
+            # Add a separator
+            formatted_parts.append("")
+
+        return "\n".join(formatted_parts)
+
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this batch tool with the MCP server.
+
+        Creates a wrapper function with explicitly defined parameters that match
+        the tool's parameter schema and registers it with the MCP server.
+
+        Args:
+            mcp_server: The FastMCP server instance
+        """
+        tool_self = self  # Create a reference to self for use in the closure
+
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def batch(
+            ctx: MCPContext,
+            description: Description,
+            invocations: Invocations,
+        ) -> str:
+            ctx = get_context()
+            return await tool_self.call(
+                ctx, description=description, invocations=invocations
+            )

hanzo_mcp/tools/common/context.py

@@ -11,7 +11,7 @@ from collections.abc import Iterable
 from pathlib import Path
 from typing import Any, ClassVar, final
 
-from mcp.server.fastmcp import Context as MCPContext
+from fastmcp import Context as MCPContext
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 
 
@@ -87,7 +87,11 @@ class ToolContext:
         Args:
             message: The message to log
         """
-        await self._mcp_context.info(self._format_message(message))
+        try:
+            await self._mcp_context.info(self._format_message(message))
+        except Exception:
+            # Silently ignore errors when client has disconnected
+            pass
 
     async def debug(self, message: str) -> None:
         """Log a debug message.
@@ -95,7 +99,11 @@ class ToolContext:
         Args:
             message: The message to log
         """
-        await self._mcp_context.debug(self._format_message(message))
+        try:
+            await self._mcp_context.debug(self._format_message(message))
+        except Exception:
+            # Silently ignore errors when client has disconnected
+            pass
 
     async def warning(self, message: str) -> None:
         """Log a warning message.
@@ -103,7 +111,11 @@ class ToolContext:
         Args:
             message: The message to log
         """
-        await self._mcp_context.warning(self._format_message(message))
+        try:
+            await self._mcp_context.warning(self._format_message(message))
+        except Exception:
+            # Silently ignore errors when client has disconnected
+            pass
 
     async def error(self, message: str) -> None:
         """Log an error message.
@@ -111,7 +123,11 @@ class ToolContext:
         Args:
             message: The message to log
         """
-        await self._mcp_context.error(self._format_message(message))
+        try:
+            await self._mcp_context.error(self._format_message(message))
+        except Exception:
+            # Silently ignore errors when client has disconnected
+            pass
 
     def _format_message(self, message: str) -> str:
         """Format a message with tool information if available.
@@ -135,7 +151,11 @@ class ToolContext:
             current: Current progress value
             total: Total progress value
         """
-        await self._mcp_context.report_progress(current, total)
+        try:
+            await self._mcp_context.report_progress(current, total)
+        except Exception:
+            # Silently ignore errors when client has disconnected
+            pass
 
     async def read_resource(self, uri: str) -> Iterable[ReadResourceContents]:
         """Read a resource via the MCP protocol.
@@ -160,289 +180,3 @@ def create_tool_context(mcp_context: MCPContext) -> ToolContext:
         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
-        """
-        # Expand user path (e.g., ~/ or $HOME)
-        expanded_path = os.path.expanduser(path)
-        resolved_path: Path = Path(expanded_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
-        """
-        # Expand user path (e.g., ~/ or $HOME)
-        expanded_path = os.path.expanduser(path)
-        resolved_path: Path = Path(expanded_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