hanzo-mcp 0.7.6__py3-none-any.whl → 0.8.0__py3-none-any.whl

hanzo_mcp/tools/shell/streaming_command.py

@@ -1,25 +1,22 @@
 """Streaming command execution with disk-based logging and session management."""
 
-import asyncio
-import json
 import os
-import re
-import shutil
-import subprocess
-import tempfile
+import json
 import time
 import uuid
-from datetime import datetime, timedelta
+import shutil
+import asyncio
+import subprocess
+from typing import Any, Dict, List, Union, Optional
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from datetime import datetime, timedelta
 
-from hanzo_mcp.tools.common.base import BaseTool
 from hanzo_mcp.tools.shell.base_process import BaseProcessTool
 
 
 class StreamingCommandTool(BaseProcessTool):
     """Execute commands with disk-based streaming and session persistence.
-    
+
     Features:
     - All output streamed directly to disk (no memory usage)
     - Session-based organization of logs
@@ -27,46 +24,46 @@ class StreamingCommandTool(BaseProcessTool):
     - Forgiving parameter handling for AI usage
     - Automatic session detection from MCP context
     """
-    
+
     name = "streaming_command"
     description = "Run commands with disk-based output streaming and easy resumption"
-    
+
     # Base directory for all session data
     SESSION_BASE_DIR = Path.home() / ".hanzo" / "sessions"
-    
+
     # Chunk size for streaming (25k tokens ≈ 100KB)
     STREAM_CHUNK_SIZE = 100_000
-    
+
     # Session retention
     SESSION_RETENTION_DAYS = 30
-    
+
     def __init__(self):
         """Initialize the streaming command tool."""
         super().__init__()
         self.session_id = self._get_or_create_session()
         self.session_dir = self.SESSION_BASE_DIR / self.session_id
         self.session_dir.mkdir(parents=True, exist_ok=True)
-        
+
         # Create subdirectories
         self.commands_dir = self.session_dir / "commands"
         self.commands_dir.mkdir(exist_ok=True)
-        
+
         # Session metadata file
         self.session_meta_file = self.session_dir / "session.json"
         self._update_session_metadata()
-        
+
         # Cleanup old sessions on init
         self._cleanup_old_sessions()
-    
+
     def _get_or_create_session(self) -> str:
         """Get session ID from MCP context or create a new one.
-        
+
         Returns:
             Session ID string
         """
         # Try to get from environment (MCP might set this)
         session_id = os.environ.get("MCP_SESSION_ID")
-        
+
         if not session_id:
             # Try to get from Claude Desktop session marker
             claude_session = os.environ.get("CLAUDE_SESSION_ID")
@@ -76,9 +73,9 @@ class StreamingCommandTool(BaseProcessTool):
                 # Generate new session ID with timestamp
                 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                 session_id = f"session_{timestamp}_{uuid.uuid4().hex[:8]}"
-        
+
         return session_id
-    
+
     def _update_session_metadata(self) -> None:
         """Update session metadata file."""
         metadata = {
@@ -89,9 +86,9 @@ class StreamingCommandTool(BaseProcessTool):
                 "session_id": os.environ.get("MCP_SESSION_ID"),
                 "claude_session": os.environ.get("CLAUDE_SESSION_ID"),
                 "user": os.environ.get("USER"),
-            }
+            },
         }
-        
+
         # Merge with existing metadata if present
         if self.session_meta_file.exists():
             try:
@@ -100,37 +97,39 @@ class StreamingCommandTool(BaseProcessTool):
                     metadata["created"] = existing.get("created", metadata["created"])
             except Exception:
                 pass
-        
+
         with open(self.session_meta_file, "w") as f:
             json.dump(metadata, f, indent=2)
-    
+
     def _cleanup_old_sessions(self) -> None:
         """Remove sessions older than retention period."""
         if not self.SESSION_BASE_DIR.exists():
             return
-            
+
         cutoff = datetime.now() - timedelta(days=self.SESSION_RETENTION_DAYS)
