Agent-System/.cursor/hooks/_lib.py

#!/usr/bin/env python3
"""Shared helpers for project Cursor hooks (advisory gates + audit/probe)."""
from __future__ import annotations

import json
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

KNOWN_ROLES = frozenset(
    {
        "research",
        "planning",
        "requirements_review",
        "plan_review",
        "code_review",
        "diagnosis",
        "test_run",
        "acceptance_check",
        "stage_notes",
        "report",
        "auxiliary",
    }
)
TEST_ROLES = frozenset({"test_run", "acceptance_check"})
FULL_SHA = re.compile(r"^[0-9a-fA-F]{40}$")
REQUEST_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")


def repo_root() -> Path:
    env = os.environ.get("CURSOR_PROJECT_DIR") or os.environ.get("CURSOR_WORKSPACE_ROOT")
    if env:
        return Path(env).resolve()
    # hooks run with cwd = project root
    return Path.cwd().resolve()


def audit_path() -> Path:
    return repo_root() / ".cursor" / "hook-audit.jsonl"


def probe_path() -> Path:
    return repo_root() / ".cursor" / "hook-probe.jsonl"


def inventory_dir() -> Path:
    return repo_root() / ".cursor" / "hook-audit"


def append_jsonl(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    record = dict(record)
    record.setdefault("ts", datetime.now(timezone.utc).isoformat())
    with path.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(record, ensure_ascii=False) + "\n")


def read_stdin() -> str:
    return sys.stdin.read()


def maybe_probe(raw: str, hook: str) -> None:
    if os.environ.get("CURSOR_HOOKS_PROBE", "").strip() not in ("1", "true", "TRUE", "yes"):
        return
    append_jsonl(
        probe_path(),
        {"hook": hook, "raw": raw[:65536], "raw_len": len(raw)},
    )


def parse_json_loose(raw: str) -> dict[str, Any] | None:
    raw = (raw or "").strip()
    if not raw:
        return None
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None
    return data if isinstance(data, dict) else None


def first_key(data: dict[str, Any], *names: str) -> Any:
    for name in names:
        if name in data and data[name] is not None:
            return data[name]
    return None


def normalize_tool_name(name: str | None) -> str | None:
    """Strip Cursor MCP prefixes such as 'MCP:' / 'MCP: pi-subagents/'."""
    if name is None:
        return None
    n = str(name).strip()
    # Observed IDE postToolUse: tool_name == "MCP:pi_subagent_roles"
    if n.upper().startswith("MCP:"):
        n = n[4:].lstrip()
    # Rare: "MCP: pi-subagents/pi_subagent_run"
    if "/" in n and n.rsplit("/", 1)[-1].startswith("pi_subagent_"):
        n = n.rsplit("/", 1)[-1]
    return n or None


def tool_name(data: dict[str, Any]) -> str | None:
    v = first_key(data, "tool_name", "mcp_tool_name", "name", "toolName")
    return normalize_tool_name(str(v) if v is not None else None)


def mcp_server(data: dict[str, Any]) -> str | None:
    v = first_key(data, "mcp_server_name", "server", "mcpServer", "server_name")
    return str(v) if v is not None else None


def tool_input(data: dict[str, Any]) -> dict[str, Any] | None:
    v = first_key(data, "tool_input", "arguments", "input", "toolInput", "params")
    if isinstance(v, dict):
        return v
    if isinstance(v, str):
        try:
            parsed = json.loads(v)
            return parsed if isinstance(parsed, dict) else None
        except json.JSONDecodeError:
            return None
    return None


def shell_command(data: dict[str, Any]) -> str:
    v = first_key(data, "command", "shell_command", "cmd")
    return str(v) if v is not None else ""


def emit(obj: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(obj, ensure_ascii=False))
    sys.stdout.flush()


def allow() -> None:
    emit({"permission": "allow"})
    raise SystemExit(0)


def ask(user_message: str, agent_message: str) -> None:
    emit(
        {
            "permission": "ask",
            "user_message": user_message,
            "agent_message": agent_message,
        }
    )
    raise SystemExit(0)


def additional_context(text: str) -> None:
    emit({"additional_context": text})
    raise SystemExit(0)


def audit(hook: str, code: str, **extra: Any) -> None:
    rec = {"hook": hook, "code": code}
    rec.update(extra)
    try:
        append_jsonl(audit_path(), rec)
    except OSError:
        pass