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

hanzo_mcp/tools/filesystem/write.py

@@ -0,0 +1,156 @@
+"""Write file tool implementation.
+
+This module provides the Write tool for creating or overwriting files.
+"""
+
+from pathlib import Path
+from typing import Annotated, 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.filesystem.base import FilesystemBaseTool
+
+FilePath = Annotated[
+    str,
+    Field(
+        description="The absolute path to the file to write (must be absolute, not relative)",
+        min_length=1,
+    ),
+]
+
+Content = Annotated[
+    str,
+    Field(
+        description="The content to write to the file",
+        min_length=1,
+    ),
+]
+
+
+class WriteToolParams(TypedDict):
+    """Parameters for the Write tool.
+
+    Attributes:
+        file_path: The absolute path to the file to write (must be absolute, not relative)
+        content: The content to write to the file
+    """
+
+    file_path: FilePath
+    content: Content
+
+
+@final
+class Write(FilesystemBaseTool):
+    """Tool for writing file contents."""
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name.
+
+        Returns:
+            Tool name
+        """
+        return "write"
+
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description.
+
+        Returns:
+            Tool description
+        """
+        return """Writes a file to the local filesystem.
+
+Usage:
+- This tool will overwrite the existing file if there is one at the provided path.
+- If this is an existing file, you MUST use the Read tool first to read the file's contents. This tool will fail if you did not read the file first.
+- ALWAYS prefer editing existing files in the codebase. NEVER write new files unless explicitly required.
+- NEVER proactively create documentation files (*.md) or README files. Only create documentation files if explicitly requested by the User."""
+
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[WriteToolParams],
+    ) -> str:
+        """Execute the tool with the given parameters.
+
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+
+        Returns:
+            Tool result
+        """
+        tool_ctx = self.create_tool_context(ctx)
+        self.set_tool_context_info(tool_ctx)
+
+        # Extract parameters
+        file_path: FilePath = params["file_path"]
+        content: Content = params["content"]
+
+        # Validate parameters
+        path_validation = self.validate_path(file_path)
+        if path_validation.is_error:
+            await tool_ctx.error(path_validation.error_message)
+            return f"Error: {path_validation.error_message}"
+
+        await tool_ctx.info(f"Writing file: {file_path}")
+
+        # Check if file is allowed to be written
+        allowed, error_msg = await self.check_path_allowed(file_path, tool_ctx)
+        if not allowed:
+            return error_msg
+
+        # Additional check already verified by is_path_allowed above
+        await tool_ctx.info(f"Writing file: {file_path}")
+
+        try:
+            path_obj = Path(file_path)
+
+            # Check if parent directory is allowed
+            parent_dir = str(path_obj.parent)
+            if not self.is_path_allowed(parent_dir):
+                await tool_ctx.error(f"Parent directory not allowed: {parent_dir}")
+                return f"Error: Parent directory not allowed: {parent_dir}"
+
+            # Create parent directories if they don't exist
+            path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+            # Write the file
+            with open(path_obj, "w", encoding="utf-8") as f:
+                f.write(content)
+
+            await tool_ctx.info(
+                f"Successfully wrote file: {file_path} ({len(content)} bytes)"
+            )
+            return f"Successfully wrote file: {file_path} ({len(content)} bytes)"
+        except Exception as e:
+            await tool_ctx.error(f"Error writing file: {str(e)}")
+            return f"Error writing file: {str(e)}"
+
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this 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 write(
+            ctx: MCPContext,
+            file_path: FilePath,
+            content: Content,
+        ) -> str:
+            ctx = get_context()
+            return await tool_self.call(ctx, file_path=file_path, content=content)

hanzo_mcp/tools/jupyter/__init__.py