-        
+
         for session_dir in self.SESSION_BASE_DIR.iterdir():
             if not session_dir.is_dir():
                 continue
-                
+
             meta_file = session_dir / "session.json"
             if meta_file.exists():
                 try:
                     with open(meta_file, "r") as f:
                         meta = json.load(f)
-                    last_accessed = datetime.fromisoformat(meta.get("last_accessed", ""))
+                    last_accessed = datetime.fromisoformat(
+                        meta.get("last_accessed", "")
+                    )
                     if last_accessed < cutoff:
                         shutil.rmtree(session_dir)
                 except Exception:
                     # If we can't read metadata, check directory mtime
                     if datetime.fromtimestamp(session_dir.stat().st_mtime) < cutoff:
                         shutil.rmtree(session_dir)
-    
+
     def _normalize_command_ref(self, ref: Union[str, int, None]) -> Optional[str]:
         """Normalize various command reference formats.
-        
+
         Args:
             ref: Command reference - can be:
                 - Full command ID (UUID)
@@ -138,15 +137,15 @@ class StreamingCommandTool(BaseProcessTool):
                 - Index number (1, 2, 3...)
                 - "last" or "latest"
                 - None
-                
+
         Returns:
             Full command ID or None
         """
         if not ref:
             return None
-            
+
         ref_str = str(ref).strip().lower()
-        
+
         # Handle special cases
         if ref_str in ["last", "latest", "recent"]:
             # Get most recent command
@@ -155,37 +154,39 @@ class StreamingCommandTool(BaseProcessTool):
                 return None
             latest = max(commands, key=lambda p: p.stat().st_mtime)
             return latest.parent.name
-        
+
         # Handle numeric index (1-based for user friendliness)
         if ref_str.isdigit():
             index = int(ref_str) - 1
-            commands = sorted(self.commands_dir.glob("*/metadata.json"), 
-                            key=lambda p: p.stat().st_mtime)
+            commands = sorted(
+                self.commands_dir.glob("*/metadata.json"),
+                key=lambda p: p.stat().st_mtime,
+            )
             if 0 <= index < len(commands):
                 return commands[index].parent.name
             return None
-        
+
         # Handle short ID (first 8 chars)
         if len(ref_str) >= 8:
             # Could be short or full ID
             for cmd_dir in self.commands_dir.iterdir():
                 if cmd_dir.name.startswith(ref_str):
                     return cmd_dir.name
-        
+
         return None
-    
+
     async def call(self, ctx: Any, **kwargs) -> Dict[str, Any]:
         """MCP tool entry point.
-        
+
         Args:
             ctx: MCP context
             **kwargs: Tool arguments
-            
+
         Returns:
             Tool result
         """
         return await self.run(**kwargs)
-    
+
     async def run(
         self,
         command: Optional[str] = None,
@@ -199,7 +200,7 @@ class StreamingCommandTool(BaseProcessTool):
         chunk_size: Optional[Union[int, str]] = None,
     ) -> Dict[str, Any]:
         """Execute or continue reading a command with maximum forgiveness.
-        
+
         Args:
             command/cmd: The command to execute (either works)
             working_dir/cwd: Directory to run in (either works)
@@ -207,7 +208,7 @@ class StreamingCommandTool(BaseProcessTool):
             continue_from/resume: Continue reading output from a command
             from_byte: Specific byte position to read from
             chunk_size: Custom chunk size for this read
-            
+
         Returns:
             Command output with metadata for easy continuation
         """
@@ -215,7 +216,7 @@ class StreamingCommandTool(BaseProcessTool):
         command = command or cmd
         working_dir = working_dir or cwd
         continue_from = continue_from or resume
-        
+
         # Convert string numbers to int
         if isinstance(timeout, str) and timeout.isdigit():
             timeout = int(timeout)
@@ -223,13 +224,13 @@ class StreamingCommandTool(BaseProcessTool):
             from_byte = int(from_byte)
         if isinstance(chunk_size, str) and chunk_size.isdigit():
             chunk_size = int(chunk_size)
