@trac3r/oh-my-god 2.2.11

package/hooks/_protected_context.py

@@ -0,0 +1,150 @@
+"""Protected context registry for PreCompact hook.
+
+Reads .claude-context-protect entries (file paths, regex patterns, literal strings)
+and collects protected context items to re-inject via additionalContext during compaction.
+
+Default protections (when no .claude-context-protect exists):
+- CLAUDE.md content
+- Active task definitions (## Task:, - [ ])
+- Recent error messages (Error:, Exception:, FAILED)
+
+Pure stdlib — no external dependencies.
+"""
+import os
+import re
+
+
+PROTECT_FILE_NAME = ".claude-context-protect"
+
+# Default protection patterns (used when no protect file exists)
+_DEFAULT_TASK_PATTERNS = [
+    re.compile(r"^## Task:"),
+    re.compile(r"^- \[ \]"),
+]
+_DEFAULT_ERROR_KEYWORDS = ("Error:", "Exception:", "FAILED")
+
+
+def load_protect_entries(project_dir):
+    """Read .claude-context-protect file, return list of entries or None if missing.
+
+    Returns:
+        list[str] | None: List of non-empty, non-comment lines. None if file missing.
+    """
+    protect_path = os.path.join(project_dir, PROTECT_FILE_NAME)
+    if not os.path.isfile(protect_path):
+        return None
+
+    try:
+        with open(protect_path, "r", encoding="utf-8", errors="ignore") as f:
+            lines = f.readlines()
+    except Exception:
+        return None
+
+    entries = []
+    for line in lines:
+        stripped = line.strip()
+        if stripped and not stripped.startswith("#"):
+            entries.append(stripped)
+    return entries
+
+
+def _read_file_content(file_path):
+    """Read file content. Returns stripped string or None on failure."""
+    try:
+        if not os.path.isfile(file_path):
+            return None
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            content = f.read().strip()
+        return content if content else None
+    except Exception:
+        return None
+
+
+def _match_entry_against_lines(entry, context_lines):
+    """Match entry against context lines. Tries regex first, falls back to literal.
+
+    Returns:
+        list[str]: Matching lines.
+    """
+    try:
+        pattern = re.compile(entry)
+        return [line for line in context_lines if pattern.search(line)]
+    except re.error:
+        # Invalid regex — fall back to literal substring match
+        return [line for line in context_lines if entry in line]
+
+
+def _process_entry(entry, project_dir, context_lines):
+    """Process a single protect entry. Returns list of protected strings.
+
+    Resolution order:
+    1. If entry resolves to an existing file → include file content
+    2. Otherwise, try regex match against context lines
+    3. If regex fails (re.error), fall back to literal substring match
+    """
+    # 1. Try as file path
+    file_path = os.path.join(project_dir, entry)
+    content = _read_file_content(file_path)
+    if content is not None:
+        return [content]
+
+    # 2. Try as regex/literal against context lines
+    return _match_entry_against_lines(entry, context_lines)
+
+
+def _get_default_protections(project_dir, context_lines):
+    """Apply default protections when no .claude-context-protect exists.
+
+    Default protected items:
+    - CLAUDE.md content (if file exists)
+    - Active task definitions (## Task:, - [ ])
+    - Recent error messages (Error:, Exception:, FAILED)
+    """
+    parts = []
+
+    # 1. CLAUDE.md content
+    claude_md_path = os.path.join(project_dir, "CLAUDE.md")
+    claude_content = _read_file_content(claude_md_path)
+    if claude_content:
+        parts.append(claude_content)
+
+    # 2. Active task definitions
+    for line in context_lines:
+        for pat in _DEFAULT_TASK_PATTERNS:
+            if pat.search(line):
+                parts.append(line)
+                break
+
+    # 3. Recent error messages
+    for line in context_lines:
+        if any(kw in line for kw in _DEFAULT_ERROR_KEYWORDS):
+            parts.append(line)
+
+    return parts
+
+
+def collect_protected_context(project_dir, context_text=""):
+    """Collect all protected context items and return as a single string.
+
+    Args:
+        project_dir: Project root directory.
+        context_text: Current context text to scan for regex/literal matches.
+
+    Returns:
+        str: Protected context items joined by newlines. Empty string if nothing.
+    """
+    context_lines = [l for l in context_text.split("\n") if l.strip()] if context_text else []
+    protected_parts = []
+
+    entries = load_protect_entries(project_dir)
+
+    if entries is None:
+        # No protect file — use defaults
+        protected_parts = _get_default_protections(project_dir, context_lines)
+    else:
+        # Process each entry from protect file
+        for entry in entries:
+            matched = _process_entry(entry, project_dir, context_lines)
+            protected_parts.extend(matched)
+
+    return "\n".join(protected_parts) if protected_parts else ""

