legal-ai/scripts/agent_tool_grants_guard.py

#!/usr/bin/env python3
"""INV-AG3 guard — every MCP tool an agent is TOLD to run must be GRANTED to it.

The canonical checker for INV-AG3 (docs/spec/X4-agents.md §2א): a Claude-Code
subagent's ``tools:`` frontmatter is a CLOSED allow-list. A tool that is
registered on the MCP server but absent from that list is *not callable* by the
agent, however well the server is connected.

Why this exists — the failure it is built to catch (2026-08-04):
``analyze_protocol`` shipped on 2026-06-30 (24e3e2f) touching 9 files, none of
them ``.claude/agents/*``. Later ``wake_analyst_for_protocol_analysis``
(web/paperclip_client.py, #226) started writing "הרץ
``mcp__legal-ai__analyze_protocol(...)``" straight into the analyst's issue.
The analyst therefore received an explicit instruction to run a tool it was
never granted, reported "tools exist on the connected server but aren't exposed
as callable in this session", and burned two runs working around it via raw
psql + a hand-written script. INV-AG3 already covered this on paper; its
enforcement was deferred ("אכיפה אוטומטית עתידית"), so the drift went unnoticed
for five weeks. This script is that deferred enforcement.

Three HARD rules:

  1. **Backend delegation.** Every ``mcp__legal-ai__X`` named inside ``web/``
     (the backend telling an agent what to run) must be granted to at least one
     agent. This is the rule that catches the 2026-08-04 failure.

  2. **Per-agent instructions.** Every ``mcp__legal-ai__X`` in an agent file's
     BODY must be granted in that same file's frontmatter. Prefixed mentions are
     imperative by convention ("הרץ `mcp__legal-ai__…`").

  3. **No phantom grants.** Every granted tool must actually be registered on
     the MCP server — catches typos and tools deleted out from under an agent.

Plus one reviewed-exception rule:

  4. **Bare tool names.** An agent body may name a tool in backticks without the
     ``mcp__legal-ai__`` prefix (```get_legal_arguments```). Those are
     ambiguous: some are real requirements, others are deliberately contrastive
     ("**לא** דרך `precedent_library_upload`"), a pointer at *another* agent's
     job, or a DB column that merely shares a tool's name. Each is classified
     once in ``CONTRASTIVE_OK`` below; anything new fails until reviewed.

And one host-only rule, opt-in via ``--check-bindings``:

  5. **Bindings.** Rules 1–4 compare files to files, which says nothing about
     whether an allow-list is *enforced*. It is only enforced when the runtime
     selects that agent (``--agent <name>``); without the flag the same file is
     delivered as ``--append-system-prompt-file`` — prose, not a gate — and every
     tool stays reachable. Found on 2026-08-05: one agent declared 41 grants with
     no ``--agent`` flag, so the largest allow-list in the system was inert while
     this guard reported OK. The binding lives in the platform DB, so CI cannot
     see it; without the flag the guard now says so out loud instead of implying
     enforcement it never verified.

NOT AGENTS (no frontmatter by design — the adapter sends the file as a raw
prompt, so YAML would leak into it): ``hermes-curator.md`` (deepseek_local),
``legal-analyst-gemini-critique.md`` (gemini_local). ``HEARTBEAT.md`` is a
shared checklist, not an agent. All three are skipped.

Usage:
    agent_tool_grants_guard.py    # exit 1 on any violation
"""
from __future__ import annotations

import json
import os
import subprocess
import re
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent

AGENTS_DIR = REPO / ".claude" / "agents"
MCP_SRC = REPO / "mcp-server" / "src"
BACKEND_DIR = REPO / "web"

# Files under .claude/agents/ that are not claude_local subagent definitions.
NOT_AGENTS = {
    "HEARTBEAT.md",
    "hermes-curator.md",
    "legal-analyst-gemini-critique.md",
}

# Bare (unprefixed) tool names in an agent body that are NOT requirements.
# Each entry is (agent file, tool, why) — reviewed 2026-08-04. Adding to this
# map is a deliberate act: it asserts "the agent is not being told to call this".
CONTRASTIVE_OK = {
    ("legal-analyst.md", "case_create"): "prose about the cases.practice_area CHECK constraint, not a call",
    ("legal-analyst.md", "search_internal_decisions"): "names the filter surface when contrasting Axis A/B",
    ("legal-ceo.md", "search_decisions"): "contrast — 'search_decisions = only Dafna' vs the granted search_internal_decisions",
    ("legal-ceo.md", "precedent_library_upload"): "explicitly the forbidden path ('לא דרך …', citation guard rejects)",
    ("legal-ceo.md", "document_update"): "describes the tagging chaim must fix, not a CEO call",
    ("legal-proofreader.md", "extraction_status"): "the documents.extraction_status DB column — name collides with a tool",
    ("legal-qa.md", "precedent_attach"): "explicitly the researcher's job ('דרך precedent_attach של ה-researcher')",
    ("legal-writer.md", "revise_draft"): "the CEO calls it ('CEO יקרא ל-revise_draft'), not the writer",
    ("legal-writer.md", "search_case_precedents"): "a do-not-confuse disambiguation note ('שונה! … לא לבלבל')",
}

