140 lines
5.3 KiB
Python
140 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
import subprocess
|
|
import os
|
|
import re
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from tools.capabilities import Capability
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ToolResult:
|
|
returncode: int
|
|
stdout: str
|
|
stderr: str
|
|
|
|
|
|
class CapabilityError(PermissionError):
|
|
pass
|
|
|
|
|
|
class WorktreeTools:
|
|
def __init__(self, worktree_path: Path, capabilities: set[Capability]) -> None:
|
|
self.worktree_path = worktree_path.resolve()
|
|
self.capabilities = capabilities
|
|
|
|
def _require(self, capability: Capability) -> None:
|
|
if capability not in self.capabilities:
|
|
raise CapabilityError(f"Missing capability: {capability.value}")
|
|
|
|
def _safe_path(self, relative_path: str) -> Path:
|
|
path = (self.worktree_path / relative_path).resolve()
|
|
if self.worktree_path not in path.parents and path != self.worktree_path:
|
|
raise CapabilityError("Path escapes worktree")
|
|
return path
|
|
|
|
def read_text(self, relative_path: str) -> str:
|
|
self._require(Capability.READ_REPOSITORY)
|
|
return self._safe_path(relative_path).read_text(encoding="utf-8")
|
|
|
|
def list_directory(self, relative_path: str = ".") -> list[str]:
|
|
self._require(Capability.INVESTIGATE_WORKTREE)
|
|
path = self._safe_path(relative_path)
|
|
if not path.is_dir():
|
|
raise CapabilityError("Path is not a directory")
|
|
entries: list[str] = []
|
|
for child in sorted(path.iterdir(), key=lambda item: item.name):
|
|
if child.name == ".git" or child.name == "__pycache__":
|
|
continue
|
|
suffix = "/" if child.is_dir() else ""
|
|
entries.append(f"{child.relative_to(self.worktree_path).as_posix()}{suffix}")
|
|
return entries
|
|
|
|
def read_file(self, relative_path: str, max_chars: int = 12000) -> str:
|
|
self._require(Capability.INVESTIGATE_WORKTREE)
|
|
path = self._safe_path(relative_path)
|
|
if not path.is_file():
|
|
raise CapabilityError("Path is not a file")
|
|
return path.read_text(encoding="utf-8")[:max_chars]
|
|
|
|
def search_code(self, pattern: str, include: str = "*.py", max_results: int = 50) -> list[dict[str, object]]:
|
|
self._require(Capability.INVESTIGATE_WORKTREE)
|
|
regex = re.compile(pattern)
|
|
results: list[dict[str, object]] = []
|
|
for path in self.worktree_path.rglob(include):
|
|
if ".git" in path.parts or "__pycache__" in path.parts or not path.is_file():
|
|
continue
|
|
relative = path.relative_to(self.worktree_path).as_posix()
|
|
for number, line in enumerate(path.read_text(encoding="utf-8", errors="ignore").splitlines(), start=1):
|
|
if regex.search(line):
|
|
results.append({"path": relative, "line": number, "text": line[:300]})
|
|
if len(results) >= max_results:
|
|
return results
|
|
return results
|
|
|
|
def find_symbol(self, symbol: str, max_results: int = 50) -> list[dict[str, object]]:
|
|
escaped = re.escape(symbol)
|
|
return self.search_code(rf"^\s*(class|def)\s+{escaped}\b|\b{escaped}\b", max_results=max_results)
|
|
|
|
def write_text(self, relative_path: str, content: str) -> None:
|
|
self._require(Capability.WRITE_WORKTREE)
|
|
path = self._safe_path(relative_path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(content, encoding="utf-8")
|
|
|
|
def run(self, command: list[str], timeout: int = 60) -> ToolResult:
|
|
self._require(Capability.RUN_TESTS)
|
|
env = os.environ.copy()
|
|
for key in ["DJANGO_SETTINGS_MODULE", "DATABASE_URL"]:
|
|
env.pop(key, None)
|
|
completed = subprocess.run(
|
|
command,
|
|
cwd=self.worktree_path,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
env=env,
|
|
)
|
|
return ToolResult(completed.returncode, completed.stdout, completed.stderr)
|
|
|
|
def git(self, args: list[str]) -> ToolResult:
|
|
required = Capability.COMMIT_CHANGES if args and args[0] == "commit" else Capability.READ_REPOSITORY
|
|
self._require(required)
|
|
completed = subprocess.run(
|
|
["git", *args],
|
|
cwd=self.worktree_path,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
return ToolResult(completed.returncode, completed.stdout, completed.stderr)
|
|
|
|
def status(self) -> str:
|
|
return self.git(["status", "--short"]).stdout
|
|
|
|
def diff(self) -> str:
|
|
return self.git(["diff", "--", "."]).stdout
|
|
|
|
def git_status(self) -> str:
|
|
self._require(Capability.INVESTIGATE_WORKTREE)
|
|
return self.status()
|
|
|
|
def git_diff(self) -> str:
|
|
self._require(Capability.INVESTIGATE_WORKTREE)
|
|
return self.diff()
|
|
|
|
def commit_all(self, message: str) -> str:
|
|
self._require(Capability.COMMIT_CHANGES)
|
|
add = self.git(["add", "."])
|
|
if add.returncode != 0:
|
|
raise RuntimeError(add.stderr)
|
|
commit = self.git(["-c", "user.name=Artifex", "-c", "user.email=artifex@example.invalid", "commit", "-m", message])
|
|
if commit.returncode != 0:
|
|
raise RuntimeError(commit.stderr or commit.stdout)
|
|
rev = self.git(["rev-parse", "HEAD"])
|
|
if rev.returncode != 0:
|
|
raise RuntimeError(rev.stderr)
|
|
return rev.stdout.strip()
|