codetrail-cli 1.4.1__py3-none-any.whl

codetrail/assistant/guard.py

@@ -0,0 +1,80 @@
+"""The tool guard: the assistant may only read `source/`, with Read, Grep and Glob (design section 6.9).
+
+It runs as a PreToolUse hook, which sees every tool call, read-only ones included, before it runs. Anything it can't
+place inside `source/` is refused; every file Read opens is recorded, and becomes the page's `files`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+ALLOWED_TOOLS = ("Read", "Grep", "Glob")
+# The channel structured answers come back through; it reads nothing and takes no path.
+ANSWER_TOOL = "StructuredOutput"
+
+
+class ToolGuard:
+    def __init__(self, source_root: Path) -> None:
+        self.root = source_root.resolve()
+        self.files_read: list[str] = []
+
+    def decide(self, tool: str, arguments: Mapping[str, Any]) -> str | None:
+        """None to allow the call, or the reason it is refused; any error refuses (fails closed)."""
+        try:
+            return self._decide(tool, arguments)
+        except Exception:
+            return "The request couldn't be checked, so it is refused."
+
+    def _decide(self, tool: str, arguments: Mapping[str, Any]) -> str | None:
+        if tool == ANSWER_TOOL:
+            return None
+        if tool not in ALLOWED_TOOLS:
+            return f"{tool} isn't available; only Read, Grep and Glob inside the repository are."
+        if tool == "Read":
+            path = arguments.get("file_path")
+            if not isinstance(path, str):
+                return "Read needs a file path."
+            inside = self._inside(path)
+            if inside is None or not inside.is_file():
+                return f"{path} isn't a file in the repository."
+            relative = inside.relative_to(self.root).as_posix()
+            if relative not in self.files_read:
+                self.files_read.append(relative)
+            return None
+        folder = arguments.get("path")
+        if folder is not None and (not isinstance(folder, str) or self._inside(folder) is None):
+            return f"{folder} isn't inside the repository."
+        for key in ("pattern", "glob") if tool == "Glob" else ("glob",):
+            pattern = arguments.get(key)
+            # Braces could expand to .. or an absolute path before the pattern is resolved; refuse them outright.
+            if isinstance(pattern, str) and (
+                pattern.startswith(("/", "~")) or ".." in pattern or any(char in pattern for char in "{}\\\0")
+            ):
+                return f"The pattern {pattern} could reach outside the repository."
+        return None
+
+    def _inside(self, path: str) -> Path | None:
+        if "\0" in path or path.startswith("~"):
+            return None
+        candidate = Path(path)
+        resolved = (candidate if candidate.is_absolute() else self.root / candidate).resolve()
+        return resolved if resolved == self.root or resolved.is_relative_to(self.root) else None
+
+    async def hook(self, input_data: Mapping[str, Any], tool_use_id: str | None, context: Any) -> dict[str, Any]:
+        arguments = input_data.get("tool_input")
+        if not isinstance(arguments, Mapping):
+            arguments = {"invalid": True}
+        reason = self.decide(str(input_data.get("tool_name", "")), arguments)
+        if reason is None and arguments.get("invalid") is True and len(arguments) == 1:
+            reason = "The request couldn't be checked, so it is refused."
+        if reason is None:
+            return {}
+        return {
+            "hookSpecificOutput": {
+                "hookEventName": "PreToolUse",
+                "permissionDecision": "deny",
+                "permissionDecisionReason": reason,
+            }
+        }

codetrail/assistant/guard_hook.py

@@ -0,0 +1,58 @@
+"""The tool guard as Claude Code's PreToolUse hook: `python -m codetrail.assistant.guard_hook <source> <read-log>`.
+
+Claude Code runs it before every tool call, with the call as JSON on stdin. A refusal is printed as the hook's deny
+decision; each allowed Read's path is appended to the read log, which becomes the page's `files`. Claude Code runs a
+tool when a hook fails in any way other than exit code 2, so every error of the hook's own exits with code 2, which
+refuses the call (design section 15.1). Claude Code's permission rules confine reads independently of this hook.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Mapping
+from pathlib import Path
+
+
+def main(arguments: list[str]) -> int:
+    from codetrail.assistant.guard import ToolGuard
+
+    source, log = Path(arguments[0]), Path(arguments[1])
+    payload = json.loads(sys.stdin.read())
+    guard = ToolGuard(source)
+    if not isinstance(payload, Mapping):
+        raise ValueError("The hook's input isn't a JSON object.")
+    tool_input = payload.get("tool_input")
+    if not isinstance(tool_input, Mapping):
+        raise ValueError("The tool's input isn't a JSON object.")
+    reason = guard.decide(str(payload.get("tool_name", "")), tool_input)
+    if reason is not None:
+        print(
+            json.dumps(
+                {
+                    "hookSpecificOutput": {
+                        "hookEventName": "PreToolUse",
+                        "permissionDecision": "deny",
+                        "permissionDecisionReason": reason,
+                    }
+                }
+            )
+        )
+        return 0
+    if guard.files_read:
+        with log.open("a", encoding="utf-8") as handle:
+            handle.writelines(f"{path}\n" for path in guard.files_read)
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main(sys.argv[1:]))
+    except BaseException as error:  # any failure must refuse the call, never let it through
+        if isinstance(error, SystemExit) and error.code == 0:
+            raise
+        print(
+            f"Codetrail's tool guard couldn't check this call ({type(error).__name__}), so it is refused.",
+            file=sys.stderr,
+        )
+        sys.exit(2)

