beaver-gateway/src/beaver_gateway/agents/policy.py

"""PreToolUse policy - the boundary without permission prompts (§3.7).

``bypassPermissions`` everywhere; what a model may do is decided by
mounts, ``disallowed_tools`` and the rules here. A rule is a callable
``(ToolCall) -> Deny | None`` declared per agent (``ClaudeAgent.policy``);
the SDK backend registers one in-process ``PreToolUse`` hook that runs
the rules in order and turns the first :class:`Deny` into a hook deny
whose reason the model reads as the tool result. Rules never see
secrets and never prompt - they only say no, with a reason.

A rule that raises is a deny too: the boundary fails closed, the
traceback lands in the log.

Every tool call - allowed or denied - is reported to the audit sink the
backend was given, so the admin audit page shows what the model touched.
"""

from __future__ import annotations

import inspect
import json
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, MutableMapping

__all__ = [
    "AuditSink",
    "Decision",
    "Deny",
    "PolicyRule",
    "ToolAudit",
    "ToolCall",
    "brief",
    "evaluate",
    "hook_output",
]

_log = logging.getLogger("beaver_gateway.policy")

PATH_KEYS: tuple[str, ...] = ("file_path", "notebook_path", "path")
"""Input keys built-in file tools use for their target."""

Decision = Literal["allow", "deny"]


@dataclass(frozen=True, slots=True)
class Deny:
    """Refusal with the reason the model reads as the tool result."""

    reason: str


@dataclass(frozen=True, slots=True)
class ToolCall:
    """One tool invocation as the hook sees it.

    ``state`` is a mutable per-session dict rules may write to (which
    skills were opened, what was already asked); it lives as long as the
    claude process and is empty again after a resume.
    """

    tool: str
    input: Mapping[str, Any]
    agent: str
    kind: str
    conversation: str
    cwd: Path
    state: MutableMapping[str, Any] = field(default_factory=dict)

    @property
    def mcp(self) -> tuple[str, str] | None:
        """``(server, tool)`` for ``mcp__<server>__<tool>`` names, else ``None``."""
        if not self.tool.startswith("mcp__"):
            return None
        _, _, rest = self.tool.partition("__")
        server, sep, name = rest.partition("__")
        return (server, name) if sep else (server, "")

    def path(self, *keys: str) -> Path | None:
        """Target of a file tool, absolute; relative paths resolve against ``cwd``."""
        for key in keys or PATH_KEYS:
            raw = self.input.get(key)
            if isinstance(raw, str) and raw:
                return self.resolve(raw)
        return None

    def resolve(self, raw: str) -> Path:
        path = Path(raw).expanduser()
        if not path.is_absolute():
            path = self.cwd / path
        # ``resolve`` would follow symlinks on the gateway host, which may
        # not be the model's view; normalise lexically instead.
        return Path(*_normalize(path.parts))


def _normalize(parts: Iterable[str]) -> list[str]:
    out: list[str] = []
    for part in parts:
        if part == "..":
            if len(out) > 1:
                out.pop()
        elif part != ".":
            out.append(part)
    return out


type PolicyRule = Callable[[ToolCall], Deny | None | Awaitable[Deny | None]]
"""A rule; sync or async. First ``Deny`` wins, ``None`` passes to the next."""


@dataclass(frozen=True, slots=True)
class ToolAudit:
    agent: str
    conversation: str
    kind: str
    tool: str
    decision: Decision
    reason: str | None
    brief: str


AuditSink = "Callable[[ToolAudit], Awaitable[None]]"


async def evaluate(
    rules: Iterable[Callable[[ToolCall], Any]], call: ToolCall
) -> Deny | None:
    """Run ``rules`` in order; the first deny wins, an exception is a deny."""
    for rule in rules:
        try:
            verdict = rule(call)
            if inspect.isawaitable(verdict):
                verdict = await verdict
        except Exception:  # noqa: BLE001 - a broken rule must fail closed
            name = getattr(rule, "__name__", repr(rule))
            _log.exception("policy rule %s failed on %s", name, call.tool)
            return Deny(reason=f"policy rule {name} failed; the call is refused")
        if isinstance(verdict, Deny):
            return verdict
    return None


def hook_output(deny: Deny | None) -> dict[str, Any]:
    """``PreToolUse`` hook JSON for a verdict: empty means allow."""
    if deny is None:
        return {}
    return {
        "hookSpecificOutput": {
            "hookEventName": "PreToolUse",
            "permissionDecision": "deny",
            "permissionDecisionReason": deny.reason,
        }
    }


def brief(tool_input: Mapping[str, Any], *, limit: int = 240) -> str:
    """Short, log-safe summary of a call: the path, the command, or the args."""
    for key in (*PATH_KEYS, "command", "skill", "pattern", "query", "url"):
        value = tool_input.get(key)
        if isinstance(value, str) and value:
            return _cut(value, limit)
    try:
        return _cut(json.dumps(tool_input, ensure_ascii=False, sort_keys=True), limit)
    except (TypeError, ValueError):
        return _cut(repr(tool_input), limit)


def _cut(text: str, limit: int) -> str:
    return text if len(text) <= limit else text[: limit - 1] + "…"