mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user