codetrail/assistant/local.py

@@ -0,0 +1,245 @@
+"""The local adapter: a model on the reader's machine, through Codetrail's own read-only tools (design 15.1, 15.5).
+
+It talks to an OpenAI-compatible chat endpoint on loopback (Ollama or LM Studio) with httpx, ignoring proxy settings
+and never following redirects. The model may call Read, Grep and Glob, which Codetrail runs over `source/` through the
+tool guard. A structured answer must be JSON with the schema's keys; one retry is given with the errors. Turns,
+tokens and time are limited by configuration. Answers are buffered, so the bridge can scan them before the reader sees
+them. Local models cost nothing.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Final
+
+import anyio
+import httpx
+
+from codetrail.assistant import (
+    AnswerChunk,
+    AssistantError,
+    DigestDraft,
+    DigestRequest,
+    GradeRequest,
+    PageDraft,
+    PageRequest,
+    PlanDraft,
+    PlanRequest,
+    QuestionRequest,
+    Usage,
+    Verdict,
+)
+from codetrail.assistant.local_tools import SCHEMAS, LocalTools
+from codetrail.assistant.prompts import (
+    ANSWER_RULES,
+    DIGEST_SCHEMA,
+    GRADE_RULES,
+    GRADE_SCHEMA,
+    GROUND_RULES,
+    PAGE_SYNTAX,
+    PLAN_SCHEMA,
+    answer_prompt,
+    digest_prompt,
+    grade_prompt,
+    neutralize,
+    page_task,
+    plan_prompt,
+)
+from codetrail.config import GenerationSettings, LocalSettings
+
+PROVIDER: Final = "local"
+MAX_TOOL_CALLS_PER_TURN = 20
+TOOL_RULES = "Use the read, grep and glob tools to look at the repository's files before you answer."
+
+
+@dataclass
+class _Result:
+    text: str
+    files_read: list[str]
+    usage: Usage
+
+
+class LocalAssistant:
+    def __init__(
+        self,
+        source_root: Path,
+        models: Mapping[str, str],
+        settings: LocalSettings,
+        limits: GenerationSettings,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self.root = source_root.resolve()
+        self.models = dict(models)
+        self.settings = settings
+        self.limits = limits
+        self.client = client or httpx.AsyncClient(
+            base_url=settings.base_url,
+            trust_env=False,
+            follow_redirects=False,
+            timeout=httpx.Timeout(settings.timeout_seconds),
+        )
+
+    async def plan(self, request: PlanRequest) -> PlanDraft:
+        data, result = await self._structured("plan", GROUND_RULES, plan_prompt(request), PLAN_SCHEMA)
+        return PlanDraft(list(data.get("pages", [])), list(data.get("paths", [])), result.files_read, 0.0, result.usage)
+
+    async def write_page(self, request: PageRequest) -> PageDraft:
+        prompt, schema = page_task(request)  # a new page, or a revision's sections
+        data, result = await self._structured("write", GROUND_RULES + "\n" + PAGE_SYNTAX, prompt, schema)
+        return PageDraft(str(data.get("body", "")), list(data.get("checks", [])), result.files_read, 0.0, result.usage,
+                         list(data.get("sections", [])), str(data.get("brief", "")))  # fmt: skip
+
+    async def write_digest(self, request: DigestRequest) -> DigestDraft:
+        data, result = await self._structured(
+            "digest", GROUND_RULES + "\n" + PAGE_SYNTAX, digest_prompt(request), DIGEST_SCHEMA
+        )
+        return DigestDraft(str(data.get("title", "")), str(data.get("body", "")), result.files_read, 0.0, result.usage)
+
+    async def grade(self, request: GradeRequest) -> Verdict:
+        data, result = await self._structured("grade", GRADE_RULES, grade_prompt(request), GRADE_SCHEMA, tools=False)
+        missed = [str(item) for item in data.get("missed", [])]
+        return Verdict(str(data.get("verdict", "")), missed, str(data.get("feedback", "")), 0.0, result.usage)
+
+    async def answer(self, request: QuestionRequest) -> AsyncIterator[AnswerChunk]:
+        rules = GROUND_RULES + "\n" + PAGE_SYNTAX + "\n" + ANSWER_RULES
+        result = await self._loop("answer", rules, answer_prompt(request), None)
+        # One final chunk: the bridge shows it only after scanning it (design section 15.5).
+        yield AnswerChunk(result.text, True, result.files_read, usage=result.usage)
+
+    async def _structured(
+        self, kind: str, rules: str, prompt: str, schema: dict[str, Any], tools: bool = True
+    ) -> tuple[dict[str, Any], _Result]:
+        result = await self._loop(kind, rules, prompt, schema, tools)
+        data = _parse(result.text, schema)
+        if isinstance(data, str):
+            raise AssistantError(f"The local model's answer wasn't usable: {data}")
+        return data, result
+
+    async def _loop(
+        self, kind: str, rules: str, prompt: str, schema: dict[str, Any] | None, tools: bool = True
+    ) -> _Result:
+        model = self.models.get(kind, "")
+        files = LocalTools(self.root, self.settings.max_read_bytes)
+        system = rules
+        if tools:
+            system += "\n" + TOOL_RULES
+        if schema is not None:
+            system += "\nFinish with only a JSON object that follows this schema, and nothing else:\n" + json.dumps(
+                schema
+            )
+        messages: list[dict[str, Any]] = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": neutralize(prompt)},
+        ]
+        tokens_in = tokens_out = 0
+        retried = False
+        with anyio.fail_after(self.settings.timeout_seconds):
+            for _turn in range(self.settings.max_turns):
+                body: dict[str, Any] = {"model": model, "messages": messages, "temperature": 0}
+                if tools:
+                    body["tools"] = SCHEMAS
+                reply = await self._chat(body)
+                usage = reply.get("usage") or {}
+                tokens_in += _count(usage, "prompt_tokens")
+                tokens_out += _count(usage, "completion_tokens")
+                if tokens_in + tokens_out > self.settings.max_tokens_per_call:
+                    raise AssistantError(
+                        f"The local model passed its token limit ({self.settings.max_tokens_per_call})."
+                    )
+                message = _message(reply)
+                calls = message.get("tool_calls") or []
+                if calls and tools:
+                    messages.append({"role": "assistant", "content": message.get("content") or "", "tool_calls": calls})
+                    for index, call in enumerate(calls):
+                        content = (
+                            await _run_tool(files, call)
+                            if index < MAX_TOOL_CALLS_PER_TURN
+                            else f"Refused: too many tool calls in one turn (at most {MAX_TOOL_CALLS_PER_TURN})."
+                        )
+                        messages.append({"role": "tool", "tool_call_id": str(call.get("id", "")), "content": content})
+                    continue
+                text = str(message.get("content") or "")
+                if schema is not None and isinstance(problem := _parse(text, schema), str) and not retried:
+                    retried = True
+                    messages += [
+                        {"role": "assistant", "content": text},
+                        {
+                            "role": "user",
+                            "content": f"That wasn't usable: {problem} Answer again with only the JSON object.",
+                        },
+                    ]
+                    continue
+                usage_record = Usage(PROVIDER, model, tokens_in, 0, tokens_out, 0.0)
+                return _Result(text, list(files.files_read), usage_record)
+        raise AssistantError(f"The local model didn't finish within {self.settings.max_turns} turns.")
+
+    async def _chat(self, body: dict[str, Any]) -> dict[str, Any]:
+        try:
+            response = await self.client.post("chat/completions", json=body)
+        except httpx.HTTPError as error:
+            raise AssistantError(
+                f"The local model at {self.settings.base_url} couldn't be reached ({type(error).__name__})."
+            ) from error
+        if response.status_code != 200:
+            raise AssistantError(f"The local model answered with HTTP {response.status_code}.")
+        try:
+            reply = response.json()
+        except ValueError as error:
+            raise AssistantError("The local model's reply wasn't JSON.") from error
+        if not isinstance(reply, dict):
+            raise AssistantError("The local model's reply wasn't a JSON object.")
+        return reply
+
+
+async def _run_tool(files: LocalTools, call: Mapping[str, Any]) -> str:
+    function = call.get("function") or {}
+    try:
+        arguments = json.loads(function.get("arguments") or "{}")
+    except ValueError:
+        return "Refused: the tool's arguments weren't JSON."
+    if not isinstance(arguments, dict):
+        return "Refused: the tool's arguments weren't a JSON object."
+    # In a worker thread, so a slow search can't block the event loop or outlast the time limit.
+    return await anyio.to_thread.run_sync(files.run, str(function.get("name", "")), arguments, abandon_on_cancel=True)
+
+
+def _message(reply: Mapping[str, Any]) -> dict[str, Any]:
+    choices = reply.get("choices")
+    if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict):
+        raise AssistantError("The local model's reply had no message.")
+    message = choices[0].get("message")
+    if not isinstance(message, dict):
+        raise AssistantError("The local model's reply had no message.")
+    return message
+
+
+def _unfenced(text: str) -> str:
+    """The text inside a ```json fence, if the whole answer is one; plain string work, so it stays linear."""
+    stripped = text.strip()
+    if stripped.startswith("```") and stripped.endswith("```") and len(stripped) >= 6:
+        inner = stripped[3:-3]
+        first_line, _newline, rest = inner.partition("\n")
+        return rest if first_line.strip() in ("", "json") else inner
+    return stripped
+
+
+def _parse(text: str, schema: Mapping[str, Any]) -> dict[str, Any] | str:
+    """The answer as a JSON object with the schema's required keys, or what is wrong with it."""
+    try:
+        data = json.loads(_unfenced(text))
+    except ValueError:
+        return "it isn't valid JSON."
+    if not isinstance(data, dict):
+        return "it isn't a JSON object."
+    missing = [key for key in schema.get("required", []) if key not in data]
+    if missing:
+        return f"it lacks {', '.join(missing)}."
+    return data
+
+
+def _count(usage: Mapping[str, Any], key: str) -> int:
+    value = usage.get(key, 0)
+    return int(value) if isinstance(value, int | float) else 0

codetrail/assistant/local_tools.py

@@ -0,0 +1,177 @@
+"""The local provider's file tools, implemented by Codetrail: Read, Grep and Glob, only inside `source/` (design 15.1).
+
+Every path goes through the tool guard. Grep searches for plain text, never a regular expression, so no pattern from
+the model can run away; Glob patterns are short, with few wildcards, for the same reason. A refusal comes back to the
+model as text starting with "Refused:", so it can try something else.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+from codetrail.assistant.guard import ToolGuard
+
+MAX_GREP_LINES = 200
+MAX_GLOB_PATHS = 500
+MAX_PATTERN_CHARS = 200
+MAX_WILDCARDS = 10
+MAX_LINE_CHARS = 300
+
+SCHEMAS: list[dict[str, Any]] = [
+    {
+        "type": "function",
+        "function": {
+            "name": "read",
+            "description": "Read a text file in the repository, with line numbers.",
+            "parameters": {
+                "type": "object",
+                "required": ["path"],
+                "properties": {
+                    "path": {"type": "string", "description": "A path relative to the repository's root."},
+                    "offset": {"type": "integer", "description": "The first line to read, from 1."},
+                    "limit": {"type": "integer", "description": "How many lines to read."},
+                },
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "grep",
+            "description": "Find lines containing some text (plain text, not case-sensitive).",
+            "parameters": {
+                "type": "object",
+                "required": ["text"],
+                "properties": {
+                    "text": {"type": "string"},
+                    "path": {"type": "string", "description": "A folder to search in; the root by default."},
+                    "glob": {"type": "string", "description": "Only files whose path matches this pattern, like *.py."},
+                },
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "glob",
+            "description": "List the files whose paths match a pattern, like **/*.py.",
+            "parameters": {
+                "type": "object",
+                "required": ["pattern"],
+                "properties": {
+                    "pattern": {"type": "string"},
+                    "path": {"type": "string", "description": "A folder to search in; the root by default."},
+                },
+            },
+        },
+    },
+]
+
+
+class _Refused(Exception):
+    pass
+
+
+class LocalTools:
+    def __init__(self, source_root: Path, max_read_bytes: int) -> None:
+        self.guard = ToolGuard(source_root)
+        self.root = self.guard.root
+        self.max_read_bytes = max_read_bytes
+
+    @property
+    def files_read(self) -> list[str]:
+        return self.guard.files_read
+
+    def run(self, name: str, arguments: Mapping[str, Any]) -> str:
+        try:
+            if name == "read":
+                return self._read(
+                    _text(arguments, "path"), _number(arguments, "offset", 1), _number(arguments, "limit", 2000)
+                )
+            if name == "grep":
+                return self._grep(_text(arguments, "text"), _text(arguments, "path", "."), arguments.get("glob"))
+            if name == "glob":
+                return self._glob(_text(arguments, "pattern"), _text(arguments, "path", "."))
+            raise _Refused(f"{name} isn't a tool here; only read, grep and glob are.")
+        except _Refused as refusal:
+            return f"Refused: {refusal}"
+
+    def _read(self, path: str, offset: int, limit: int) -> str:
+        reason = self.guard.decide("Read", {"file_path": path})
+        if reason is not None:
+            raise _Refused(reason)
+        with (self.root / path).open("rb") as handle:
+            data = handle.read(self.max_read_bytes)
+        lines = data.decode("utf-8", "replace").splitlines()
+        start = max(offset, 1)
+        shown = lines[start - 1 : start - 1 + max(limit, 1)]
+        return "\n".join(f"{start + index}\t{line}" for index, line in enumerate(shown))
+
+    def _grep(self, text: str, path: str, glob: object) -> str:
+        if not text:
+            raise _Refused("grep needs some text to find.")
+        folder = self._folder(path)
+        pattern = _pattern(glob) if glob is not None else None
+        needle = text.lower()
+        found: list[str] = []
+        for file in sorted(folder.rglob("*")):
+            if not self._plain_file(file) or (pattern and not fnmatch.fnmatch(file.name, pattern)):
+                continue
+            relative = file.relative_to(self.root).as_posix()
+            with file.open("rb") as handle:
+                content = handle.read(self.max_read_bytes).decode("utf-8", "replace")
+            for number, line in enumerate(content.splitlines(), start=1):
+                if needle in line.lower():
+                    found.append(f"{relative}:{number}: {line.strip()[:MAX_LINE_CHARS]}")
+                    if len(found) >= MAX_GREP_LINES:
+                        return "\n".join(found)
+        return "\n".join(found) or "No matches."
+
+    def _glob(self, pattern: str, path: str) -> str:
+        reason = self.guard.decide("Glob", {"pattern": _pattern(pattern), "path": path})
+        if reason is not None:
+            raise _Refused(reason)
+        folder = self._folder(path)
+        paths = sorted(
+            file.relative_to(self.root).as_posix() for file in folder.glob(pattern) if self._plain_file(file)
+        )
+        return "\n".join(paths[:MAX_GLOB_PATHS]) or "No matches."
+
+    def _plain_file(self, file: Path) -> bool:
+        """A regular file inside the root, reached through no symlink: source/ never holds one, but don't rely on it."""
+        if file.is_symlink() or not file.is_file():
+            return False
+        resolved = file.resolve()
+        return resolved.is_relative_to(self.root) and resolved == file.absolute()
+
+    def _folder(self, path: str) -> Path:
+        reason = self.guard.decide("Grep", {"path": path})
+        folder = (self.root / path).resolve()
+        if reason is not None or not folder.is_dir():
+            raise _Refused(reason or f"{path} isn't a folder in the repository.")
+        return folder
+
+
+def _pattern(value: object) -> str:
+    if not isinstance(value, str) or not value:
+        raise _Refused("A pattern must be some text.")
+    if len(value) > MAX_PATTERN_CHARS or value.count("*") + value.count("?") > MAX_WILDCARDS:
+        raise _Refused(f"Patterns are limited to {MAX_PATTERN_CHARS} characters and {MAX_WILDCARDS} wildcards.")
+    return value
+
+
+def _text(arguments: Mapping[str, Any], key: str, default: str | None = None) -> str:
+    value = arguments.get(key, default)
+    if not isinstance(value, str):
+        raise _Refused(f"{key} must be text.")
+    return value
+
+
+def _number(arguments: Mapping[str, Any], key: str, default: int) -> int:
+    value = arguments.get(key, default)
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise _Refused(f"{key} must be a whole number.")
+    return value