Worktree Guard

Event

PreToolUse

Trigger

tool-name — Bash

Language

python

Side effects

• blocks new-branch creation (checkout -b / switch -c) in the primary worktree with worktree-add guidance
• blocks non-canonical worktree placement with path guidance

Platforms

linuxmacoswindows

Notes

Logic is cross-platform. Wiring: use ‘ai hooks run worktree-guard’ in settings.json — the ai binary discovers Python on each OS. Inspects git command strings via Python. Path separators handled by pathlib. Works on all platforms.

Depends on

hook/lib

Script · worktree-guard.py

#!/usr/bin/env python3
"""hooks/worktree-guard.py — enforce worktree-based feature development.

Two rules, both implementing Common.md §U17:

1. BRANCH-IN-PRIMARY GUARD
   Block `git checkout -b <branch>` and `git switch -c <branch>` when
   running in the primary worktree. Feature work must live in a linked
   worktree so the primary repo stays on main. Redirect to
   `git worktree add`.

2. PLACEMENT GUARD
   When `git worktree add <path>` IS used, enforce canonical placement:
     <repo>/.worktrees/<name>/         (single-repo, dies with the repo)
     <repo>/.claude/worktrees/<name>/  (Claude Code native EnterWorktree)
     ~/.ai/worktrees/<name>/           (cross-repo, persistent)
   Ad-hoc placement (../<branch>/, /tmp/worktree-X/, etc.) is forbidden.

Known limitations (PreToolUse static parser): a `git checkout -b` or
`git worktree add` hidden inside a subshell (`(...)`), command
substitution (`$(...)` / backticks), or a delegated interpreter
(`bash -c "..."`) is not caught by static tokenization. The
client-agnostic wrapper mode (`ai wrap git`) is the backstop for those
— it inspects the real git argv.

Two invocation modes:
  PreToolUse  — reads a JSON event from stdin; on violation emits a
                JSON permissionDecision deny on stdout and exits 0.
                (Claude Code path via hooks.PreToolUse in settings.json)
  wrapper     — invoked as `python <hook> --mode=wrapper` by `ai wrap git`;
                reads git argv from $WRAPPED_ARGV; on violation writes to
                stderr and exits non-zero to block the real git binary.

Self-check:
  --self-check
"""
from __future__ import annotations

import json
import os
import shlex
import subprocess
import sys
from pathlib import Path
from typing import List, Optional

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _lib  # noqa: E402

# Flags on `git worktree add` that consume the next token as a value.
VALUE_FLAGS = {"-b", "-B", "--reason"}

# Shell operators that separate one command from the next.
SHELL_SEPARATORS = {";", "&&", "||", "|", "&"}

# `git` global options (before the subcommand) that consume the next token.
GIT_GLOBAL_VALUE_OPTS = {
    "-C", "-c", "--git-dir", "--work-tree",
    "--namespace", "--exec-path", "--super-prefix",
}


def deny_pretooluse(reason: str) -> None:
    """Emit a PreToolUse permission-deny decision (stdout) and exit 0."""
    print(json.dumps({
        "hookSpecificOutput": {
            "hookEventName": "PreToolUse",
            "permissionDecision": "deny",
            "permissionDecisionReason": reason,
        }
    }))
    sys.exit(0)


def ai_worktrees_root() -> Path:
    """Return the cross-repo canonical worktree root."""
    ai_root = os.environ.get("AI_ROOT", str(Path.home() / ".ai"))
    return Path(ai_root) / "worktrees"


