@trac3er/oh-my-god 2.0.4 → 2.0.7

package/build/lib/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 Exception:
+        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/build/lib/hooks/branch_manager.py

@@ -0,0 +1,236 @@
+#!/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'))
+"""
+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 _create_branch(project_dir: str, branch_name: str) -> bool:
+    """Create and checkout a new branch. Returns True on success."""
+    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()

package/build/lib/hooks/budget_governor.py

@@ -0,0 +1,232 @@
+#!/usr/bin/env python3
+"""PostToolUse budget governor (BATS-style additionalContext injection)."""
+from __future__ import annotations
+
+import importlib.util
+import json
+import os
+import sys
+from typing import Any
+
+HOOKS_DIR = os.path.dirname(__file__)
+
+
+def _load_module(module_name: str, filename: str):
+    path = os.path.join(HOOKS_DIR, filename)
+    spec = importlib.util.spec_from_file_location(module_name, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Unable to load module: {filename}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+_common = _load_module("_common", "_common.py")
+_cost_ledger = _load_module("_cost_ledger", "_cost_ledger.py")
+_token_counter = _load_module("_token_counter", "_token_counter.py")
+
+get_feature_flag = _common.get_feature_flag
+get_project_dir = _common.get_project_dir
+json_input = _common.json_input
+setup_crash_handler = _common.setup_crash_handler
+read_cost_summary = _cost_ledger.read_cost_summary
+estimate_tokens = _token_counter.estimate_tokens
+
+DEFAULT_SESSION_LIMIT_USD = 5.0
+DEFAULT_INPUT_PER_MTOK = 3.0
+DEFAULT_OUTPUT_PER_MTOK = 15.0
+DEFAULT_PROJECTED_TOOL_CALLS = 50
+DEFAULT_THRESHOLDS = [50, 80, 95]
+THRESHOLD_STATE_FILE = ".omg/state/.cost-threshold-state.json"
+
+
+def _safe_float(value, default: float) -> float:
+    try:
+        parsed = float(value)
+        if parsed <= 0:
+            return default
+        return parsed
+    except (TypeError, ValueError):
+        return default
+
+
+def _read_budget_config(project_dir: str) -> tuple[float, float, float]:
+    session_limit = DEFAULT_SESSION_LIMIT_USD
+    input_per_mtok = DEFAULT_INPUT_PER_MTOK
+    output_per_mtok = DEFAULT_OUTPUT_PER_MTOK
+
+    settings_path = os.path.join(project_dir, "settings.json")
+    try:
+        with open(settings_path, "r", encoding="utf-8") as f:
+            settings = json.load(f)
+        budget_cfg = settings.get("_omg", {}).get("cost_budget", {})
+        pricing = budget_cfg.get("pricing", {})
+        session_limit = _safe_float(budget_cfg.get("session_limit_usd"), DEFAULT_SESSION_LIMIT_USD)
+        input_per_mtok = _safe_float(pricing.get("input_per_mtok"), DEFAULT_INPUT_PER_MTOK)
+        output_per_mtok = _safe_float(pricing.get("output_per_mtok"), DEFAULT_OUTPUT_PER_MTOK)
+    except Exception:
+        pass
+
+    return session_limit, input_per_mtok, output_per_mtok
+
+
+def _to_text(value) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    try:
+        return json.dumps(value, ensure_ascii=True, sort_keys=True)
+    except Exception:
+        return str(value)
+
+
+def _estimate_call_cost(tool_input, tool_response, input_per_mtok: float, output_per_mtok: float) -> float:
+    input_text = _to_text(tool_input)
+    output_text = _to_text(tool_response)
+
+    tokens_in = estimate_tokens(input_text, tier=2)
+    tokens_out = estimate_tokens(output_text, tier=2)
+
+    cost_in = (tokens_in / 1_000_000.0) * input_per_mtok
+    cost_out = (tokens_out / 1_000_000.0) * output_per_mtok
+    return max(0.0, cost_in + cost_out)
+
+
+def _project_total_calls(used_cost_usd: float, used_calls: int, session_limit_usd: float) -> int:
+    if used_calls <= 0 or used_cost_usd <= 0:
+        return DEFAULT_PROJECTED_TOOL_CALLS
+    avg_cost = used_cost_usd / float(used_calls)
+    if avg_cost <= 0:
+        return DEFAULT_PROJECTED_TOOL_CALLS
+    projected = max(used_calls, int(round(session_limit_usd / avg_cost)))
+    if projected > (DEFAULT_PROJECTED_TOOL_CALLS * 10):
+        return DEFAULT_PROJECTED_TOOL_CALLS
+    rounded = int(round(projected / 10.0) * 10)
+    return max(10, rounded)
+
+
+def _build_context(used_cost_usd: float, session_limit_usd: float, used_calls: int, projected_calls: int) -> str:
+    remaining_ratio = 1.0 - (used_cost_usd / session_limit_usd)
+    remaining_pct = int(round(max(0.0, min(1.0, remaining_ratio)) * 100))
+    return (
+        f"Budget: {remaining_pct}% remaining | "
+        f"${used_cost_usd:.2f} of ${session_limit_usd:.2f} used | "
+        f"{used_calls} tool calls of ~{projected_calls}"
+    )
+
+
+def _get_threshold_message(pct: int) -> str:
+    if pct >= 95:
+        return (
+            f"@cost-limit: {pct}% budget used. "
+            "Complete current task and stop. Do NOT start new tasks."
+        )
+    if pct >= 80:
+        return (
+            f"@cost-critical: {pct}% budget used. "
+            "Be efficient \u2014 minimize unnecessary tool calls, "
+            "batch operations where possible."
+        )
+    return f"@cost-warning: {pct}% budget used"
+
+
+def _read_thresholds_config(project_dir: str) -> list[int]:
+    try:
+        settings_path = os.path.join(project_dir, "settings.json")
+        with open(settings_path, "r", encoding="utf-8") as f:
+            settings = json.load(f)
+        raw = settings.get("_omg", {}).get("cost_budget", {}).get("thresholds")
+        if isinstance(raw, list) and all(isinstance(t, (int, float)) for t in raw):
+            return sorted(int(t) for t in raw)
+    except Exception:
+        pass
+    return list(DEFAULT_THRESHOLDS)
+
+
+def _read_threshold_state(project_dir: str) -> dict[str, Any]:
+    path = os.path.join(project_dir, THRESHOLD_STATE_FILE)
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {"session_id": "", "fired": []}
+
+
+def _write_threshold_state(project_dir: str, state: dict[str, Any]) -> None:
+    path = os.path.join(project_dir, THRESHOLD_STATE_FILE)
+    try:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(state, f, separators=(",", ":"))
+    except Exception:
+        pass
+
+
+def _check_thresholds(
+    used_pct: float, project_dir: str, session_id: str
+) -> list[str]:
+    thresholds = _read_thresholds_config(project_dir)
+    state = _read_threshold_state(project_dir)
+
+    if state.get("session_id", "") != session_id:
+        state = {"session_id": session_id, "fired": []}
+
+    already_fired = set(state.get("fired", []))
+    new_messages: list[str] = []
+    new_fired: list[int] = []
+
+    for threshold in thresholds:
+        if used_pct >= threshold and threshold not in already_fired:
+            new_messages.append(_get_threshold_message(threshold))
+            new_fired.append(threshold)
+
+    if new_fired:
+        state["fired"] = sorted(list(already_fired | set(new_fired)))
+        state["session_id"] = session_id
+        _write_threshold_state(project_dir, state)
+
+    return new_messages
+
+
+def main() -> None:
+    setup_crash_handler("budget-governor", fail_closed=False)
+
+    payload = json_input()
+    if not get_feature_flag("COST_TRACKING", default=False):
+        sys.exit(0)
+
+    project_dir = get_project_dir()
+    session_limit_usd, input_per_mtok, output_per_mtok = _read_budget_config(project_dir)
+    summary = read_cost_summary(project_dir)
+
+    estimated_current_cost = _estimate_call_cost(
+        payload.get("tool_input", {}),
+        payload.get("tool_response", {}),
+        input_per_mtok,
+        output_per_mtok,
+    )
+
+    used_cost_usd = float(summary.get("total_cost_usd", 0.0)) + estimated_current_cost
+    used_calls = int(summary.get("entry_count", 0)) + 1
+    projected_calls = _project_total_calls(used_cost_usd, used_calls, session_limit_usd)
+
+    context = _build_context(
+        used_cost_usd=used_cost_usd,
+        session_limit_usd=session_limit_usd,
+        used_calls=used_calls,
+        projected_calls=projected_calls,
+    )
+
+    used_pct = (used_cost_usd / session_limit_usd * 100) if session_limit_usd > 0 else 0.0
+    session_id = os.environ.get("CLAUDE_SESSION_ID", "")
+    threshold_alerts = _check_thresholds(used_pct, project_dir, session_id)
+    if threshold_alerts:
+        context += "\n" + "\n".join(threshold_alerts)
+
+    json.dump({"additionalContext": context}, sys.stdout)
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()