feat(tools): add Rust<>Python parity for plan, patch, wait tools

- plan: add update/get/clear actions for step tracking (Rust update_plan parity)
- fs: add patch action with Rust grammar format parser (*** Begin Patch)
- proc: add wait action for background polling, array command support

Aligns Python MCP tools with Rust dev tool schemas for cross-language compatibility.
This commit is contained in:
2026-02-06 07:25:02 -08:00
parent ab926d31bd
commit b35a5ca4e0
3 changed files with 505 additions and 3 deletions
@@ -13,6 +13,7 @@ Following Unix philosophy: one tool for the Bytes + Paths axis.
"""
import os
import re
import json
import fnmatch
import hashlib
@@ -20,6 +21,8 @@ import aiofiles
from typing import Any, Optional, ClassVar
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass, field
from enum import Enum
from mcp.server import FastMCP
from mcp.server.fastmcp import Context as MCPContext
@@ -36,6 +39,119 @@ from hanzo_tools.core import (
)
class PatchOp(str, Enum):
"""Patch operation type (Rust parity)."""
ADD = "add"
UPDATE = "update"
DELETE = "delete"
@dataclass
class PatchHunk:
"""A single hunk in a patch (Rust parity)."""
context: str = "" # @@ header context
old_lines: list[str] = field(default_factory=list)
new_lines: list[str] = field(default_factory=list)
@dataclass
class PatchFile:
"""A single file operation in a patch (Rust parity)."""
op: PatchOp
path: str
hunks: list[PatchHunk] = field(default_factory=list)
content: str = "" # For ADD operations
def parse_patch_format(patch_text: str) -> list[PatchFile]:
"""Parse Rust-style patch format.
Format:
*** Begin Patch
*** Add File: path
+content lines
*** Update File: path
@@ context header
-old line
+new line
*** Delete File: path
*** End Patch
"""
files = []
current_file: PatchFile | None = None
current_hunk: PatchHunk | None = None
lines = patch_text.strip().splitlines()
for line in lines:
# Skip begin/end markers
if line.strip() in ("*** Begin Patch", "*** End Patch"):
continue
# File operations
if line.startswith("*** Add File:"):
if current_file:
if current_hunk:
current_file.hunks.append(current_hunk)
files.append(current_file)
path = line.replace("*** Add File:", "").strip()
current_file = PatchFile(op=PatchOp.ADD, path=path)
current_hunk = None
elif line.startswith("*** Update File:"):
if current_file:
if current_hunk:
current_file.hunks.append(current_hunk)
files.append(current_file)
path = line.replace("*** Update File:", "").strip()
current_file = PatchFile(op=PatchOp.UPDATE, path=path)
current_hunk = None
elif line.startswith("*** Delete File:"):
if current_file:
if current_hunk:
current_file.hunks.append(current_hunk)
files.append(current_file)
path = line.replace("*** Delete File:", "").strip()
current_file = PatchFile(op=PatchOp.DELETE, path=path)
current_hunk = None
files.append(current_file)
current_file = None
# Hunk header
elif line.startswith("@@"):
if current_file and current_hunk:
current_file.hunks.append(current_hunk)
context = line.strip()
current_hunk = PatchHunk(context=context)
# Content lines
elif current_file:
if current_file.op == PatchOp.ADD:
# For ADD, lines starting with + are content
if line.startswith("+"):
current_file.content += line[1:] + "\n"
else:
current_file.content += line + "\n"
elif current_hunk:
if line.startswith("-"):
current_hunk.old_lines.append(line[1:])
elif line.startswith("+"):
current_hunk.new_lines.append(line[1:])
elif line.startswith(" "):
# Context line - appears in both old and new
current_hunk.old_lines.append(line[1:])
current_hunk.new_lines.append(line[1:])
# Don't forget the last file
if current_file:
if current_hunk:
current_file.hunks.append(current_hunk)
files.append(current_file)
return files
class FsTool(BaseTool):
"""Unified filesystem tool (HIP-0300).
@@ -73,11 +189,13 @@ Actions:
- stat: File metadata including hash
- list: Directory listing
- apply_patch: Edit with base_hash precondition
- patch: Apply Rust-style patch format (Rust parity)
- search_text: Text search
- mkdir: Create directory
- rm: Remove (requires confirm=true)
IMPORTANT: apply_patch is the ONLY way to edit existing files.
patch supports Rust grammar format: *** Begin Patch / *** Update File: / @@ / -old +new
"""
def _validate_path(self, path: str) -> Path:
@@ -383,6 +501,136 @@ IMPORTANT: apply_patch is the ONLY way to edit existing files.
"previous_hash": current_hash,
}
@self.action("patch", "Apply Rust-style patch format (Rust parity)")
async def patch(
ctx: MCPContext,
input: str,
) -> dict:
"""Apply a patch in Rust grammar format.
This provides parity with the Rust apply_patch tool.
Format:
*** Begin Patch
*** Add File: path/to/new/file.py
+new file content
+line 2
*** Update File: path/to/existing/file.py
@@ context header
-old line to remove
+new line to add
*** Delete File: path/to/delete.py
*** End Patch
Args:
input: The patch text in Rust grammar format
Returns:
Results for each file operation
"""
if not input or not input.strip():
raise InvalidParamsError("Patch input is required", param="input")
# Parse the patch
try:
patch_files = parse_patch_format(input)
except Exception as e:
raise InvalidParamsError(f"Invalid patch format: {e}", param="input")
if not patch_files:
raise InvalidParamsError("No file operations found in patch", param="input")
results = []
for pf in patch_files:
path = Path(pf.path)
# Make path absolute if relative
if not path.is_absolute():
# Use current working directory
path = Path.cwd() / path
self._check_permission(str(path))
try:
if pf.op == PatchOp.ADD:
# Create new file
if path.exists():
raise ConflictError(f"File already exists: {path}")
path.parent.mkdir(parents=True, exist_ok=True)
async with aiofiles.open(path, "w", encoding="utf-8") as f:
await f.write(pf.content)
results.append({
"op": "add",
"path": str(path),
"hash": self._compute_hash(pf.content),
"success": True,
})
elif pf.op == PatchOp.UPDATE:
# Update existing file
if not path.exists():
raise NotFoundError(f"File not found: {path}")
async with aiofiles.open(path, "r", encoding="utf-8") as f:
content = await f.read()
# Apply each hunk
for hunk in pf.hunks:
old_text = "\n".join(hunk.old_lines)
new_text = "\n".join(hunk.new_lines)
if old_text and old_text in content:
content = content.replace(old_text, new_text, 1)
elif not old_text and new_text:
# Pure addition - append to end
content += "\n" + new_text
async with aiofiles.open(path, "w", encoding="utf-8") as f:
await f.write(content)
results.append({
"op": "update",
"path": str(path),
"hash": self._compute_hash(content),
"hunks_applied": len(pf.hunks),
"success": True,
})
elif pf.op == PatchOp.DELETE:
# Delete file
if not path.exists():
results.append({
"op": "delete",
"path": str(path),
"success": True,
"message": "File already deleted",
})
else:
path.unlink()
results.append({
"op": "delete",
"path": str(path),
"success": True,
})
except Exception as e:
results.append({
"op": pf.op.value,
"path": str(path),
"success": False,
"error": str(e),
})
return {
"results": results,
"total": len(results),
"success": all(r.get("success") for r in results),
}
@self.action("search_text", "Search file contents")
async def search_text(
ctx: MCPContext,
@@ -18,6 +18,7 @@ import os
import json
import hashlib
import re
from datetime import datetime, timezone
from typing import Any, ClassVar, Literal
from dataclasses import dataclass, field
from enum import Enum
@@ -40,6 +41,29 @@ class EffectClass(str, Enum):
NONDETERMINISTIC = "nondeterministic"
class StepStatus(str, Enum):
"""Status for plan steps (Rust parity)."""
PENDING = "pending"
IN_PROGRESS = "in_progress"
COMPLETED = "completed"
@dataclass
class TrackedStep:
"""A tracked plan step with status (Rust parity)."""
step: str
status: StepStatus = StepStatus.PENDING
@dataclass
class TrackedPlan:
"""A tracked plan with name and steps (Rust parity)."""
name: str | None = None
steps: list[TrackedStep] = field(default_factory=list)
created_at: str | None = None
updated_at: str | None = None
class Scope(str, Enum):
"""Scope lattice for operators."""
SPAN = "span"
@@ -184,16 +208,21 @@ class PlanTool(BaseTool):
- intent: Parse NL → IntentIR
- route: IntentIR → Plan
- compose: Plan → ExecGraph
- update: Track plan with step status (Rust parity)
- execute: Run with audit (optional)
This is the "permissive input → strict output" layer.
"""
name: ClassVar[str] = "plan"
VERSION: ClassVar[str] = "0.1.0"
VERSION: ClassVar[str] = "0.2.0"
# In-memory tracked plan (matches Rust update_plan behavior)
_tracked_plan: TrackedPlan | None = None
def __init__(self):
super().__init__()
self._tracked_plan = None
self._register_plan_actions()
@property
@@ -201,6 +230,8 @@ class PlanTool(BaseTool):
return """Unified plan orchestration tool (HIP-0300).
Actions:
- update: Track plan with step status (Rust parity - name, steps with status)
- get: Get current tracked plan
- intent: Parse natural language → IntentIR
- route: Map IntentIR → Plan (canonical operator chain)
- compose: Plan → ExecGraph (execution-ready DAG)
@@ -324,6 +355,122 @@ Turns permissive natural language input into strict canonical operator chains.
def _register_plan_actions(self):
"""Register all plan actions."""
@self.action("update", "Update tracked plan with steps and status (Rust parity)")
async def update(
ctx: MCPContext,
name: str | None = None,
plan: list[dict] | None = None,
) -> dict:
"""Update the tracked plan with steps and their status.
This matches the Rust update_plan tool schema exactly.
Args:
name: Optional plan title (2-5 words)
plan: List of steps, each with:
- step: Step description (required)
- status: "pending" | "in_progress" | "completed"
Returns:
{"message": "Plan updated", "plan": {...}}
Example:
plan(action="update", name="Fix auth bug", plan=[
{"step": "Identify the root cause", "status": "completed"},
{"step": "Write failing test", "status": "in_progress"},
{"step": "Implement fix", "status": "pending"},
{"step": "Verify all tests pass", "status": "pending"}
])
"""
now = datetime.now(timezone.utc).isoformat()
# Parse steps
steps = []
if plan:
for item in plan:
step_text = item.get("step", "")
status_str = item.get("status", "pending")
try:
status = StepStatus(status_str)
except ValueError:
status = StepStatus.PENDING
steps.append(TrackedStep(step=step_text, status=status))
# Update or create tracked plan
if self._tracked_plan is None:
self._tracked_plan = TrackedPlan(
name=name,
steps=steps,
created_at=now,
updated_at=now,
)
else:
if name is not None:
self._tracked_plan.name = name
if steps:
self._tracked_plan.steps = steps
self._tracked_plan.updated_at = now
return {
"message": "Plan updated",
"plan": {
"name": self._tracked_plan.name,
"steps": [
{"step": s.step, "status": s.status.value}
for s in self._tracked_plan.steps
],
"created_at": self._tracked_plan.created_at,
"updated_at": self._tracked_plan.updated_at,
},
}
@self.action("get", "Get current tracked plan")
async def get(
ctx: MCPContext,
) -> dict:
"""Get the current tracked plan.
Returns:
Current plan with name, steps, and timestamps
"""
if self._tracked_plan is None:
return {"plan": None, "message": "No plan currently tracked"}
return {
"plan": {
"name": self._tracked_plan.name,
"steps": [
{"step": s.step, "status": s.status.value}
for s in self._tracked_plan.steps
],
"created_at": self._tracked_plan.created_at,
"updated_at": self._tracked_plan.updated_at,
},
"progress": {
"total": len(self._tracked_plan.steps),
"completed": sum(
1 for s in self._tracked_plan.steps
if s.status == StepStatus.COMPLETED
),
"in_progress": sum(
1 for s in self._tracked_plan.steps
if s.status == StepStatus.IN_PROGRESS
),
"pending": sum(
1 for s in self._tracked_plan.steps
if s.status == StepStatus.PENDING
),
},
}
@self.action("clear", "Clear tracked plan")
async def clear(
ctx: MCPContext,
) -> dict:
"""Clear the current tracked plan."""
self._tracked_plan = None
return {"message": "Plan cleared"}
@self.action("intent", "Parse natural language to IntentIR")
async def intent(
ctx: MCPContext,
@@ -84,6 +84,7 @@ class ProcTool(BaseTool):
Actions:
- exec: Execute command (shell: {shell_name})
- wait: Wait for background process to complete (Rust parity)
- ps: List processes
- kill: Kill process
- logs: Get process logs
@@ -98,8 +99,9 @@ Auto-backgrounds commands after {AUTO_BACKGROUND_TIMEOUT}s.
@self.action("exec", "Execute command")
async def exec_cmd(
ctx: MCPContext,
command: str,
command: str | list[str],
cwd: str | None = None,
workdir: str | None = None, # Rust parity alias
env: dict | None = None,
timeout: int | None = None,
shell: str | None = None,
@@ -109,8 +111,9 @@ Auto-backgrounds commands after {AUTO_BACKGROUND_TIMEOUT}s.
This is the ONE execution primitive per HIP-0300.
Args:
command: Command to execute
command: Command to execute (string or array of strings for Rust parity)
cwd: Working directory
workdir: Alias for cwd (Rust parity)
env: Environment variables
timeout: Timeout in seconds (default: auto-background at 45s)
shell: Shell to use (default: detected shell)
@@ -121,6 +124,16 @@ Auto-backgrounds commands after {AUTO_BACKGROUND_TIMEOUT}s.
if not command:
raise InvalidParamsError("Command is required", param="command")
# Support array format for Rust parity
if isinstance(command, list):
# Join array into shell command
import shlex
command = " ".join(shlex.quote(arg) for arg in command)
# Support workdir alias for Rust parity
if workdir and not cwd:
cwd = workdir
# Resolve shell
shell_path = shell or self._shell
@@ -179,6 +192,100 @@ Auto-backgrounds commands after {AUTO_BACKGROUND_TIMEOUT}s.
message=f"Execution failed: {e}",
)
@self.action("wait", "Wait for background process to complete (Rust parity)")
async def wait(
ctx: MCPContext,
proc_id: str,
timeout_ms: int | None = None,
) -> dict:
"""Wait for a background process to complete.
This matches the Rust wait tool schema.
Args:
proc_id: Process ID to wait for (call_id in Rust)
timeout_ms: Maximum wait time in milliseconds (default: 600000, max: 3600000)
Returns:
Process result when complete or timeout
"""
if not proc_id:
raise InvalidParamsError("proc_id is required", param="proc_id")
# Default and max timeout
max_timeout_ms = 3600000 # 1 hour max (matches Rust)
default_timeout_ms = 600000 # 10 minutes default
timeout = timeout_ms or default_timeout_ms
timeout = min(timeout, max_timeout_ms)
timeout_sec = timeout / 1000
# Get process info
procs = self.process_manager.list_processes()
if proc_id not in procs:
raise NotFoundError(f"Process not found: {proc_id}")
info = procs[proc_id]
# If already completed, return immediately
if not info.get("running", False):
log_file = info.get("log_file")
output = ""
if log_file and Path(log_file).exists():
try:
from hanzo_async import read_file
output = await read_file(log_file, encoding="utf-8", errors="replace")
except Exception:
pass
return {
"proc_id": proc_id,
"exit_code": info.get("return_code", 0),
"output": output,
"status": "completed",
}
# Poll until complete or timeout
start_time = datetime.now()
poll_interval = 0.5 # 500ms
while True:
elapsed = (datetime.now() - start_time).total_seconds()
if elapsed >= timeout_sec:
return {
"proc_id": proc_id,
"exit_code": None,
"output": "",
"status": "timeout",
"message": f"Timed out after {timeout_ms}ms",
}
# Check if process completed
procs = self.process_manager.list_processes()
if proc_id not in procs:
raise NotFoundError(f"Process disappeared: {proc_id}")
info = procs[proc_id]
if not info.get("running", False):
# Process completed
log_file = info.get("log_file")
output = ""
if log_file and Path(log_file).exists():
try:
from hanzo_async import read_file
output = await read_file(log_file, encoding="utf-8", errors="replace")
except Exception:
pass
return {
"proc_id": proc_id,
"exit_code": info.get("return_code", 0),
"output": output,
"status": "completed",
"duration_ms": int(elapsed * 1000),
}
await asyncio.sleep(poll_interval)
@self.action("ps", "List processes")
async def ps(
ctx: MCPContext,