@@ -4,73 +4,85 @@ This package provides tools for working with Jupyter notebooks (.ipynb files),
 including reading and editing notebook cells.
 """
 
-from mcp.server.fastmcp import FastMCP
+from fastmcp import FastMCP
 
 from hanzo_mcp.tools.common.base import BaseTool, ToolRegistry
-from hanzo_mcp.tools.common.context import DocumentContext
 from hanzo_mcp.tools.common.permissions import PermissionManager
-from hanzo_mcp.tools.jupyter.edit_notebook import EditNotebookTool
-from hanzo_mcp.tools.jupyter.read_notebook import ReadNotebookTool
+from hanzo_mcp.tools.jupyter.notebook_edit import NoteBookEditTool
+from hanzo_mcp.tools.jupyter.notebook_read import NotebookReadTool
 
 # Export all tool classes
 __all__ = [
-    "ReadNotebookTool",
-    "EditNotebookTool",
+    "NotebookReadTool",
+    "NoteBookEditTool",
     "get_jupyter_tools",
     "register_jupyter_tools",
 ]
 
+
 def get_read_only_jupyter_tools(
-    document_context: DocumentContext, permission_manager: PermissionManager
+    permission_manager: PermissionManager,
 ) -> list[BaseTool]:
     """Create instances of read only Jupyter notebook tools.
-    
+
     Args:
-        document_context: Document context for tracking file contents
         permission_manager: Permission manager for access control
-        
+
     Returns:
         List of Jupyter notebook tool instances
     """
     return [
-        ReadNotebookTool(document_context, permission_manager),
+        NotebookReadTool(permission_manager),
     ]
 
 
-def get_jupyter_tools(
-    document_context: DocumentContext, permission_manager: PermissionManager
-) -> list[BaseTool]:
+def get_jupyter_tools(permission_manager: PermissionManager) -> list[BaseTool]:
     """Create instances of all Jupyter notebook tools.
-    
+
     Args:
-        document_context: Document context for tracking file contents
         permission_manager: Permission manager for access control
-        
+
     Returns:
         List of Jupyter notebook tool instances
     """
     return [
-        ReadNotebookTool(document_context, permission_manager),
-        EditNotebookTool(document_context, permission_manager),
+        NotebookReadTool(permission_manager),
+        NoteBookEditTool(permission_manager),
     ]
 
 
 def register_jupyter_tools(
     mcp_server: FastMCP,
-    document_context: DocumentContext,
     permission_manager: PermissionManager,
-    disable_write_tools: bool = False,
-) -> None:
-    """Register all Jupyter notebook tools with the MCP server.
-    
+    enabled_tools: dict[str, bool] | None = None,
+) -> list[BaseTool]:
+    """Register Jupyter notebook tools with the MCP server.
+
     Args:
         mcp_server: The FastMCP server instance
-        document_context: Document context for tracking file contents
         permission_manager: Permission manager for access control
-        disable_write_tools: Whether to disable write/edit tools (default: False)
+        enabled_tools: Dictionary of individual tool enable states (default: None)
+
+    Returns:
+        List of registered tools
     """
-    if disable_write_tools:
-        tools = get_read_only_jupyter_tools(document_context, permission_manager)
+    # Define tool mapping
+    tool_classes = {
+        "notebook_read": NotebookReadTool,
+        "notebook_edit": NoteBookEditTool,
+    }
+    
+    tools = []
+    
+    if enabled_tools:
+        # Use individual tool configuration
+        for tool_name, enabled in enabled_tools.items():
+            if enabled and tool_name in tool_classes:
+                tool_class = tool_classes[tool_name]
+                tools.append(tool_class(permission_manager))
     else:
-        tools = get_jupyter_tools(document_context, permission_manager)
+        # Use all tools (backward compatibility)
+        tools = get_jupyter_tools(permission_manager)
+    
     ToolRegistry.register_tools(mcp_server, tools)
+    return tools

hanzo_mcp/tools/jupyter/base.py

@@ -10,35 +10,36 @@ import re
 from pathlib import Path
 from typing import Any, final
 
-from mcp.server.fastmcp import Context as MCPContext
+from fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.base import FileSystemTool
 from hanzo_mcp.tools.common.context import ToolContext, create_tool_context
 
 
 # Pattern to match ANSI escape sequences
