rein-openweights/src/rein_openweights/tools.py

"""Tool surface — mirrors rein-aharness's `green-commit-only` allow-list.

Read/Write/Edit/Glob/Grep + local git status/diff/log/add+commit. No
push, no network, no shell escape hatch — the same bounded surface, just
invoked by a different rein.
"""

from __future__ import annotations

import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable


class ToolPolicyError(RuntimeError):
    pass


def _resolve(repo_root: Path, rel_path: str) -> Path:
    repo_root = repo_root.resolve()
    candidate = (repo_root / rel_path).resolve()
    if candidate != repo_root and repo_root not in candidate.parents:
        raise ToolPolicyError(f"path escapes repo root: {rel_path}")
    return candidate


def read_file(repo_root: Path, path: str) -> str:
    return _resolve(repo_root, path).read_text()


def write_file(repo_root: Path, path: str, content: str) -> str:
    target = _resolve(repo_root, path)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(content)
    return f"wrote {len(content)} bytes to {path}"


def edit_file(repo_root: Path, path: str, old: str, new: str) -> str:
    target = _resolve(repo_root, path)
    text = target.read_text()
    if old not in text:
        raise ToolPolicyError(f"old string not found in {path}")
    target.write_text(text.replace(old, new, 1))
    return f"replaced 1 occurrence in {path}"


def glob_files(repo_root: Path, pattern: str) -> str:
    repo_root = repo_root.resolve()
    matches = sorted(str(p.relative_to(repo_root)) for p in repo_root.rglob(pattern) if p.is_file())
    return "\n".join(matches) or "(no matches)"


def grep_files(repo_root: Path, pattern: str, path: str = ".") -> str:
    base = _resolve(repo_root, path)
    hits: list[str] = []
    for file in sorted(base.rglob("*")):
        if not file.is_file() or ".git" in file.parts:
            continue
        try:
            text = file.read_text()
        except (UnicodeDecodeError, OSError):
            continue
        for lineno, line in enumerate(text.splitlines(), start=1):
            if pattern in line:
                hits.append(f"{file.relative_to(repo_root)}:{lineno}:{line}")
    return "\n".join(hits) or "(no matches)"


def _git(repo_root: Path, *args: str) -> str:
    proc = subprocess.run(["git", "-C", str(repo_root), *args], capture_output=True, text=True)
    if proc.returncode != 0:
        raise ToolPolicyError(f"git {' '.join(args)} failed: {proc.stderr.strip()}")
    return proc.stdout


def git_status(repo_root: Path) -> str:
    return _git(repo_root, "status", "--short") or "(clean)"


def git_diff(repo_root: Path) -> str:
    return _git(repo_root, "diff") or "(no diff)"


def git_log(repo_root: Path, max_count: int = 10) -> str:
    return _git(repo_root, "log", f"-{max_count}", "--oneline")


def git_add_commit(repo_root: Path, message: str) -> str:
    _git(repo_root, "add", "-A")
    return _git(repo_root, "commit", "-m", message)


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict[str, Any]
    fn: Callable[..., str]


TOOLS: list[Tool] = [
    Tool(
        "read_file",
        "Read a file's contents.",
        {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]},
        read_file,
    ),
    Tool(
        "write_file",
        "Write (overwrite) a file's contents, creating parent dirs as needed.",
        {
            "type": "object",
            "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
            "required": ["path", "content"],
        },
        write_file,
    ),
    Tool(
        "edit_file",
        "Replace the first occurrence of `old` with `new` in a file.",
        {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "old": {"type": "string"},
                "new": {"type": "string"},
            },
            "required": ["path", "old", "new"],
        },
        edit_file,
    ),
    Tool(
        "glob_files",
        "List files under the repo matching a glob pattern (e.g. '**/*.py').",
        {"type": "object", "properties": {"pattern": {"type": "string"}}, "required": ["pattern"]},
        glob_files,
    ),
    Tool(
        "grep_files",
        "Search file contents under `path` for a literal substring.",
        {
            "type": "object",
            "properties": {"pattern": {"type": "string"}, "path": {"type": "string"}},
            "required": ["pattern"],
        },
        grep_files,
    ),
    Tool("git_status", "Show `git status --short`.", {"type": "object", "properties": {}}, git_status),
    Tool("git_diff", "Show `git diff`.", {"type": "object", "properties": {}}, git_diff),
    Tool(
        "git_log",
        "Show recent commit log (oneline).",
        {"type": "object", "properties": {"max_count": {"type": "integer"}}},
        git_log,
    ),
    Tool(
        "git_add_commit",
        "Stage all changes (`git add -A`) and commit with a message. "
        "Call this once the task is complete.",
        {"type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"]},
        git_add_commit,
    ),
]

TOOLS_BY_NAME = {t.name: t for t in TOOLS}


def openai_tool_schemas() -> list[dict[str, Any]]:
    return [
        {
            "type": "function",
            "function": {
                "name": t.name,
                "description": t.description,
                "parameters": t.parameters,
            },
        }
        for t in TOOLS
    ]


def call_tool(repo_root: Path, name: str, args: dict[str, Any]) -> str:
    tool = TOOLS_BY_NAME.get(name)
    if tool is None:
        raise ToolPolicyError(f"unknown tool: {name}")
    return tool.fn(repo_root, **args)