def _clean_path_env() -> dict:
    """Return os.environ with AI shim bin dirs stripped from PATH.

    resolve_repo_root and is_primary_worktree must reach the real git,
    not the ~/.ai/bin/git governance shim, to avoid re-entering the
    `ai wrap` pipeline which can deadlock or timeout.
    """
    canonical_ai_bin = Path.home() / ".ai" / "bin"
    env_ai_root = os.environ.get("AI_ROOT", "")
    override_ai_bin = Path(env_ai_root) / "bin" if env_ai_root else None

    strip = {str(canonical_ai_bin.resolve())}
    if override_ai_bin is not None:
        try:
            strip.add(str(override_ai_bin.resolve()))
        except Exception:
            strip.add(str(override_ai_bin))

    raw_path = os.environ.get("PATH", "")
    cleaned_parts = []
    for p in raw_path.split(os.pathsep):
        if not p:
            continue
        try:
            resolved = str(Path(p).resolve())
        except Exception:
            resolved = p
        if resolved not in strip and p not in strip:
            cleaned_parts.append(p)
    env = dict(os.environ)
    if cleaned_parts:
        env["PATH"] = os.pathsep.join(cleaned_parts)
    return env


def resolve_repo_root(cwd: str) -> Optional[Path]:
    """Return the repo root for the given working directory, or None."""
    try:
        out = subprocess.run(
            ["git", "-C", cwd, "rev-parse", "--show-toplevel"],
            capture_output=True, text=True, timeout=5,
            env=_clean_path_env(),
        )
        if out.returncode == 0 and out.stdout.strip():
            return Path(out.stdout.strip()).resolve()
    except Exception:
        pass
    return None


def is_primary_worktree(cwd: str) -> bool:
    """Return True if cwd is inside the primary (main) worktree.

    In a linked worktree, --git-dir and --git-common-dir differ. In the
    primary worktree they resolve to the same path.
    """
    try:
        git_dir = subprocess.run(
            ["git", "-C", cwd, "rev-parse", "--git-dir"],
            capture_output=True, text=True, timeout=5,
            env=_clean_path_env(),
        )
        common_dir = subprocess.run(
            ["git", "-C", cwd, "rev-parse", "--git-common-dir"],
            capture_output=True, text=True, timeout=5,
            env=_clean_path_env(),
        )
        if git_dir.returncode != 0 or common_dir.returncode != 0:
            return False
        gd = Path(git_dir.stdout.strip()).resolve()
        cd = Path(common_dir.stdout.strip()).resolve()
        return gd == cd
    except Exception:
        return False


def _normalize_tokens(tokens: List[str]) -> List[str]:
    """Expand tokens that have trailing semicolons glued on."""
    result: List[str] = []
    for tok in tokens:
        if tok.endswith(";") and tok != ";":
            result.append(tok[:-1])
            result.append(";")
        else:
            result.append(tok)
    return result


def _split_segments(tokens: List[str]) -> List[List[str]]:
    """Split a flat token list into command segments on shell operators."""
    segments: List[List[str]] = []
    current: List[str] = []
    for tok in _normalize_tokens(tokens):
        if tok in SHELL_SEPARATORS:
            if current:
                segments.append(current)
                current = []
        else:
            current.append(tok)
    if current:
        segments.append(current)
    return segments


def _is_env_assignment(tok: str) -> bool:
    """True for a leading `VAR=value` shell environment assignment."""
    if tok.startswith("-") or "=" not in tok:
        return False
    return tok.split("=", 1)[0].isidentifier()


def _skip_git_preamble(seg: List[str], start: int) -> int:
    """Skip env assignments and git global options; return index of subcommand."""
    i, n = start, len(seg)
    while i < n and _is_env_assignment(seg[i]):
        i += 1
    if i >= n or seg[i] != "git":
        return n  # not a git command
    i += 1
    while i < n and seg[i].startswith("-"):
        opt = seg[i]
        if "=" in opt:
            i += 1
        elif opt in GIT_GLOBAL_VALUE_OPTS and i + 1 < n:
            i += 2
        else:
            i += 1
    return i


# ── Rule 1: new-branch-in-primary detection ──────────────────────────────────

def _is_new_branch_in_segment(seg: List[str]) -> bool:
    """True if this segment creates a new git branch."""
    i = _skip_git_preamble(seg, 0)
    if i >= len(seg):
        return False
    subcmd = seg[i]
    flags = set(seg[i + 1:])
    if subcmd == "checkout" and (flags & {"-b", "-B"}):
        return True
    if subcmd == "switch" and (flags & {"-c", "-C"}):
        return True
    return False


