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.

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 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 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."
            )

    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)"
    )
    return 0


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