-            
+
         chunk_size = chunk_size or self.STREAM_CHUNK_SIZE
-        
+
         # Handle continuation
         if continue_from:
             return await self._continue_reading(continue_from, from_byte, chunk_size)
-        
+
         # Need a command for new execution
         if not command:
             return {
@@ -237,10 +238,12 @@ class StreamingCommandTool(BaseProcessTool):
                 "hint": "To continue a previous command, use 'continue_from' with command ID or number.",
                 "recent_commands": await self._get_recent_commands(),
             }
-        
+
         # Execute new command
-        return await self._execute_new_command(command, working_dir, timeout, chunk_size)
-    
+        return await self._execute_new_command(
+            command, working_dir, timeout, chunk_size
+        )
+
     async def _execute_new_command(
         self,
         command: str,
@@ -253,12 +256,12 @@ class StreamingCommandTool(BaseProcessTool):
         cmd_id = str(uuid.uuid4())
         cmd_dir = self.commands_dir / cmd_id
         cmd_dir.mkdir()
-        
+
         # File paths
         output_file = cmd_dir / "output.log"
         error_file = cmd_dir / "error.log"
         metadata_file = cmd_dir / "metadata.json"
-        
+
         # Save metadata
         metadata = {
             "command_id": cmd_id,
@@ -268,10 +271,10 @@ class StreamingCommandTool(BaseProcessTool):
             "timeout": timeout,
             "status": "running",
         }
-        
+
         with open(metadata_file, "w") as f:
             json.dump(metadata, f, indent=2)
-        
+
         # Start process with output redirection
         try:
             process = await asyncio.create_subprocess_shell(
@@ -280,7 +283,7 @@ class StreamingCommandTool(BaseProcessTool):
                 stderr=asyncio.subprocess.PIPE,
                 cwd=working_dir,
             )
-            
+
             # Create tasks for streaming stdout and stderr to files
             async def stream_to_file(stream, file_path):
                 """Stream from async pipe to file."""
@@ -291,32 +294,38 @@ class StreamingCommandTool(BaseProcessTool):
                             break
                         f.write(chunk)
                         f.flush()  # Ensure immediate write
-            
+
             # Start streaming tasks
-            stdout_task = asyncio.create_task(stream_to_file(process.stdout, output_file))
-            stderr_task = asyncio.create_task(stream_to_file(process.stderr, error_file))
-            
+            stdout_task = asyncio.create_task(
+                stream_to_file(process.stdout, output_file)
+            )
+            stderr_task = asyncio.create_task(
+                stream_to_file(process.stderr, error_file)
+            )
+
             # Wait for initial output or timeout
             start_time = time.time()
-            initial_timeout = min(timeout or 5, 5)  # Wait max 5 seconds for initial output
-            
+            initial_timeout = min(
+                timeout or 5, 5
+            )  # Wait max 5 seconds for initial output
+
             while time.time() - start_time < initial_timeout:
                 if output_file.stat().st_size > 0 or error_file.stat().st_size > 0:
                     break
                 await asyncio.sleep(0.1)
-            
+
             # Read initial chunk
             output_content = ""
             error_content = ""
-            
+
             if output_file.exists() and output_file.stat().st_size > 0:
                 with open(output_file, "r", errors="replace") as f:
                     output_content = f.read(chunk_size)
-            
+
             if error_file.exists() and error_file.stat().st_size > 0:
                 with open(error_file, "r", errors="replace") as f:
                     error_content = f.read(1000)  # Just first 1KB of errors
-            
+
             # Check if process completed quickly
             try:
                 await asyncio.wait_for(process.wait(), timeout=0.1)
@@ -325,16 +334,16 @@ class StreamingCommandTool(BaseProcessTool):
             except asyncio.TimeoutError:
                 exit_code = None
                 status = "running"
-            
+
             # Update metadata
             metadata["status"] = status
             if exit_code is not None:
                 metadata["exit_code"] = exit_code
                 metadata["end_time"] = datetime.now().isoformat()
-            
+
             with open(metadata_file, "w") as f:
                 json.dump(metadata, f, indent=2)
-            
+
             # Build response
             result = {
                 "command_id": cmd_id,
@@ -345,13 +354,13 @@ class StreamingCommandTool(BaseProcessTool):
                 "bytes_read": len(output_content),
                 "session_path": str(cmd_dir),
             }
-            
+
             if error_content:
                 result["stderr"] = error_content
-            
+
             if exit_code is not None:
                 result["exit_code"] = exit_code
-            
+
             # Add continuation info if more output available
             total_size = output_file.stat().st_size
             if total_size > len(output_content) or status == "running":
@@ -366,30 +375,30 @@ class StreamingCommandTool(BaseProcessTool):
                     f"Command {'is still running' if status == 'running' else 'has more output'}. "
                     f"Use any of: {', '.join(result['continue_hints'])}"
                 )
