mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
refactor(shell): merge dag into zsh for unified shell tool
- ZshTool now includes full DAG functionality (serial, parallel, graph)
- Single "zsh" command for all shell operations
- Usage: zsh("ls") or zsh(["cmd1", "cmd2"]) or zsh([...], parallel=True)
- DagTool kept in codebase for backwards compatibility
- Removed DagTool from default TOOLS list
- Bump hanzo-tools-shell to v0.3.1
This commit is contained in:
@@ -3,9 +3,9 @@
|
||||
Minimal, orthogonal shell execution.
|
||||
|
||||
Core tools:
|
||||
- dag: Execute commands/tools with DAG flow control (serial, parallel, graph)
|
||||
- zsh: Primary shell with full DAG support (serial, parallel, graph)
|
||||
- ps: Process management (list, kill, logs)
|
||||
- zsh: Zsh shell execution with auto-backgrounding
|
||||
- shell: Smart shell alias (zsh > bash fallback)
|
||||
|
||||
HTTP/Data tools:
|
||||
- curl: HTTP client without shell escaping issues
|
||||
@@ -16,6 +16,10 @@ Convenience tools:
|
||||
- npx: Node package execution with auto-backgrounding
|
||||
- uvx: Python package execution with auto-backgrounding
|
||||
- open: Open files/URLs in system apps
|
||||
|
||||
Note: 'dag' is now merged into 'zsh'. Use zsh(["cmd1", "cmd2"]) for serial,
|
||||
zsh([...], parallel=True) for parallel, and full DAG syntax is supported.
|
||||
DagTool kept for backwards compatibility.
|
||||
"""
|
||||
|
||||
from mcp.server import FastMCP
|
||||
@@ -48,7 +52,8 @@ from hanzo_tools.shell.base_process import (
|
||||
)
|
||||
|
||||
# Tools list for entry point discovery
|
||||
TOOLS = [DagTool, PsTool, ZshTool, ShellTool, NpxTool, UvxTool, OpenTool, CurlTool, JqTool, WgetTool]
|
||||
# Note: ZshTool now includes full DAG functionality. DagTool kept for backwards compatibility.
|
||||
TOOLS = [ZshTool, ShellTool, PsTool, NpxTool, UvxTool, OpenTool, CurlTool, JqTool, WgetTool]
|
||||
|
||||
__all__ = [
|
||||
# Base classes
|
||||
@@ -96,21 +101,20 @@ def get_shell_tools(
|
||||
|
||||
Args:
|
||||
permission_manager: Permission manager for access control
|
||||
all_tools: Dict of all registered tools (for dag tool invocations)
|
||||
all_tools: Dict of all registered tools (for zsh tool invocations)
|
||||
|
||||
Returns:
|
||||
List of shell tool instances
|
||||
"""
|
||||
# Create dag tool with access to other tools
|
||||
dag = create_dag_tool(tools=all_tools or {})
|
||||
# Create zsh tool with access to other tools (for DAG tool invocations)
|
||||
zsh = ZshTool(tools=all_tools or {})
|
||||
|
||||
# Set permission manager for convenience tools
|
||||
npx_tool.permission_manager = permission_manager
|
||||
uvx_tool.permission_manager = permission_manager
|
||||
|
||||
return [
|
||||
dag, # DAG execution for complex workflows
|
||||
zsh_tool, # Zsh shell execution (with auto-backgrounding)
|
||||
zsh, # Primary shell with DAG support
|
||||
ps_tool, # Process management
|
||||
npx_tool, # Node packages
|
||||
uvx_tool, # Python packages
|
||||
|
||||
@@ -1,230 +1,587 @@
|
||||
"""Zsh shell tool for command execution with enhanced features."""
|
||||
"""Unified Zsh shell tool with DAG execution support.
|
||||
|
||||
Combines simple shell execution with DAG (directed acyclic graph) semantics.
|
||||
Supports serial, parallel, and complex mixed execution patterns.
|
||||
|
||||
Usage:
|
||||
zsh("ls -la") # Single command
|
||||
zsh(["ls", "pwd", "git status"]) # Serial execution
|
||||
zsh(["npm i", "cargo build"], parallel=True) # Parallel
|
||||
zsh([
|
||||
"mkdir dist",
|
||||
{"parallel": ["cp a.txt dist/", "cp b.txt dist/"]},
|
||||
"zip -r out.zip dist/"
|
||||
]) # Mixed DAG
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
import shutil
|
||||
import asyncio
|
||||
import platform
|
||||
from typing import Optional, override
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Union, Optional, Annotated, override
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from dataclasses import field, dataclass
|
||||
|
||||
import aiofiles
|
||||
from pydantic import Field
|
||||
from mcp.server import FastMCP
|
||||
from mcp.server.fastmcp import Context as MCPContext
|
||||
|
||||
from hanzo_tools.core import auto_timeout
|
||||
from hanzo_tools.shell.base_process import BaseScriptTool
|
||||
from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
|
||||
from hanzo_tools.shell.base_process import ProcessManager
|
||||
|
||||
|
||||
class ZshTool(BaseScriptTool):
|
||||
"""Tool for running commands in Zsh shell with enhanced features."""
|
||||
class NodeStatus(Enum):
|
||||
"""Execution status of a DAG node."""
|
||||
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
SUCCESS = "success"
|
||||
FAILED = "failed"
|
||||
SKIPPED = "skipped"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DagResult:
|
||||
"""Result from a single DAG node execution."""
|
||||
|
||||
node_id: str
|
||||
command: str
|
||||
stdout: str
|
||||
stderr: str
|
||||
status: NodeStatus
|
||||
exit_code: int
|
||||
duration_ms: int = 0
|
||||
node_type: str = "shell"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DagNode:
|
||||
"""A node in the execution DAG."""
|
||||
|
||||
id: str
|
||||
command: Union[str, Dict[str, Any]]
|
||||
depends_on: List[str] = field(default_factory=list)
|
||||
status: NodeStatus = NodeStatus.PENDING
|
||||
result: Optional[DagResult] = None
|
||||
|
||||
|
||||
Command = Union[str, Dict[str, Any], List[Any]]
|
||||
|
||||
|
||||
class ZshTool(BaseTool):
|
||||
"""Zsh shell with full DAG execution support.
|
||||
|
||||
The primary shell tool for hanzo-mcp. Supports:
|
||||
- Simple command execution: zsh("ls -la")
|
||||
- Serial execution: zsh(["cmd1", "cmd2"])
|
||||
- Parallel execution: zsh(["cmd1", "cmd2"], parallel=True)
|
||||
- Mixed DAG patterns: zsh(["setup", {"parallel": ["a", "b"]}, "finish"])
|
||||
- Tool invocations: zsh([{"tool": "search", "input": {"pattern": "TODO"}}])
|
||||
- Named dependencies: zsh([{"id": "a", "run": "...", "after": ["b"]}])
|
||||
|
||||
Uses zsh by default, falls back to bash.
|
||||
Auto-backgrounds long-running processes.
|
||||
"""
|
||||
|
||||
name = "zsh"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
def __init__(self, tools: Optional[Dict[str, BaseTool]] = None):
|
||||
"""Initialize Zsh tool with optional tool registry for DAG tool invocations."""
|
||||
super().__init__()
|
||||
self.tools = tools or {}
|
||||
self._shell = self._resolve_shell()
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def zsh(
|
||||
ctx: MCPContext,
|
||||
command: str,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
timeout: Optional[int] = None,
|
||||
) -> str:
|
||||
return await tool_self.run(ctx, command=command, cwd=cwd, env=env, timeout=timeout)
|
||||
|
||||
@auto_timeout("zsh")
|
||||
async def call(self, ctx: MCPContext, **params) -> str:
|
||||
"""Call the tool with arguments."""
|
||||
return await self.run(
|
||||
ctx,
|
||||
command=params["command"],
|
||||
cwd=params.get("cwd"),
|
||||
env=params.get("env"),
|
||||
timeout=params.get("timeout"),
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return """Run commands in Zsh shell with enhanced features like better globbing and completion.
|
||||
|
||||
Zsh provides advanced features over bash:
|
||||
- Extended globbing patterns
|
||||
- Better tab completion
|
||||
- Array and associative array support
|
||||
- Powerful command line editing
|
||||
- Plugin ecosystem (oh-my-zsh, etc.)
|
||||
|
||||
Commands that run for more than 2 minutes will automatically continue in the background.
|
||||
|
||||
Usage:
|
||||
zsh "ls -la"
|
||||
zsh "echo $ZSH_VERSION"
|
||||
zsh "git status && git diff"
|
||||
zsh "npm run dev" --cwd ./frontend # Auto-backgrounds if needed"""
|
||||
|
||||
@override
|
||||
def get_interpreter(self) -> str:
|
||||
"""Get the zsh interpreter path."""
|
||||
def _resolve_shell(self) -> str:
|
||||
"""Resolve shell - prefer zsh, fallback to bash."""
|
||||
# Allow override
|
||||
force_shell = os.environ.get("HANZO_MCP_FORCE_SHELL")
|
||||
if force_shell:
|
||||
return force_shell
|
||||
|
||||
if platform.system() == "Windows":
|
||||
zsh_paths = [
|
||||
"C:\\Program Files\\Git\\usr\\bin\\zsh.exe",
|
||||
"C:\\cygwin64\\bin\\zsh.exe",
|
||||
"C:\\msys64\\usr\\bin\\zsh.exe",
|
||||
]
|
||||
for path in zsh_paths:
|
||||
if Path(path).exists():
|
||||
return path
|
||||
return "bash"
|
||||
shell_priority = ["zsh", "bash"]
|
||||
search_paths = [
|
||||
"/opt/homebrew/bin",
|
||||
"/usr/local/bin",
|
||||
"/bin",
|
||||
"/usr/bin",
|
||||
]
|
||||
|
||||
zsh_path = shutil.which("zsh")
|
||||
if zsh_path:
|
||||
return zsh_path
|
||||
for shell in shell_priority:
|
||||
for prefix in search_paths:
|
||||
full_path = f"{prefix}/{shell}"
|
||||
if os.path.isfile(full_path) and os.access(full_path, os.X_OK):
|
||||
return full_path
|
||||
found = shutil.which(shell)
|
||||
if found:
|
||||
return found
|
||||
|
||||
return "bash"
|
||||
return "sh"
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
shell_name = os.path.basename(self._shell)
|
||||
return f"""Zsh shell with DAG execution support (using: {shell_name}).
|
||||
|
||||
MODES:
|
||||
|
||||
Simple command:
|
||||
zsh("ls -la")
|
||||
zsh("git status && git diff")
|
||||
|
||||
Serial execution (default):
|
||||
zsh(["ls", "pwd", "git status"])
|
||||
|
||||
Parallel execution:
|
||||
zsh(["npm install", "cargo build"], parallel=True)
|
||||
|
||||
Mixed DAG:
|
||||
zsh([
|
||||
"mkdir -p dist",
|
||||
{{"parallel": ["cp a.txt dist/", "cp b.txt dist/"]}},
|
||||
"zip -r out.zip dist/"
|
||||
])
|
||||
|
||||
Tool invocations:
|
||||
zsh([{{"tool": "search", "input": {{"pattern": "TODO"}}}}])
|
||||
|
||||
Named with dependencies:
|
||||
zsh([
|
||||
{{"id": "build", "run": "make build"}},
|
||||
{{"id": "test", "run": "make test", "after": ["build"]}},
|
||||
])
|
||||
|
||||
OPTIONS:
|
||||
parallel: Run all commands concurrently (default: False)
|
||||
strict: Stop on first error (default: False)
|
||||
quiet: Suppress stdout (default: False)
|
||||
timeout: Per-command timeout in seconds (default: 120)
|
||||
cwd: Working directory
|
||||
env: Additional environment variables
|
||||
|
||||
AUTO-BACKGROUNDING: Commands exceeding timeout auto-background.
|
||||
Use ps tool to monitor: ps --logs <id>, ps --kill <id>"""
|
||||
|
||||
async def _run_shell(
|
||||
self,
|
||||
cmd: str,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
timeout: int = 120,
|
||||
) -> DagResult:
|
||||
"""Run a shell command with auto-backgrounding on timeout."""
|
||||
start_time = datetime.now()
|
||||
node_id = f"shell_{id(cmd)}"
|
||||
process_manager = ProcessManager()
|
||||
|
||||
try:
|
||||
run_env = os.environ.copy()
|
||||
if env:
|
||||
run_env.update(env)
|
||||
|
||||
process_id = f"zsh_{uuid.uuid4().hex[:8]}"
|
||||
log_file = await process_manager.create_log_file(process_id)
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
self._shell,
|
||||
"-c",
|
||||
cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
cwd=cwd,
|
||||
env=run_env,
|
||||
)
|
||||
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
||||
|
||||
duration = int((datetime.now() - start_time).total_seconds() * 1000)
|
||||
exit_code = proc.returncode or 0
|
||||
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=cmd,
|
||||
stdout=stdout.decode("utf-8", errors="replace"),
|
||||
stderr=stderr.decode("utf-8", errors="replace"),
|
||||
status=NodeStatus.SUCCESS if exit_code == 0 else NodeStatus.FAILED,
|
||||
exit_code=exit_code,
|
||||
duration_ms=duration,
|
||||
node_type="shell",
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
duration = int((datetime.now() - start_time).total_seconds() * 1000)
|
||||
|
||||
async with aiofiles.open(log_file, "w") as f:
|
||||
await f.write(f"[zsh] Command backgrounded after {timeout}s timeout\n")
|
||||
await f.write(f"[zsh] Command: {cmd}\n")
|
||||
await f.write(f"[zsh] PID: {proc.pid}\n")
|
||||
await f.write(f"[zsh] Started: {start_time.isoformat()}\n")
|
||||
await f.write("-" * 40 + "\n")
|
||||
|
||||
process_manager.add_process(process_id, proc, str(log_file))
|
||||
asyncio.create_task(self._capture_background_output(proc, log_file))
|
||||
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=cmd,
|
||||
stdout=f"[backgrounded] Process {process_id} (PID {proc.pid}) running in background.\n"
|
||||
f"Use: ps --logs {process_id} # view output\n"
|
||||
f"Use: ps --kill {process_id} # stop process",
|
||||
stderr="",
|
||||
status=NodeStatus.SUCCESS,
|
||||
exit_code=0,
|
||||
duration_ms=duration,
|
||||
node_type="shell",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
duration = int((datetime.now() - start_time).total_seconds() * 1000)
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=cmd,
|
||||
stdout="",
|
||||
stderr=str(e),
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=duration,
|
||||
node_type="shell",
|
||||
)
|
||||
|
||||
async def _capture_background_output(self, proc: asyncio.subprocess.Process, log_file: Path) -> None:
|
||||
"""Capture output from backgrounded process to log file."""
|
||||
try:
|
||||
async with aiofiles.open(log_file, "a") as f:
|
||||
|
||||
async def read_stream(stream, prefix: str):
|
||||
if stream:
|
||||
while True:
|
||||
line = await stream.readline()
|
||||
if not line:
|
||||
break
|
||||
await f.write(f"{prefix}{line.decode('utf-8', errors='replace')}")
|
||||
await f.flush()
|
||||
|
||||
await asyncio.gather(
|
||||
read_stream(proc.stdout, ""),
|
||||
read_stream(proc.stderr, "[stderr] "),
|
||||
)
|
||||
|
||||
await proc.wait()
|
||||
await f.write(f"\n[zsh] Process exited with code {proc.returncode}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _run_tool(self, tool_name: str, tool_input: Dict[str, Any], ctx: MCPContext) -> DagResult:
|
||||
"""Run an MCP tool invocation."""
|
||||
start_time = datetime.now()
|
||||
node_id = f"tool_{tool_name}"
|
||||
|
||||
if tool_name not in self.tools:
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=f"tool:{tool_name}",
|
||||
stdout="",
|
||||
stderr=f"Tool '{tool_name}' not found. Available: {list(self.tools.keys())}",
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=0,
|
||||
node_type="tool",
|
||||
)
|
||||
|
||||
try:
|
||||
tool = self.tools[tool_name]
|
||||
result = await tool.call(ctx, **tool_input)
|
||||
duration = int((datetime.now() - start_time).total_seconds() * 1000)
|
||||
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=f"tool:{tool_name}",
|
||||
stdout=str(result) if result else "",
|
||||
stderr="",
|
||||
status=NodeStatus.SUCCESS,
|
||||
exit_code=0,
|
||||
duration_ms=duration,
|
||||
node_type="tool",
|
||||
)
|
||||
except Exception as e:
|
||||
duration = int((datetime.now() - start_time).total_seconds() * 1000)
|
||||
return DagResult(
|
||||
node_id=node_id,
|
||||
command=f"tool:{tool_name}",
|
||||
stdout="",
|
||||
stderr=str(e),
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=duration,
|
||||
node_type="tool",
|
||||
)
|
||||
|
||||
async def _execute_node(
|
||||
self,
|
||||
cmd: Command,
|
||||
ctx: MCPContext,
|
||||
cwd: Optional[str],
|
||||
env: Optional[Dict[str, str]],
|
||||
timeout: int,
|
||||
) -> DagResult:
|
||||
"""Execute a single DAG node."""
|
||||
|
||||
if isinstance(cmd, str):
|
||||
return await self._run_shell(cmd, cwd, env, timeout)
|
||||
|
||||
if isinstance(cmd, dict):
|
||||
if "tool" in cmd:
|
||||
return await self._run_tool(cmd["tool"], cmd.get("input", {}), ctx)
|
||||
|
||||
if "parallel" in cmd:
|
||||
parallel_cmds = cmd["parallel"]
|
||||
tasks = [self._execute_node(c, ctx, cwd, env, timeout) for c in parallel_cmds]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
combined_stdout = []
|
||||
combined_stderr = []
|
||||
all_success = True
|
||||
total_duration = 0
|
||||
|
||||
for i, r in enumerate(results):
|
||||
if isinstance(r, Exception):
|
||||
combined_stderr.append(f"[{i}] Error: {r}")
|
||||
all_success = False
|
||||
else:
|
||||
if r.stdout:
|
||||
combined_stdout.append(f"[{i}] {r.stdout.rstrip()}")
|
||||
if r.stderr:
|
||||
combined_stderr.append(f"[{i}] {r.stderr.rstrip()}")
|
||||
if r.status != NodeStatus.SUCCESS:
|
||||
all_success = False
|
||||
total_duration = max(total_duration, r.duration_ms)
|
||||
|
||||
return DagResult(
|
||||
node_id=f"parallel_{len(parallel_cmds)}",
|
||||
command=f"parallel[{len(parallel_cmds)} tasks]",
|
||||
stdout="\n".join(combined_stdout),
|
||||
stderr="\n".join(combined_stderr),
|
||||
status=NodeStatus.SUCCESS if all_success else NodeStatus.FAILED,
|
||||
exit_code=0 if all_success else 1,
|
||||
duration_ms=total_duration,
|
||||
node_type="parallel",
|
||||
)
|
||||
|
||||
if "run" in cmd:
|
||||
run_cmd = cmd["run"]
|
||||
return await self._execute_node(run_cmd, ctx, cwd, env, timeout)
|
||||
|
||||
return DagResult(
|
||||
node_id="unknown",
|
||||
command=str(cmd),
|
||||
stdout="",
|
||||
stderr=f"Unknown command format: {cmd}",
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=0,
|
||||
node_type="unknown",
|
||||
)
|
||||
|
||||
return DagResult(
|
||||
node_id="unknown",
|
||||
command=str(cmd),
|
||||
stdout="",
|
||||
stderr=f"Unknown command type: {type(cmd).__name__}",
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=0,
|
||||
node_type="unknown",
|
||||
)
|
||||
|
||||
def _format_output(self, results: List[DagResult], quiet: bool = False) -> str:
|
||||
"""Format DAG execution results."""
|
||||
output_parts = []
|
||||
total_duration = 0
|
||||
failed_count = 0
|
||||
|
||||
for r in results:
|
||||
total_duration += r.duration_ms
|
||||
|
||||
if r.status == NodeStatus.FAILED:
|
||||
failed_count += 1
|
||||
|
||||
if r.stdout and not quiet:
|
||||
output_parts.append(r.stdout.rstrip())
|
||||
|
||||
if r.stderr:
|
||||
output_parts.append(f"[stderr] {r.stderr.rstrip()}")
|
||||
|
||||
if len(results) > 1:
|
||||
status = "✓" if failed_count == 0 else f"✗ ({failed_count} failed)"
|
||||
output_parts.append(f"\n[zsh] {len(results)} nodes, {total_duration}ms, {status}")
|
||||
|
||||
return "\n".join(output_parts) if output_parts else "(no output)"
|
||||
|
||||
@override
|
||||
def get_script_flags(self) -> list[str]:
|
||||
"""Get interpreter flags."""
|
||||
if platform.system() == "Windows" and self.get_interpreter().endswith(".exe"):
|
||||
return ["-c"]
|
||||
return ["-c"]
|
||||
|
||||
@override
|
||||
def get_tool_name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "zsh"
|
||||
|
||||
@override
|
||||
async def run(
|
||||
@auto_timeout("zsh")
|
||||
async def call(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
command: str,
|
||||
command: Optional[str] = None,
|
||||
commands: Optional[List[Command]] = None,
|
||||
parallel: bool = False,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
timeout: Optional[int] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
timeout: int = 120,
|
||||
strict: bool = False,
|
||||
quiet: bool = False,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""Run a zsh command with auto-backgrounding."""
|
||||
if not shutil.which("zsh") and platform.system() != "Windows":
|
||||
return "Error: Zsh is not installed. Please install zsh first."
|
||||
"""Execute shell commands with optional DAG semantics.
|
||||
|
||||
Args:
|
||||
command: Single command string (simple mode)
|
||||
commands: List of commands for DAG mode
|
||||
parallel: Run all commands concurrently
|
||||
cwd: Working directory
|
||||
env: Additional environment variables
|
||||
timeout: Per-command timeout in seconds
|
||||
strict: Stop on first error
|
||||
quiet: Suppress stdout
|
||||
"""
|
||||
tool_ctx = create_tool_context(ctx)
|
||||
await tool_ctx.set_tool_info(self.name)
|
||||
|
||||
work_dir = Path(cwd).resolve() if cwd else Path.cwd()
|
||||
# Single command mode
|
||||
if command and not commands:
|
||||
result = await self._run_shell(command, cwd, env, timeout)
|
||||
if result.status == NodeStatus.SUCCESS:
|
||||
return result.stdout if result.stdout else "(no output)"
|
||||
return f"{result.stdout}\n[stderr] {result.stderr}" if result.stderr else result.stdout
|
||||
|
||||
output = await self.execute_sync(command, cwd=work_dir, env=env, timeout=timeout)
|
||||
return output if output else "Command completed successfully (no output)"
|
||||
# DAG mode
|
||||
if not commands:
|
||||
return "Error: Provide either 'command' (string) or 'commands' (list)"
|
||||
|
||||
results: List[DagResult] = []
|
||||
|
||||
if parallel:
|
||||
tasks = [self._execute_node(cmd, ctx, cwd, env, timeout) for cmd in commands]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
results = [
|
||||
r
|
||||
if not isinstance(r, Exception)
|
||||
else DagResult(
|
||||
node_id=f"error_{i}",
|
||||
command=str(commands[i]),
|
||||
stdout="",
|
||||
stderr=str(r),
|
||||
status=NodeStatus.FAILED,
|
||||
exit_code=1,
|
||||
duration_ms=0,
|
||||
node_type="error",
|
||||
)
|
||||
for i, r in enumerate(results)
|
||||
]
|
||||
else:
|
||||
for cmd in commands:
|
||||
result = await self._execute_node(cmd, ctx, cwd, env, timeout)
|
||||
results.append(result)
|
||||
|
||||
if strict and result.status == NodeStatus.FAILED:
|
||||
break
|
||||
|
||||
return self._format_output(results, quiet)
|
||||
|
||||
@override
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register zsh tool with MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def zsh(
|
||||
command: Annotated[
|
||||
Optional[str], Field(description="Single command to execute", default=None)
|
||||
] = None,
|
||||
commands: Annotated[
|
||||
Optional[List[Any]], Field(description="List of commands for DAG execution", default=None)
|
||||
] = None,
|
||||
parallel: Annotated[bool, Field(description="Run all commands in parallel", default=False)] = False,
|
||||
cwd: Annotated[Optional[str], Field(description="Working directory", default=None)] = None,
|
||||
env: Annotated[Optional[Dict[str, str]], Field(description="Environment variables", default=None)] = None,
|
||||
timeout: Annotated[int, Field(description="Timeout per command (seconds)", default=120)] = 120,
|
||||
strict: Annotated[bool, Field(description="Stop on first error", default=False)] = False,
|
||||
quiet: Annotated[bool, Field(description="Suppress stdout", default=False)] = False,
|
||||
ctx: MCPContext = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
command=command,
|
||||
commands=commands,
|
||||
parallel=parallel,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
timeout=timeout,
|
||||
strict=strict,
|
||||
quiet=quiet,
|
||||
)
|
||||
|
||||
|
||||
class ShellTool(BaseScriptTool):
|
||||
"""Smart shell tool that uses the best available shell (zsh > bash)."""
|
||||
class ShellTool(ZshTool):
|
||||
"""Alias for ZshTool - smart shell that uses zsh > bash."""
|
||||
|
||||
name = "shell"
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize and detect the best shell."""
|
||||
super().__init__()
|
||||
self._best_shell = self._detect_best_shell()
|
||||
|
||||
def _detect_best_shell(self) -> str:
|
||||
"""Detect the best available shell."""
|
||||
force_shell = os.environ.get("HANZO_MCP_FORCE_SHELL")
|
||||
if force_shell:
|
||||
return force_shell
|
||||
|
||||
if shutil.which("zsh"):
|
||||
if (Path.home() / ".zshrc").exists():
|
||||
return "zsh"
|
||||
|
||||
user_shell = os.environ.get("SHELL", "")
|
||||
if user_shell and Path(user_shell).exists():
|
||||
return user_shell
|
||||
|
||||
return "bash"
|
||||
|
||||
def register(self, server: FastMCP) -> None:
|
||||
"""Register the tool with the MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@server.tool(name=self.name, description=self.description)
|
||||
async def shell(
|
||||
ctx: MCPContext,
|
||||
command: str,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
timeout: Optional[int] = None,
|
||||
) -> str:
|
||||
return await tool_self.run(ctx, command=command, cwd=cwd, env=env, timeout=timeout)
|
||||
|
||||
@auto_timeout("shell")
|
||||
async def call(self, ctx: MCPContext, **params) -> str:
|
||||
"""Call the tool with arguments."""
|
||||
return await self.run(
|
||||
ctx,
|
||||
command=params["command"],
|
||||
cwd=params.get("cwd"),
|
||||
env=params.get("env"),
|
||||
timeout=params.get("timeout"),
|
||||
)
|
||||
|
||||
@property
|
||||
@override
|
||||
def description(self) -> str:
|
||||
"""Get the tool description."""
|
||||
return f"""Run shell commands using the best available shell (currently: {os.path.basename(self._best_shell)}).
|
||||
shell_name = os.path.basename(self._shell)
|
||||
return f"""Run shell commands using the best available shell (currently: {shell_name}).
|
||||
|
||||
Automatically selects:
|
||||
- Zsh if available (with .zshrc)
|
||||
- User's preferred shell ($SHELL)
|
||||
- Bash as fallback
|
||||
|
||||
Commands that run for more than 2 minutes will automatically continue in the background.
|
||||
Supports all DAG modes (see: zsh --help).
|
||||
|
||||
Usage:
|
||||
shell "ls -la"
|
||||
shell "echo $SHELL" # Shows which shell is being used
|
||||
shell "git status && git diff"
|
||||
shell "npm run dev" --cwd ./frontend # Auto-backgrounds if needed"""
|
||||
shell("ls -la")
|
||||
shell(["cmd1", "cmd2"]) # Serial
|
||||
shell(["a", "b"], parallel=True) # Parallel"""
|
||||
|
||||
@override
|
||||
def get_interpreter(self) -> str:
|
||||
"""Get the best shell interpreter."""
|
||||
return self._best_shell
|
||||
def register(self, mcp_server: FastMCP) -> None:
|
||||
"""Register shell tool with MCP server."""
|
||||
tool_self = self
|
||||
|
||||
@override
|
||||
def get_script_flags(self) -> list[str]:
|
||||
"""Get interpreter flags."""
|
||||
if platform.system() == "Windows":
|
||||
return ["/c"] if self._best_shell == "cmd.exe" else ["-c"]
|
||||
return ["-c"]
|
||||
|
||||
@override
|
||||
def get_tool_name(self) -> str:
|
||||
"""Get the tool name."""
|
||||
return "shell"
|
||||
|
||||
@override
|
||||
async def run(
|
||||
self,
|
||||
ctx: MCPContext,
|
||||
command: str,
|
||||
cwd: Optional[str] = None,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
timeout: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Run a shell command with auto-backgrounding."""
|
||||
work_dir = Path(cwd).resolve() if cwd else Path.cwd()
|
||||
shell_name = os.path.basename(self._best_shell)
|
||||
|
||||
output = await self.execute_sync(command, cwd=work_dir, env=env, timeout=timeout)
|
||||
|
||||
if output:
|
||||
return output
|
||||
else:
|
||||
return f"Command completed successfully in {shell_name} (no output)"
|
||||
@mcp_server.tool(name=self.name, description=self.description)
|
||||
async def shell(
|
||||
command: Annotated[
|
||||
Optional[str], Field(description="Single command to execute", default=None)
|
||||
] = None,
|
||||
commands: Annotated[
|
||||
Optional[List[Any]], Field(description="List of commands for DAG execution", default=None)
|
||||
] = None,
|
||||
parallel: Annotated[bool, Field(description="Run all commands in parallel", default=False)] = False,
|
||||
cwd: Annotated[Optional[str], Field(description="Working directory", default=None)] = None,
|
||||
env: Annotated[Optional[Dict[str, str]], Field(description="Environment variables", default=None)] = None,
|
||||
timeout: Annotated[int, Field(description="Timeout per command (seconds)", default=120)] = 120,
|
||||
strict: Annotated[bool, Field(description="Stop on first error", default=False)] = False,
|
||||
quiet: Annotated[bool, Field(description="Suppress stdout", default=False)] = False,
|
||||
ctx: MCPContext = None,
|
||||
) -> str:
|
||||
return await tool_self.call(
|
||||
ctx,
|
||||
command=command,
|
||||
commands=commands,
|
||||
parallel=parallel,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
timeout=timeout,
|
||||
strict=strict,
|
||||
quiet=quiet,
|
||||
)
|
||||
|
||||
|
||||
# Create tool instances
|
||||
|
||||
@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hanzo-tools-shell"
|
||||
version = "0.3.0"
|
||||
description = "Shell/command execution tools - dag, ps, zsh, npx, uvx, open, curl, jq, wget"
|
||||
version = "0.3.1"
|
||||
description = "Shell/command execution tools - zsh (with DAG), ps, npx, uvx, open, curl, jq, wget"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
license = { text = "MIT" }
|
||||
|
||||
Reference in New Issue
Block a user