def _has_new_branch_command(tokens: List[str]) -> bool:
    """True if any segment in tokens creates a new git branch."""
    return any(_is_new_branch_in_segment(seg) for seg in _split_segments(tokens))


def new_branch_deny_message(cwd: str) -> str:
    """Build the Rule 1 deny message with worktree-add guidance."""
    repo_root = resolve_repo_root(cwd)
    ai_root = ai_worktrees_root()
    hints = []
    if repo_root:
        hints.append(
            f"  git worktree add {repo_root}/.worktrees/<name> -b <branch>"
            "  (single-repo)"
        )
    hints.append(
        f"  git worktree add {ai_root}/<name> -b <branch>"
        "  (cross-repo / persistent)"
    )
    return (
        "Creating a branch in the primary worktree violates Common.md §U17.\n"
        "Feature work must live in a linked worktree so the primary repo "
        "stays on main.\n"
        "\nUse a worktree instead:\n"
        + "\n".join(hints) + "\n"
        "\nOr: `ai worktree add <name>` (with `--global` for cross-repo)."
    )


# ── Rule 2: placement detection ───────────────────────────────────────────────

def _path_after_add(seg: List[str], i: int) -> Optional[str]:
    """Return the positional path argument following `worktree add`."""
    n = len(seg)
    while i < n:
        tok = seg[i]
        if tok == "--":
            i += 1
            break
        if not tok.startswith("-"):
            return tok
        if tok in VALUE_FLAGS and i + 1 < n:
            i += 2
            continue
        i += 1
    if i < n:
        return seg[i]
    return None


def _worktree_add_path_in_segment(seg: List[str]) -> Optional[str]:
    """Return the target path if this segment is `git worktree add`."""
    i = _skip_git_preamble(seg, 0)
    n = len(seg)
    if i + 1 >= n or seg[i] != "worktree" or seg[i + 1] != "add":
        return None
    return _path_after_add(seg, i + 2)


def _worktree_add_paths(tokens: List[str]) -> List[str]:
    """Every `git worktree add` target path across the token segments."""
    paths: List[str] = []
    for seg in _split_segments(tokens):
        path = _worktree_add_path_in_segment(seg)
        if path is not None:
            paths.append(path)
    return paths


def _worktree_add_paths_in_command(command: str) -> List[str]:
    """All `git worktree add` paths in a raw shell command string."""
    paths: List[str] = []
    for line in command.splitlines():
        line_tokens = shlex.split(line, comments=False, posix=True)
        paths.extend(_worktree_add_paths(line_tokens))
    return paths


def is_canonical(raw_path: str, cwd: str) -> bool:
    """Return True if raw_path resolves to a canonical worktree location."""
    target = Path(raw_path)
    if not target.is_absolute():
        target = Path(cwd) / target
    try:
        resolved = target.resolve()
    except Exception:
        resolved = Path(os.path.normpath(str(target)))

    ai_root = ai_worktrees_root()
    try:
        resolved.relative_to(ai_root.resolve())
        return True
    except (ValueError, FileNotFoundError):
        pass

    repo_root = resolve_repo_root(cwd)
    if repo_root is not None:
        try:
            resolved.relative_to(repo_root / ".worktrees")
            return True
        except ValueError:
            pass
        try:
            # Claude Code's native EnterWorktree tool places worktrees here;
            # it can't be redirected to `ai worktree add`, so this harness-
            # managed root is canonical too.
            resolved.relative_to(repo_root / ".claude" / "worktrees")
            return True
        except ValueError:
            pass

    return False