-            
+
             # Ensure tasks complete
             if status == "completed":
                 await stdout_task
                 await stderr_task
-            
+
             return result
-            
+
         except Exception as e:
             # Update metadata with error
             metadata["status"] = "error"
             metadata["error"] = str(e)
             metadata["end_time"] = datetime.now().isoformat()
-            
+
             with open(metadata_file, "w") as f:
                 json.dump(metadata, f, indent=2)
-            
+
             return {
                 "error": str(e),
                 "command_id": cmd_id,
                 "short_id": cmd_id[:8],
                 "command": command,
             }
-    
+
     async def _continue_reading(
         self,
         ref: Union[str, int],
@@ -399,33 +408,33 @@ class StreamingCommandTool(BaseProcessTool):
         """Continue reading output from a previous command."""
         # Normalize reference
         cmd_id = self._normalize_command_ref(ref)
-        
+
         if not cmd_id:
             return {
                 "error": f"Command not found: {ref}",
                 "hint": "Use 'list' to see available commands",
                 "recent_commands": await self._get_recent_commands(),
             }
-        
+
         cmd_dir = self.commands_dir / cmd_id
         if not cmd_dir.exists():
             return {"error": f"Command directory not found: {cmd_id}"}
-        
+
         # Load metadata
         metadata_file = cmd_dir / "metadata.json"
         with open(metadata_file, "r") as f:
             metadata = json.load(f)
-        
+
         # Determine start position
         output_file = cmd_dir / "output.log"
         if not output_file.exists():
             return {"error": "No output file found"}
-        
+
         # If no from_byte specified, read from where we left off
         if from_byte is None:
             # Try to determine from previous reads (could track this)
             from_byte = 0  # For now, start from beginning if not specified
-        
+
         # Read chunk
         try:
             with open(output_file, "r", errors="replace") as f:
@@ -433,10 +442,10 @@ class StreamingCommandTool(BaseProcessTool):
                 content = f.read(chunk_size)
                 new_position = f.tell()
                 file_size = output_file.stat().st_size
-            
+
             # Check if process is still running
             status = metadata.get("status", "unknown")
-            
+
             # Build response
             result = {
                 "command_id": cmd_id,
@@ -449,13 +458,13 @@ class StreamingCommandTool(BaseProcessTool):
                 "read_to": new_position,
                 "total_bytes": file_size,
             }
-            
+
             # Add stderr if needed
             error_file = cmd_dir / "error.log"
             if error_file.exists() and error_file.stat().st_size > 0:
                 with open(error_file, "r", errors="replace") as f:
                     result["stderr"] = f.read(1000)
-            
+
             # Add continuation info
             if new_position < file_size or status == "running":
                 result["has_more"] = True
@@ -468,82 +477,88 @@ class StreamingCommandTool(BaseProcessTool):
                     f"{file_size - new_position} bytes remaining. "
                     f"Use: {result['continue_hints'][0]}"
                 )
-            
+
             return result
-            
+
         except Exception as e:
             return {"error": f"Error reading output: {str(e)}"}