-ANSI_ESCAPE_PATTERN = re.compile(r'\x1B\[[0-9;]*[a-zA-Z]')
+ANSI_ESCAPE_PATTERN = re.compile(r"\x1B\[[0-9;]*[a-zA-Z]")
+
 
 # Function to clean ANSI escape codes from text
 def clean_ansi_escapes(text: str) -> str:
     """Remove ANSI escape sequences from text.
-    
+
     Args:
         text: Text containing ANSI escape sequences
-        
+
     Returns:
         Text with ANSI escape sequences removed
     """
-    return ANSI_ESCAPE_PATTERN.sub('', text)
+    return ANSI_ESCAPE_PATTERN.sub("", text)
 
 
 @final
 class NotebookOutputImage:
     """Representation of an image output in a notebook cell."""
-    
+
     def __init__(self, image_data: str, media_type: str):
         """Initialize a notebook output image.
-        
+
         Args:
             image_data: Base64-encoded image data
             media_type: Media type of the image (e.g., "image/png")
@@ -50,15 +51,15 @@ class NotebookOutputImage:
 @final
 class NotebookCellOutput:
     """Representation of an output from a notebook cell."""
-    
+
     def __init__(
-        self, 
+        self,
         output_type: str,
         text: str | None = None,
-        image: NotebookOutputImage | None = None
+        image: NotebookOutputImage | None = None,
     ):
         """Initialize a notebook cell output.
-        
+
         Args:
             output_type: Type of output
             text: Text output (if any)
@@ -72,7 +73,7 @@ class NotebookCellOutput:
 @final
 class NotebookCellSource:
     """Representation of a source cell from a notebook."""
-    
+
     def __init__(
         self,
         cell_index: int,
@@ -80,10 +81,10 @@ class NotebookCellSource:
         source: str,
         language: str,
         execution_count: int | None = None,
-        outputs: list[NotebookCellOutput] | None = None
+        outputs: list[NotebookCellOutput] | None = None,
     ):
         """Initialize a notebook cell source.
-        
+
         Args:
             cell_index: Index of the cell in the notebook
             cell_type: Type of cell (code or markdown)
@@ -100,81 +101,87 @@ class NotebookCellSource:
         self.outputs = outputs or []
 
 
-class JupyterBaseTool(FileSystemTool,ABC):
+class JupyterBaseTool(FileSystemTool, ABC):
     """Base class for Jupyter notebook tools.
-    
+
     Provides common functionality for working with Jupyter notebooks, including
     parsing, cell extraction, and output formatting.
     """
-    
+
     def create_tool_context(self, ctx: MCPContext) -> ToolContext:
         """Create a tool context with the tool name.
-        
+
         Args:
             ctx: MCP context
-            
+
         Returns:
             Tool context
         """
         tool_ctx = create_tool_context(ctx)
         return tool_ctx
-        
+
     def set_tool_context_info(self, tool_ctx: ToolContext) -> None:
         """Set the tool info on the context.
-        
+
         Args:
             tool_ctx: Tool context
         """
         tool_ctx.set_tool_info(self.name)
-        
-    async def parse_notebook(self, file_path: Path) -> tuple[dict[str, Any], list[NotebookCellSource]]:
+
+    async def parse_notebook(
+        self, file_path: Path
+    ) -> tuple[dict[str, Any], list[NotebookCellSource]]:
         """Parse a Jupyter notebook file.
-        
+
         Args:
             file_path: Path to the notebook file