package/hooks/_token_counter.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python3
+"""Tiered token estimation helpers for OMG hooks."""
+from __future__ import annotations
+
+import json
+import os
+import importlib
+import urllib.request
+from collections.abc import Iterable
+
+API_URL = "https://api.anthropic.com/v1/messages/count_tokens"
+API_MODEL = "claude-3-5-haiku-20241022"
+
+_FEATURE_UI_DISPLAY = "ui_display"
+_FEATURE_BUDGET_ENFORCEMENT = "budget_enforcement"
+_FEATURE_PREFLIGHT = "preflight"
+
+
+def _safe_int(value: float) -> int:
+    if value <= 0:
+        return 0
+    return int(value)
+
+
+def _extract_features(text: str) -> tuple[int, int, int]:
+    encoded = text.encode("utf-8")
+    byte_count = len(encoded)
+    word_count = len(text.split())
+    line_count = text.count("\n") + (1 if text else 0)
+    return byte_count, word_count, line_count
+
+
+def _default_coefficients() -> tuple[float, float, float, float]:
+    return (1.0, 0.19, 0.75, 1.1)
+
+
+def _gaussian_solve(matrix: list[list[float]], vector: list[float]) -> list[float] | None:
+    n = len(vector)
+    if n == 0:
+        return None
+    try:
+        for i in range(n):
+            pivot = i
+            for r in range(i + 1, n):
+                if abs(matrix[r][i]) > abs(matrix[pivot][i]):
+                    pivot = r
+            if abs(matrix[pivot][i]) < 1e-12:
+                return None
+
+            if pivot != i:
+                matrix[i], matrix[pivot] = matrix[pivot], matrix[i]
+                vector[i], vector[pivot] = vector[pivot], vector[i]
+
+            pivot_val = matrix[i][i]
+            for c in range(i, n):
+                matrix[i][c] /= pivot_val
+            vector[i] /= pivot_val
+
+            for r in range(n):
+                if r == i:
+                    continue
+                factor = matrix[r][i]
+                if factor == 0:
+                    continue
+                for c in range(i, n):
+                    matrix[r][c] -= factor * matrix[i][c]
+                vector[r] -= factor * vector[i]
+        return vector
+    except Exception:
+        return None
+
+
+def _fit_linear_coefficients(samples: Iterable[tuple[str, int]]) -> tuple[float, float, float, float]:
+    rows: list[list[float]] = []
+    targets: list[float] = []
+    for text, tokens in samples:
+        bcount, wcount, lcount = _extract_features(text)
+        rows.append([1.0, float(bcount), float(wcount), float(lcount)])
+        targets.append(float(tokens))
+
+    dim = 4
+    if len(rows) < dim:
+        return _default_coefficients()
+
+    xtx = [[0.0 for _ in range(dim)] for _ in range(dim)]
+    xty = [0.0 for _ in range(dim)]
+    for row, target in zip(rows, targets):
+        for i in range(dim):
+            xty[i] += row[i] * target
+            for j in range(dim):
+                xtx[i][j] += row[i] * row[j]
+
+    solved = _gaussian_solve(xtx, xty)
+    if solved is None:
+        return _default_coefficients()
+
+    return (float(solved[0]), float(solved[1]), float(solved[2]), float(solved[3]))
+
+
+_CALIBRATION_SAMPLES: tuple[tuple[str, int], ...] = (
+    ("ls", 5),
+    ("git status", 7),
+    ("echo 'hello world'", 10),
+    ("python3 -m pytest tests/hooks/test_feature_flags_v2.py -q", 18),
+    ("def hello():\n    return 'world'\n", 24),
+    (
+        """from pathlib import Path\nfor path in Path('hooks').glob('*.py'):\n    print(path.name)\n""",
+        44,
+    ),
+    (
+        """def estimate_tokens(text: str, tier: int = 1) -> int:\n    if tier == 1:\n        return max(1, int(len(text) / 3.5))\n    return 0\n""",
+        80,
+    ),
+    (
+        "\n".join(["line with common code and comments" for _ in range(80)]),
+        720,
+    ),
+    (
+        "\n".join(["longer source line with punctuation () {} [] == != <= >=" for _ in range(250)]),
+        1800,
+    ),
+)
+
+_COEFFICIENTS = _fit_linear_coefficients(_CALIBRATION_SAMPLES)
+
+
+def _estimate_tier1(text: str) -> int:
+    if not text:
+        return 0
+    return max(1, int(len(text) / 3.5))
+
+
+def _estimate_tier2(text: str) -> int:
+    if not text:
+        return 0
+    bias, w_bytes, w_words, w_lines = _COEFFICIENTS
+    bcount, wcount, lcount = _extract_features(text)
+    prediction = bias + (w_bytes * bcount) + (w_words * wcount) + (w_lines * lcount)
+    return max(1, _safe_int(prediction))
+
+
+def _get_anthropic_api_key() -> str | None:
+    try:
+        store_mod = importlib.import_module("credential_store")
+        key = store_mod.get_active_key("anthropic")
+        if key:
+            return key
+    except (ImportError, RuntimeError, ValueError, OSError, AttributeError):
+        pass
+    return os.environ.get("ANTHROPIC_API_KEY")
+
+
+def _estimate_tier3(text: str) -> int:
+    if not text:
+        return 0
+
+    api_key = _get_anthropic_api_key()
+    if not api_key:
+        return _estimate_tier2(text)
+
+    payload = {
+        "model": API_MODEL,
+        "messages": [{"role": "user", "content": text}],
+    }
+    body = json.dumps(payload).encode("utf-8")
+    request = urllib.request.Request(
+        API_URL,
+        data=body,
+        method="POST",
+        headers={
+            "content-type": "application/json",
+            "x-api-key": api_key,
+            "anthropic-version": "2023-06-01",
+        },
+    )
+
+    try:
+        with urllib.request.urlopen(request, timeout=8) as response:
+            raw = response.read().decode("utf-8")
+        parsed = json.loads(raw)
+        token_value = parsed.get("input_tokens")
+        if isinstance(token_value, int) and token_value >= 0:
+            return token_value
+    except Exception:
+        return _estimate_tier2(text)
+
+    return _estimate_tier2(text)
+
+
+def auto_select_tier(operation: str, text: str = "") -> int:
+    normalized = (operation or "").strip().lower()
+    if normalized == _FEATURE_UI_DISPLAY:
+        return 1
+    if normalized == _FEATURE_BUDGET_ENFORCEMENT:
+        return 2
+    if normalized == _FEATURE_PREFLIGHT:
+        if len(text) >= 8000:
+            return 3
+        if _estimate_tier2(text) >= 1000:
+            return 3
+        return 2
+    return 1
+
+
+def estimate_tokens(text: str, tier: int = 1) -> int:
+    """Estimate token count with 3 reliability/cost tiers.
+
+    Tier 1: fast heuristic (`len(text)/3.5`).
+    Tier 2: calibrated linear model using bytes, words, lines.
+    Tier 3: Anthropic count_tokens API with graceful fallback to tier 2.
+    """
+    try:
+        if tier == 1:
+            return _estimate_tier1(text)
+        if tier == 2:
+            return _estimate_tier2(text)
+        if tier == 3:
+            return _estimate_tier3(text)
+        return _estimate_tier1(text)
+    except Exception:
+        return _estimate_tier1(text)