-    
+
     async def _get_recent_commands(self, limit: int = 5) -> List[Dict[str, Any]]:
         """Get list of recent commands for hints."""
         commands = []
-        
-        for cmd_dir in sorted(self.commands_dir.iterdir(), 
-                             key=lambda p: p.stat().st_mtime, 
-                             reverse=True)[:limit]:
+
+        for cmd_dir in sorted(
+            self.commands_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True
+        )[:limit]:
             try:
                 with open(cmd_dir / "metadata.json", "r") as f:
                     meta = json.load(f)
-                    
+
                 output_size = 0
                 output_file = cmd_dir / "output.log"
                 if output_file.exists():
                     output_size = output_file.stat().st_size
-                    
-                commands.append({
-                    "id": meta["command_id"][:8],
-                    "command": meta["command"][:50] + "..." if len(meta["command"]) > 50 else meta["command"],
-                    "status": meta.get("status", "unknown"),
-                    "output_size": output_size,
-                    "time": meta.get("start_time", ""),
-                })
+
+                commands.append(
+                    {
+                        "id": meta["command_id"][:8],
+                        "command": (
+                            meta["command"][:50] + "..."
+                            if len(meta["command"]) > 50
+                            else meta["command"]
+                        ),
+                        "status": meta.get("status", "unknown"),
+                        "output_size": output_size,
+                        "time": meta.get("start_time", ""),
+                    }
+                )
             except Exception:
                 continue
-                
+
         return commands
-    
+
     async def list(self, limit: Optional[int] = 10) -> Dict[str, Any]:
         """List recent commands in this session.
-        
+
         Args:
             limit: Maximum number of commands to show
-            
+
         Returns:
             List of recent commands with details
         """
         commands = await self._get_recent_commands(limit or 10)
-        
+
         return {
             "session_id": self.session_id,
             "session_path": str(self.session_dir),
             "commands": commands,
             "hint": "Use continue_from='<id>' or resume='last' to read output",
         }
-    
+
     async def tail(
         self,
         ref: Optional[Union[str, int]] = None,
         lines: Optional[int] = 20,
     ) -> Dict[str, Any]:
         """Get the tail of a command's output (like 'tail -f').
-        
+
         Args:
             ref: Command reference (defaults to 'last')
             lines: Number of lines to show
-            
+
         Returns:
             Last N lines of output
         """
         ref = ref or "last"
         cmd_id = self._normalize_command_ref(ref)
-        
+
         if not cmd_id:
             return {"error": f"Command not found: {ref}"}
-        
+
         output_file = self.commands_dir / cmd_id / "output.log"
         if not output_file.exists():
             return {"error": "No output file found"}
-        
+
         try:
             # Use tail command for efficiency
             result = subprocess.run(
@@ -551,7 +566,7 @@ class StreamingCommandTool(BaseProcessTool):
                 capture_output=True,
                 text=True,
             )
-            
+
             return {
                 "command_id": cmd_id[:8],
                 "output": result.stdout,
@@ -559,7 +574,7 @@ class StreamingCommandTool(BaseProcessTool):
             }
         except Exception as e:
             return {"error": f"Error tailing output: {str(e)}"}
-    
+
     def get_params_schema(self) -> Dict[str, Any]:
         """Get parameter schema - very forgiving."""
         return {
@@ -604,24 +619,24 @@ class StreamingCommandTool(BaseProcessTool):
             },
             "required": [],  # No required fields for maximum forgiveness
         }
-    
+
     def get_command_args(self, command: str, **kwargs) -> List[str]:
         """Get the command arguments for subprocess.
-        
+
         Args:
             command: The command or script to run
             **kwargs: Additional arguments (not used for shell commands)
-            
+
         Returns:
             List of command arguments for subprocess
         """
         # For shell commands, we use shell=True, so return the command as-is
         return [command]
-    
+
     def get_tool_name(self) -> str:
         """Get the name of the tool being used.
-        
+
         Returns:
             Tool name
         """
-        return "streaming_command"
+        return "streaming_command"