TOOL_RE = re.compile(r"mcp__legal-ai__(\w+)")
REGISTER_RE = re.compile(r"@mcp\.tool\([^)]*\)\s*(?:async\s+)?def\s+(\w+)")
# `tool_name(` or `tool_name` inside backticks.
BARE_RE = re.compile(r"`(\w+)[(`]")


def server_tools() -> set[str]:
    """Tool names registered on the MCP server."""
    out: set[str] = set()
    for path in MCP_SRC.rglob("*.py"):
        out |= set(REGISTER_RE.findall(path.read_text(encoding="utf-8", errors="ignore")))
    return out


def split_frontmatter(text: str) -> tuple[str, str]:
    """Return (frontmatter, body). Empty frontmatter when the file has none."""
    if not text.startswith("---"):
        return "", text
    parts = text.split("---")
    if len(parts) < 3:
        return "", text
    return parts[1], "---".join(parts[2:])


def agent_files() -> list[Path]:
    return sorted(p for p in AGENTS_DIR.glob("*.md") if p.name not in NOT_AGENTS)


def check_bindings() -> list[tuple[str, str]]:
    """Return [(agent file, why)] for agents whose allow-list nothing enforces.

    A ``tools:`` list is only an allow-list when the runtime is told which agent
    to be. The local adapter enforces it under ``--agent <name>``; without that
    flag the very same file is delivered as ``--append-system-prompt-file``, i.e.
    prose the model may follow or ignore, and every tool stays reachable.

    Found the hard way on 2026-08-05: one agent carried 41 grants and no
    ``--agent`` flag, so the largest allow-list in the system was inert — and
    this guard had been reporting OK on it, because Rules 1–4 only ever compare
    files to files.

    Host-only. The binding lives in the platform's database, which CI cannot
    reach, so this shells out to ``psql`` rather than adding a driver dependency
    that would break the stdlib-only property the CI path relies on. Returns []
    when the database is unreachable — an unreachable DB is "not checked", not
    "no violations", and the caller prints that distinction.
    """

    sql = (
        "select adapter_config->>'instructionsEntryFile', "
        "coalesce(adapter_config->>'extraArgs','') "
        "from agents where adapter_type='claude_local' "
        "and adapter_config->>'instructionsEntryFile' is not null;"
    )
    try:
        out = subprocess.run(
            ["psql", "-h", "localhost", "-p", "54329", "-U", "paperclip",
             "-d", "paperclip", "-X", "-A", "-t", "-F", "\t", "-c", sql],
            capture_output=True, text=True, timeout=20,
            env={**os.environ, "PGPASSWORD": os.environ.get("PGPASSWORD", "paperclip")},
        )
    except (OSError, subprocess.SubprocessError):
        return []
    if out.returncode != 0:
        return []

    bad: dict[str, str] = {}
    for line in out.stdout.splitlines():
        if "\t" not in line:
            continue
        entry_file, extra = line.split("\t", 1)
        entry_file = entry_file.strip()
        if not entry_file or entry_file in NOT_AGENTS:
            continue
        want = entry_file[:-3] if entry_file.endswith(".md") else entry_file
        try:
            args = json.loads(extra) if extra.strip() else []
        except json.JSONDecodeError:
            args = []
        # Only a literal ["--agent", "<name>"] pair binds the allow-list.
        ok = any(
            a == "--agent" and i + 1 < len(args) and args[i + 1] == want
            for i, a in enumerate(args)
        )
        if not ok:
            bad[entry_file] = (
                "extraArgs is empty" if not args
                else f"extraArgs={extra.strip()} does not select '{want}'"
            )

    # Only report agents that actually declare grants — an agent with no tools:
    # list has nothing to enforce and is not a finding.
    result = []
    for path in agent_files():
        if path.name in bad:
            fm, _ = split_frontmatter(path.read_text(encoding="utf-8", errors="ignore"))
            if TOOL_RE.findall(fm):
                result.append((path.name, bad[path.name]))
    return sorted(result)