package/hooks/branch_manager.py

@@ -0,0 +1,255 @@
+#!/usr/bin/env python3
+"""SessionStart Hook — Smart Branch Manager.
+
+Auto-creates a feature branch when on main/master/develop.
+Extracts task description from OMG state files for branch naming.
+
+Feature-gated: OMG_GIT_WORKFLOW_ENABLED (uses get_feature_flag('GIT_WORKFLOW'))
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+import sys
+from datetime import datetime
+
+HOOKS_DIR = os.path.dirname(__file__)
+if HOOKS_DIR not in sys.path:
+    sys.path.insert(0, HOOKS_DIR)
+
+from _common import setup_crash_handler, json_input, get_feature_flag
+
+setup_crash_handler("branch-manager", fail_closed=False)
+
+# Default branches that trigger feature branch creation
+DEFAULT_BRANCHES = frozenset({"main", "master", "develop"})
+
+# Max length for the descriptive part of branch name
+MAX_BRANCH_NAME_LEN = 50
+
+
+def _get_project_dir() -> str:
+    """Get project directory from env or cwd."""
+    return os.environ.get("CLAUDE_PROJECT_DIR", os.getcwd())
+
+
+def _has_git(project_dir: str) -> bool:
+    """Check if project_dir is inside a git repo."""
+    try:
+        result = subprocess.run(
+            ["git", "-C", project_dir, "rev-parse", "--git-dir"],
+            capture_output=True, text=True, timeout=5,
+        )
+        return result.returncode == 0
+    except Exception:
+        return False
+
+
+def _current_branch(project_dir: str) -> str | None:
+    """Get current branch name. Returns None on failure."""
+    try:
+        result = subprocess.run(
+            ["git", "-C", project_dir, "rev-parse", "--abbrev-ref", "HEAD"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            return result.stdout.strip()
+    except Exception:
+        pass
+    return None
+
+
+def _sanitize_branch_name(description: str) -> str:
+    """Sanitize a description into a valid git branch name segment.
+
+    Rules:
+    - Lowercase
+    - Replace spaces/underscores with hyphens
+    - Strip special characters (keep alphanumeric and hyphens)
+    - Collapse consecutive hyphens
+    - Strip leading/trailing hyphens
+    - Max MAX_BRANCH_NAME_LEN chars
+    """
+    name = description.lower().strip()
+    # Replace spaces and underscores with hyphens
+    name = re.sub(r"[\s_]+", "-", name)
+    # Remove everything except alphanumeric and hyphens
+    name = re.sub(r"[^a-z0-9-]", "", name)
+    # Collapse consecutive hyphens
+    name = re.sub(r"-{2,}", "-", name)
+    # Strip leading/trailing hyphens
+    name = name.strip("-")
+    # Truncate to max length, but don't cut mid-word if possible
+    if len(name) > MAX_BRANCH_NAME_LEN:
+        truncated = name[:MAX_BRANCH_NAME_LEN]
+        # Try to cut at last hyphen to avoid mid-word truncation
+        last_hyphen = truncated.rfind("-")
+        if last_hyphen > 20:
+            truncated = truncated[:last_hyphen]
+        name = truncated.rstrip("-")
+    return name
+
+
+def _extract_task_description(project_dir: str) -> str | None:
+    """Extract task description from OMG state files.
+
+    Priority order:
+    (a) .omg/state/_plan.md title (first # heading)
+    (b) .omg/state/_checklist.md first item
+    (c) .omg/state/working-memory.md last entry
+    (d) fallback: None (caller uses session-{timestamp})
+    """
+    state_dir = os.path.join(project_dir, ".omg", "state")
+
+    # (a) Plan title
+    plan_path = os.path.join(state_dir, "_plan.md")
+    if os.path.isfile(plan_path):
+        try:
+            with open(plan_path, "r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    if line.startswith("# "):
+                        title = line[2:].strip()
+                        if title:
+                            return title
+        except Exception:
+            pass
+
+    # (b) Checklist first item
+    checklist_path = os.path.join(state_dir, "_checklist.md")
+    if os.path.isfile(checklist_path):
+        try:
+            with open(checklist_path, "r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    # Match markdown checkbox items: - [ ] or - [x]
+                    m = re.match(r"^-\s*\[.\]\s*(.+)$", line)
+                    if m:
+                        item = m.group(1).strip()
+                        if item:
+                            return item
+        except Exception:
+            pass
+
+    # (c) Working memory last entry
+    wm_path = os.path.join(state_dir, "working-memory.md")
+    if os.path.isfile(wm_path):
+        try:
+            with open(wm_path, "r", encoding="utf-8", errors="ignore") as f:
+                content = f.read()
+            # Split by ## headings, take last entry's content
+            sections = re.split(r"\n## ", content)
+            if len(sections) > 1:
+                # Last section: first line is heading, rest is content
+                last_lines = sections[-1].split("\n")
+                # Find first non-empty content line after heading
+                for line in last_lines[1:]:
+                    line = line.strip()
+                    if line:
+                        return line
+                # Fallback to heading if no content
+                heading = last_lines[0].strip()
+                if heading:
+                    return heading
+            elif sections:
+                # Single section — try first non-empty non-heading line
+                for line in sections[0].split("\n"):
+                    line = line.strip()
+                    if line and not line.startswith("#"):
+                        return line
+        except Exception:
+            pass
+
+    # (d) No state files found
+    return None
+
+
+def _is_merge_writer_locked(project_dir: str) -> bool:
+    try:
+        _root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        if _root not in sys.path:
+            sys.path.insert(0, _root)
+        from runtime.merge_writer import MergeWriter  # pyright: ignore[reportMissingImports]
+        return MergeWriter(project_dir).is_locked()
+    except Exception:
+        return False
+
+
+def _create_branch(project_dir: str, branch_name: str) -> bool:
+    """Create and checkout a new branch. Returns True on success."""
+    if _is_merge_writer_locked(project_dir):
+        print(
+            f"[OMG branch-manager] Branch creation blocked: merge-writer lock is held",
+            file=sys.stderr,
+        )
+        return False
+    try:
+        result = subprocess.run(
+            ["git", "-C", project_dir, "checkout", "-b", branch_name],
+            capture_output=True, text=True, timeout=10,
+        )
+        return result.returncode == 0
+    except Exception:
+        return False
+
+
+def main() -> None:
+    """Main hook entry point."""
+    data = json_input()
+
+    # Feature gate: exit silently if disabled
+    if not get_feature_flag("GIT_WORKFLOW", default=False):
+        sys.exit(0)
+
+    project_dir = _get_project_dir()
+
+    # No-op if not a git repo
+    if not _has_git(project_dir):
+        sys.exit(0)
+
+    # Get current branch
+    branch = _current_branch(project_dir)
+    if branch is None:
+        sys.exit(0)
+
+    # No-op if already on a non-default branch (feature branch, etc.)
+    if branch not in DEFAULT_BRANCHES:
+        sys.exit(0)
+
+    # Extract task description and build branch name
+    description = _extract_task_description(project_dir)
+    if description:
+        sanitized = _sanitize_branch_name(description)
+    else:
+        sanitized = ""
+
+    if not sanitized:
+        # Fallback: session-{timestamp}
+        sanitized = f"session-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+
+    target_branch = f"feature/{sanitized}"
+
+    # Dry-run mode: output what would happen without executing
+    dry_run = os.environ.get("OMG_GIT_WORKFLOW_DRY_RUN", "").lower() in ("1", "true", "yes")
+    if dry_run:
+        print(
+            f"[OMG branch-manager] DRY-RUN: Would create branch '{target_branch}' from '{branch}'",
+            file=sys.stderr,
+        )
+        sys.exit(0)
+
+    # Create the feature branch
+    success = _create_branch(project_dir, target_branch)
+    if success:
+        print(
+            f"[OMG branch-manager] Created branch '{target_branch}' from '{branch}'",
+            file=sys.stderr,
+        )
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()