-            
+
         Returns:
             Tuple of (notebook_data, processed_cells)
         """
         with open(file_path, "r", encoding="utf-8") as f:
             content = f.read()
             notebook = json.loads(content)
-        
+
         # Get notebook language
-        language = notebook.get("metadata", {}).get("language_info", {}).get("name", "python")
+        language = (
+            notebook.get("metadata", {}).get("language_info", {}).get("name", "python")
+        )
         cells = notebook.get("cells", [])
         processed_cells = []
 
         for i, cell in enumerate(cells):
             cell_type = cell.get("cell_type", "code")
-            
+
             # Skip if not code or markdown
             if cell_type not in ["code", "markdown"]:
                 continue
-                
+
             # Get source
             source = cell.get("source", "")
             if isinstance(source, list):
                 source = "".join(source)
-            
+
             # Get execution count for code cells
             execution_count = None
             if cell_type == "code":
                 execution_count = cell.get("execution_count")
-            
+
             # Process outputs for code cells
             outputs = []
             if cell_type == "code" and "outputs" in cell:
                 for output in cell["outputs"]:
                     output_type = output.get("output_type", "")
-                    
+
                     # Process different output types
                     if output_type == "stream":
                         text = output.get("text", "")
                         if isinstance(text, list):
                             text = "".join(text)
-                        outputs.append(NotebookCellOutput(output_type="stream", text=text))
-                        
+                        outputs.append(
+                            NotebookCellOutput(output_type="stream", text=text)
+                        )
+
                     elif output_type in ["execute_result", "display_data"]:
                         # Process text output
                         text = None
@@ -184,46 +191,48 @@ class JupyterBaseTool(FileSystemTool,ABC):
                                 text = "".join(text_data)
                             else:
                                 text = text_data
-                        
+
                         # Process image output
                         image = None
                         if "data" in output:
                             if "image/png" in output["data"]:
                                 image = NotebookOutputImage(
                                     image_data=output["data"]["image/png"],
-                                    media_type="image/png"
+                                    media_type="image/png",
                                 )
                             elif "image/jpeg" in output["data"]:
                                 image = NotebookOutputImage(
                                     image_data=output["data"]["image/jpeg"],
-                                    media_type="image/jpeg"
+                                    media_type="image/jpeg",
                                 )
-                        
+
                         outputs.append(
                             NotebookCellOutput(
-                                output_type=output_type, 
-                                text=text,
-                                image=image
+                                output_type=output_type, text=text, image=image
                             )
                         )
-                        
+
                     elif output_type == "error":
                         # Format error traceback
                         ename = output.get("ename", "")
                         evalue = output.get("evalue", "")
                         traceback = output.get("traceback", [])
-                        
+
                         # Handle raw text strings and lists of strings
                         if isinstance(traceback, list):
                             # Clean ANSI escape codes and join the list but preserve the formatting
-                            clean_traceback = [clean_ansi_escapes(line) for line in traceback]
+                            clean_traceback = [
+                                clean_ansi_escapes(line) for line in traceback
+                            ]
                             traceback_text = "\n".join(clean_traceback)
                         else:
                             traceback_text = clean_ansi_escapes(str(traceback))
-                        
+
                         error_text = f"{ename}: {evalue}\n{traceback_text}"
-                        outputs.append(NotebookCellOutput(output_type="error", text=error_text))
-            
+                        outputs.append(
+                            NotebookCellOutput(output_type="error", text=error_text)
+                        )
+
             # Create cell object
             processed_cell = NotebookCellSource(
                 cell_index=i,
@@ -231,19 +240,19 @@ class JupyterBaseTool(FileSystemTool,ABC):
                 source=source,
                 language=language,
                 execution_count=execution_count,
-                outputs=outputs
+                outputs=outputs,
             )
-            
+
             processed_cells.append(processed_cell)
-            
+
         return notebook, processed_cells
-        
+
     def format_notebook_cells(self, cells: list[NotebookCellSource]) -> str:
         """Format notebook cells as a readable string.
-        
+
         Args:
             cells: List of processed notebook cells
-            
+
         Returns:
             Formatted string representation of the cells
         """
@@ -255,13 +264,13 @@ class JupyterBaseTool(FileSystemTool,ABC):
                 cell_header += f" (execution_count: {cell.execution_count})"
             if cell.cell_type == "code" and cell.language != "python":
                 cell_header += f" [{cell.language}]"
-            
+
             # Add cell to result
             result.append(f"{cell_header}:")
             result.append(f"```{cell.language if cell.cell_type == 'code' else ''}")
             result.append(cell.source)
             result.append("```")
-            
+
             # Add outputs if any
             if cell.outputs:
                 result.append("Outputs:")
@@ -278,7 +287,7 @@ class JupyterBaseTool(FileSystemTool,ABC):
                         result.append("```")
                     if output.image:
                         result.append(f"[Image output: {output.image.media_type}]")
-                
+
             result.append("")  # Empty line between cells
-            
+
         return "\n".join(result)