def placement_deny_message(raw_path: str, cwd: str) -> str:
    """Build the Rule 2 deny message."""
    p = Path(raw_path)
    if not p.is_absolute():
        p = Path(cwd) / p
    try:
        shown = p.resolve()
    except Exception:
        shown = p

    repo_root = resolve_repo_root(cwd)
    ai_root = ai_worktrees_root()
    hints = []
    if repo_root:
        hints.append(f"  {repo_root}/.worktrees/<name>/  — single-repo")
        hints.append(f"  {repo_root}/.claude/worktrees/<name>/  — Claude Code native (EnterWorktree)")
    hints.append(f"  {ai_root}/<name>/  — cross-repo or persistent")
    return (
        "Worktree placement violates Common.md §U17.\n"
        f"  Target: {shown}\n"
        "Canonical roots:\n"
        + "\n".join(hints) + "\n"
        "Choose by lifecycle (§U17.1) and re-run.\n"
        "Preferred surface: `ai worktree add <name>` "
        "(with `--global` for cross-repo)."
    )


# ── Mode runners ──────────────────────────────────────────────────────────────

def run_pretooluse_mode() -> int:
    """Claude Code PreToolUse path: parse a JSON event from stdin."""
    raw = sys.stdin.read()
    if not raw.strip():
        return 0
    try:
        event = json.loads(raw)
    except json.JSONDecodeError:
        return 0

    hook_event = event.get("hookEventName") or event.get("hook_event_name") or ""
    if hook_event != "PreToolUse":
        return 0
    tool_name = event.get("tool_name") or event.get("toolName") or ""
    if tool_name not in ("Bash", "shell", "execute"):
        return 0
    tool_input = event.get("tool_input") or event.get("toolInput") or {}
    if not isinstance(tool_input, dict):
        return 0
    command = tool_input.get("command") or ""
    if not isinstance(command, str):
        return 0

    has_branch_op = "checkout" in command or "switch" in command
    has_worktree = "worktree" in command
    if not has_branch_op and not has_worktree:
        return 0

    cwd = event.get("cwd") or os.getcwd()

    # Rule 1: block new-branch creation in the primary worktree.
    if has_branch_op and is_primary_worktree(cwd):
        for line in command.splitlines():
            try:
                line_tokens = shlex.split(line, comments=False, posix=True)
            except ValueError:
                continue
            if _has_new_branch_command(line_tokens):
                deny_pretooluse(new_branch_deny_message(cwd))

    # Rule 2: enforce canonical placement when `git worktree add` is used.
    if has_worktree:
        try:
            targets = _worktree_add_paths_in_command(command)
        except ValueError:
            deny_pretooluse(
                "worktree-guard: could not parse a command containing "
                "'worktree'; denying per fail-closed policy (Common.md §U17)."
            )
            return 0
        for target in targets:
            if not is_canonical(target, cwd):
                deny_pretooluse(placement_deny_message(target, cwd))

    return 0


def run_wrapper_mode() -> int:
    """Client-agnostic `ai wrap git` path: git argv arrives in $WRAPPED_ARGV."""
    argv_json = os.environ.get("WRAPPED_ARGV", "")
    if not argv_json:
        return 0
    try:
        argv = json.loads(argv_json)
    except json.JSONDecodeError:
        return 0
    if not isinstance(argv, list):
        return 0

    tokens = ["git"] + [str(a) for a in argv]
    cwd = os.getcwd()

    # Rule 1: block new-branch creation in the primary worktree.
    if ("checkout" in tokens or "switch" in tokens):
        if _has_new_branch_command(tokens) and is_primary_worktree(cwd):
            sys.stderr.write(new_branch_deny_message(cwd) + "\n")
            return 1

    # Rule 2: enforce canonical placement.
    if "worktree" in tokens:
        for target in _worktree_add_paths(tokens):
            if not is_canonical(target, cwd):
                sys.stderr.write(placement_deny_message(target, cwd) + "\n")
                return 1

    return 0


def main() -> int:
    if "--self-check" in sys.argv:
        return _lib.self_check_ok()
    if "--mode=wrapper" in sys.argv or os.environ.get("WRAPPED_CMD"):
        return run_wrapper_mode()
    return run_pretooluse_mode()


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