def main() -> int:
    registered = server_tools()
    if not registered:
        print("agent-tool-grants: FAIL — no @mcp.tool registrations found; is the tree complete?")
        return 1

    grants: dict[str, set[str]] = {}
    bodies: dict[str, str] = {}
    for path in agent_files():
        fm, body = split_frontmatter(path.read_text(encoding="utf-8", errors="ignore"))
        grants[path.name] = set(TOOL_RE.findall(fm))
        bodies[path.name] = body

    all_granted: set[str] = set().union(*grants.values()) if grants else set()
    violations: list[str] = []

    # Rule 1 — backend delegation must land on a granted tool.
    for path in sorted(BACKEND_DIR.rglob("*.py")):
        text = path.read_text(encoding="utf-8", errors="ignore")
        for tool in sorted(set(TOOL_RE.findall(text))):
            if tool not in all_granted:
                rel = path.relative_to(REPO)
                violations.append(
                    f"[1 backend] {rel} instructs an agent to run "
                    f"mcp__legal-ai__{tool}, but NO agent grants it.\n"
                    f"           fix: add `- mcp__legal-ai__{tool}` to the tools: "
                    f"frontmatter of the agent that receives that issue."
                )

    # Rule 2 — a prefixed mention in an agent body is an instruction to that agent.
    for name, body in bodies.items():
        for tool in sorted(set(TOOL_RE.findall(body))):
            if tool not in grants[name]:
                violations.append(
                    f"[2 instructions] .claude/agents/{name} tells the agent to run "
                    f"mcp__legal-ai__{tool}, which its own tools: list omits.\n"
                    f"           fix: add `- mcp__legal-ai__{tool}` to that frontmatter."
                )

    # Rule 3 — no grant may point at a tool the server does not register.
    for name, granted in grants.items():
        for tool in sorted(granted - registered):
            violations.append(
                f"[3 phantom] .claude/agents/{name} grants mcp__legal-ai__{tool}, "
                f"which is not registered on the MCP server.\n"
                f"           fix: correct the name, or drop the grant if the tool was removed."
            )

    # Rule 4 — every bare tool name is either granted or classified as contrastive.
    for name, body in bodies.items():
        bare = {m for m in BARE_RE.findall(body) if m in registered}
        for tool in sorted(bare - grants[name]):
            if (name, tool) in CONTRASTIVE_OK:
                continue
            violations.append(
                f"[4 bare name] .claude/agents/{name} mentions `{tool}` — a real MCP "
                f"tool it is not granted.\n"
                f"           fix: grant it if the agent must call it, otherwise add "
                f"(\"{name}\", \"{tool}\") to CONTRASTIVE_OK with the reason."
            )

    # Stale exceptions: a classification that no longer matches the text is noise.
    for (name, tool), _why in sorted(CONTRASTIVE_OK.items()):
        if name not in bodies:
            violations.append(
                f"[4 stale] CONTRASTIVE_OK names {name}, which is not an agent file."
            )
        elif tool not in {m for m in BARE_RE.findall(bodies[name])}:
            violations.append(
                f"[4 stale] CONTRASTIVE_OK ({name}, {tool}) no longer appears in that "
                f"file — drop the exception."
            )

    # Rule 5 — a grant list only binds if the runtime actually selects that agent.
    unenforced = check_bindings() if "--check-bindings" in sys.argv else None
    if unenforced:
        for name, detail in unenforced:
            violations.append(
                f"[5 binding] .claude/agents/{name} declares a tools: allow-list, but "
                f"the runtime does not select that agent — {detail}.\n"
                f"           The list is inert: it is delivered as prompt text only, so "
                f"every tool remains callable.\n"
                f"           fix: set adapter_config.extraArgs to "
                f'["--agent", "{name[:-3]}"], or drop tools: and document the agent as '
                f"unrestricted. Not both."
            )

    if violations:
        print(f"INV-AG3 agent-tool-grants guard: {len(violations)} violation(s)\n")
        for v in violations:
            print(f"  ✗ {v}")
        print(
            "\ndocs/spec/X4-agents.md §2א INV-AG3 — the frontmatter tools: list is a "
            "CLOSED allow-list.\nA tool missing from it is not callable, no matter that "
            "the MCP server is connected."
        )
        return 1

    print(
        f"INV-AG3 agent-tool-grants guard: OK "
        f"({len(agent_files())} agents, {len(all_granted)} distinct grants, "
        f"{len(registered)} tools registered)"
    )
    if unenforced is None:
        # Say plainly what was NOT checked. A guard that prints a bare "OK" invites
        # the reader to conclude the allow-lists are enforced; this one has only
        # compared files to files. Enforcement is a runtime property (see Rule 5),
        # and on 2026-08-05 exactly one agent was found declaring 41 grants that
        # nothing enforces — while this guard reported OK.
        print(
            "  note: file-level only. Whether each allow-list is actually ENFORCED "
            "depends on the\n  runtime passing --agent <name>, which needs the platform "
            "DB — re-run with --check-bindings\n  on the host to verify."
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())