from __future__ import annotations import subprocess import os import re import shutil from dataclasses import dataclass from pathlib import Path from tools.capabilities import Capability @dataclass(frozen=True) class ToolResult: returncode: int stdout: str stderr: str @dataclass(frozen=True) class MutationResult: success: bool operation: str files_changed: list[str] hunks_applied: int = 0 error: 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 write_file(self, relative_path: str, content: str) -> MutationResult: try: self.write_text(relative_path, content) except Exception as exc: return MutationResult(False, "write_file", [], error=str(exc)) return MutationResult(True, "write_file", [relative_path]) def create_directory(self, relative_path: str) -> MutationResult: self._require(Capability.WRITE_WORKTREE) try: path = self._safe_path(relative_path) if path.exists() and not path.is_dir(): return MutationResult(False, "create_directory", [], error="Target exists as a file") path.mkdir(parents=True, exist_ok=True) except Exception as exc: return MutationResult(False, "create_directory", [], error=str(exc)) return MutationResult(True, "create_directory", [relative_path]) def delete_file(self, relative_path: str) -> MutationResult: self._require(Capability.WRITE_WORKTREE) try: path = self._safe_path(relative_path) if not path.exists(): return MutationResult(False, "delete_file", [], error="Target file does not exist") if path.is_dir(): return MutationResult(False, "delete_file", [], error="Target is a directory") path.unlink() except Exception as exc: return MutationResult(False, "delete_file", [], error=str(exc)) return MutationResult(True, "delete_file", [relative_path]) def move_file(self, source: str, destination: str, overwrite: bool = False) -> MutationResult: self._require(Capability.WRITE_WORKTREE) try: source_path = self._safe_path(source) destination_path = self._safe_path(destination) if not source_path.exists(): return MutationResult(False, "move_file", [], error="Source file does not exist") if source_path.is_dir(): return MutationResult(False, "move_file", [], error="Source is a directory") if destination_path.exists() and not overwrite: return MutationResult(False, "move_file", [], error="Destination already exists") destination_path.parent.mkdir(parents=True, exist_ok=True) shutil.move(str(source_path), str(destination_path)) except Exception as exc: return MutationResult(False, "move_file", [], error=str(exc)) return MutationResult(True, "move_file", [source, destination]) def apply_patch(self, relative_path: str, patch: str) -> MutationResult: self._require(Capability.WRITE_WORKTREE) try: path = self._safe_path(relative_path) if not path.is_file(): return MutationResult(False, "apply_patch", [], error="Target file does not exist") original = path.read_text(encoding="utf-8").splitlines(keepends=True) updated, hunks = self._apply_unified_patch(original, patch) path.write_text("".join(updated), encoding="utf-8") except Exception as exc: return MutationResult(False, "apply_patch", [], error=str(exc)) return MutationResult(True, "apply_patch", [relative_path], hunks_applied=hunks) def _apply_unified_patch(self, original: list[str], patch: str) -> tuple[list[str], int]: patch_lines = patch.splitlines(keepends=True) output: list[str] = [] source_index = 0 index = 0 hunks = 0 while index < len(patch_lines): line = patch_lines[index] if not line.startswith("@@"): index += 1 continue match = re.match(r"@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@", line) if not match: raise ValueError(f"Invalid hunk header: {line.strip()}") old_start = int(match.group(1)) - 1 index += 1 hunk_lines: list[str] = [] while index < len(patch_lines) and not patch_lines[index].startswith("@@"): hunk_lines.append(patch_lines[index]) index += 1 old_lines = self._patch_old_lines(hunk_lines) actual_start = self._locate_hunk(original, old_lines, old_start, source_index) if actual_start is None: raise ValueError("Patch context does not match target file") if actual_start < source_index: raise ValueError("Overlapping or out-of-order patch hunk") output.extend(original[source_index:actual_start]) source_index = actual_start hunks += 1 for hunk_line in hunk_lines: if hunk_line.startswith(("---", "+++")): continue if hunk_line.startswith(" "): expected = hunk_line[1:] if source_index >= len(original) or original[source_index] != expected: raise ValueError("Patch context does not match target file") output.append(original[source_index]) source_index += 1 elif hunk_line.startswith("-"): expected = hunk_line[1:] if source_index >= len(original) or original[source_index] != expected: raise ValueError("Patch removal does not match target file") source_index += 1 elif hunk_line.startswith("+"): output.append(hunk_line[1:]) elif hunk_line.startswith("\\ No newline at end of file"): pass elif hunk_line.strip() == "": raise ValueError("Blank patch lines must be prefixed with context/add/remove marker") else: raise ValueError(f"Invalid hunk line: {hunk_line[:40]}") if hunks == 0: raise ValueError("Patch contains no hunks") output.extend(original[source_index:]) return output, hunks def _patch_old_lines(self, hunk_lines: list[str]) -> list[str]: old_lines: list[str] = [] for hunk_line in hunk_lines: if hunk_line.startswith(("---", "+++", "+", "\\ No newline at end of file")): continue if hunk_line.startswith((" ", "-")): old_lines.append(hunk_line[1:]) elif hunk_line.strip() == "": raise ValueError("Blank patch lines must be prefixed with context/add/remove marker") else: raise ValueError(f"Invalid hunk line: {hunk_line[:40]}") return old_lines def _locate_hunk(self, original: list[str], old_lines: list[str], old_start: int, source_index: int) -> int | None: if not old_lines: return max(old_start, source_index) if old_start >= source_index else None if self._matches_at(original, old_lines, old_start): return old_start window_start = max(source_index, old_start - 200) window_end = min(len(original), old_start + 200) for candidate in range(window_start, window_end + 1): if self._matches_at(original, old_lines, candidate): return candidate for candidate in range(source_index, len(original) + 1): if self._matches_at(original, old_lines, candidate): return candidate return None def _matches_at(self, original: list[str], old_lines: list[str], start: int) -> bool: if start < 0 or start + len(old_lines) > len(original): return False return original[start : start + len(old_lines)] == old_lines 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()