@trac3r/oh-my-god 2.2.11

package/hooks/budget_governor.py

@@ -0,0 +1,326 @@
+#!/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, dict[str, Any]]:
+    session_limit = DEFAULT_SESSION_LIMIT_USD
+    input_per_mtok = DEFAULT_INPUT_PER_MTOK
+    output_per_mtok = DEFAULT_OUTPUT_PER_MTOK
+    tier_info: dict[str, Any] = {
+        "tier": "free",
+        "tier_limit_usd": DEFAULT_SESSION_LIMIT_USD,
+        "tier_provenance": "default",
+    }
+
+    try:
+        parent = os.path.dirname(HOOKS_DIR)
+        if parent not in sys.path:
+            sys.path.insert(0, parent)
+        from runtime.subscription_tiers import detect_tier
+
+        provider = os.environ.get("OMG_PROVIDER", "claude")
+        tier_result = detect_tier(provider, project_dir=project_dir)
+        session_limit = _safe_float(
+            tier_result.get("budget_usd_per_session"),
+            DEFAULT_SESSION_LIMIT_USD,
+        )
+        tier_info = {
+            "tier": tier_result.get("tier", "free"),
+            "tier_limit_usd": float(
+                tier_result.get("budget_usd_per_session", DEFAULT_SESSION_LIMIT_USD)
+            ),
+            "tier_provenance": tier_result.get("provenance", "default"),
+        }
+    except Exception:
+        pass
+
+    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"), session_limit)
+        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, tier_info
+
+
+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 _check_budget_envelope(project_dir: str) -> str:
+    run_id = os.environ.get("OMG_RUN_ID", "").strip()
+    if not run_id:
+        return ""
+    try:
+        parent = os.path.dirname(HOOKS_DIR)
+        if parent not in sys.path:
+            sys.path.insert(0, parent)
+        from runtime.budget_envelopes import get_budget_envelope_manager
+
+        mgr = get_budget_envelope_manager(project_dir)
+        result = mgr.check_envelope(run_id)
+        if result.status == "ok":
+            return ""
+        action_tag = {"warn": "@envelope-warning", "reflect": "@envelope-critical", "block": "@envelope-limit"}
+        tag = action_tag.get(result.governance_action, "@envelope-warning")
+        return f"{tag}: {result.reason} [action={result.governance_action}]"
+    except Exception:
+        return ""
+
+
+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, tier_info = _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,
+    )
+
+    session_id = os.environ.get("CLAUDE_SESSION_ID", "")
+    by_session = summary.get("by_session", {})
+
+    if session_id and session_id in by_session:
+        session_data = by_session[session_id]
+        used_cost_usd = float(session_data.get("cost_usd", 0.0)) + estimated_current_cost
+        used_calls = int(session_data.get("count", 0)) + 1
+        provenance = "session"
+    else:
+        used_cost_usd = float(summary.get("total_cost_usd", 0.0)) + estimated_current_cost
+        used_calls = int(summary.get("entry_count", 0)) + 1
+        provenance = "default"
+
+    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
+    threshold_alerts = _check_thresholds(used_pct, project_dir, session_id)
+    if threshold_alerts:
+        context += "\n" + "\n".join(threshold_alerts)
+
+    envelope_context = _check_budget_envelope(project_dir)
+    if envelope_context:
+        context += "\n" + envelope_context
+
+    tier_limit_usd = tier_info["tier_limit_usd"]
+    near_limit = used_cost_usd > (0.80 * tier_limit_usd)
+    limit_reached = used_cost_usd >= tier_limit_usd
+
+    throttle_reason: str | None = None
+    if near_limit or limit_reached:
+        throttle_reason = (
+            f"tier={tier_info['tier']} limit={tier_limit_usd:.2f} "
+            f"used={used_cost_usd:.2f} provenance={tier_info['tier_provenance']}"
+        )
+
+    if limit_reached:
+        context += f"\n@tier-limit-reached: {throttle_reason}"
+    elif near_limit:
+        context += f"\n@tier-near-limit: {throttle_reason}"
+
+    output: dict[str, Any] = {
+        "additionalContext": context,
+        "provenance": provenance,
+        "tier": tier_info["tier"],
+        "tier_limit_usd": tier_limit_usd,
+        "near_limit": near_limit,
+        "limit_reached": limit_reached,
+        "throttle_reason": throttle_reason,
+    }
+
+    if tier_info["tier_provenance"] in ("cache", "default"):
+        output["tier_provenance_fallback"] = (
+            f"tier data from {tier_info['tier_provenance']} source "
+            "— may not reflect current subscription"
+        )
+
+    json.dump(output, sys.stdout)
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

package/hooks/circuit-breaker.py

@@ -0,0 +1,270 @@
+#!/usr/bin/env python3
+"""
+PostToolUse Hook: Circuit Breaker + Auto-Escalation (v4)
+Key v4 change: After 3 failures, automatically SUGGESTS Codex/Gemini escalation
+instead of just saying "stop". Actionable, not just blocking.
+"""
+import json, sys, os
+from datetime import datetime, timezone, timedelta
+
+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, _resolve_project_dir
+from state_migration import resolve_state_dir
+
+setup_crash_handler("circuit-breaker", fail_closed=False)
+
+# Domain-aware routing hints: pattern prefix → suggested model
+DOMAIN_MODEL_HINTS = {
+    'Bash:pytest': 'codex',
+    'Bash:npm': 'codex',
+    'Bash:python': 'codex',
+    'Write:': 'codex',
+    'Edit:': 'codex',
+}
+
+
+def _get_domain_hint(pk: str) -> str:
+    """Return model hint for a pattern key, or empty string."""
+    for prefix, model in DOMAIN_MODEL_HINTS.items():
+        if pk.startswith(prefix):
+            return model
+    return ''
+
+data = json_input()
+
+tool = data.get("tool_name", "")
+tool_input = data.get("tool_input", {})
+tool_response = data.get("tool_response", {})
+project_dir = _resolve_project_dir()
+
+ledger_dir = resolve_state_dir(project_dir, "state/ledger", "ledger")
+tracker_path = os.path.join(ledger_dir, "failure-tracker.json")
+os.makedirs(os.path.dirname(tracker_path), exist_ok=True)
+
+# Determine failure
+is_failure = False
+if tool == "Bash":
+    ec = None
+    if isinstance(tool_response, dict):
+        ec = tool_response.get("exitCode", tool_response.get("exit_code"))
+    if ec is not None and ec != 0:
+        is_failure = True
+elif tool in ("Write", "Edit", "MultiEdit"):
+    if isinstance(tool_response, dict) and not tool_response.get("success", True):
+        is_failure = True
+
+# Pattern key — normalized to prevent duplicates
+# "npm test" and "npm run test" should be the same failure pattern
+pattern_key = tool
+if tool == "Bash":
+    cmd = tool_input.get("command", "").strip()
+    # Normalize: strip common prefixes, reduce to base command
+    cmd_clean = cmd
+    # Strip package manager prefixes: npx, pnpm, yarn, bunx
+    cmd_clean = cmd_clean.replace("npx ", "").replace("pnpm ", "npm ").replace("yarn ", "npm ").replace("bunx ", "")
+    # Strip python -m X → X (e.g., "python3 -m pytest" → "pytest")
+    if cmd_clean.startswith("python3 -m "):
+        cmd_clean = cmd_clean.replace("python3 -m ", "", 1)
+    elif cmd_clean.startswith("python -m "):
+        cmd_clean = cmd_clean.replace("python -m ", "", 1)
+    words = cmd_clean.split()[:3]  # first 3 words for more specificity
+    # Remove common noise: run, exec, --
+    words = [w for w in words if not w.startswith("-") and w not in ("run", "exec")][:2]
+    pattern_key = f"Bash:{' '.join(words)}" if words else f"Bash:{cmd[:30]}"
+elif tool in ("Write", "Edit", "MultiEdit"):
+    fp = tool_input.get("file_path", "")
+    # Normalize: use basename to avoid path-length variants
+    pattern_key = f"{tool}:{os.path.basename(fp)}" if fp else tool
+pattern_key = pattern_key[:120].replace("\n", " ")
+
+# Load tracker
+tracker = {}
+if os.path.exists(tracker_path):
+    try:
+        with open(tracker_path, "r") as f:
+            tracker = json.load(f)
+        if not isinstance(tracker, dict):
+            tracker = {}
+    except Exception:
+        tracker = {}
+
+# Evict stale (24h) + cap 100
+now = datetime.now(timezone.utc)
+cutoff = now - timedelta(hours=24)
+
+
+def _parse_ts(ts_str):
+    """Parse ISO timestamp string to datetime, returning None on failure."""
+    try:
+        # Handle both Z-suffix and +00:00 formats
+        ts = ts_str.replace("Z", "+00:00") if ts_str.endswith("Z") else ts_str
+        return datetime.fromisoformat(ts)
+    except (ValueError, TypeError, AttributeError):
+        return None
+
+
+def _effective_count(entry: dict[str, object], now: datetime) -> float:
+    """Apply time-decay: failures >30 min old count as 0.5x."""
+    last_failure = entry.get('last_failure', '')
+    last_ts = _parse_ts(last_failure if isinstance(last_failure, str) else '')
+    raw_count = entry.get('count', 0)
+    count_value = float(raw_count) if isinstance(raw_count, (int, float)) else 0.0
+    if last_ts is None:
+        return count_value
+    age_minutes = (now - last_ts).total_seconds() / 60
+    if age_minutes > 30:
+        return count_value * 0.5
+    return count_value
+
+
+tracker = {k: v for k, v in tracker.items()
+           if isinstance(v, dict) and (_parse_ts(v.get("last_failure", "")) or cutoff) >= cutoff}
+if len(tracker) > 100:
+    for k in sorted(tracker, key=lambda x: tracker[x].get("last_failure", ""))[:-100]:
+        del tracker[k]
+
+if is_failure:
+    entry_raw = tracker.get(pattern_key, {"count": 0, "errors": []})
+    if not isinstance(entry_raw, dict):
+        entry_raw = {"count": 0, "errors": []}
+    entry: dict[str, object] = dict(entry_raw)
+    entry_count = entry.get("count", 0)
+    count_value = entry_count if isinstance(entry_count, (int, float)) else 0
+    entry["count"] = count_value + 1
+    entry["last_failure"] = now.isoformat()
+
+    err = ""
+    if isinstance(tool_response, dict):
+        err = str(tool_response.get("stderr", tool_response.get("stdout", "")))[:200].strip()
+    errors = entry.get("errors", [])
+    if not isinstance(errors, list):
+        errors = []
+    # Deduplicate: don't store the same error message twice in a row
+    if err and (not errors or errors[-1] != err):
+        errors.append(err)
+    entry["errors"] = errors[-3:]  # keep last 3 unique errors
+    tracker[pattern_key] = entry
+
+    try:
+        import fcntl
+        # Open read+write without truncating, acquire lock, THEN truncate and write.
+        # This prevents data loss if lock acquisition fails after truncation.
+        fd = open(tracker_path, "a+")
+        fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        fd.seek(0)
+        fd.truncate()
+        json.dump(tracker, fd, indent=2)
+        fd.flush()
+        fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
+        fd.close()
+    except (ImportError, BlockingIOError):
+        # Fallback: write without lock (better than losing data)
+        try:
+            with open(tracker_path, "w") as f:
+                json.dump(tracker, f, indent=2)
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+    count = entry["count"]
+    effective_count = _effective_count(entry, now)
+    domain_hint = _get_domain_hint(pattern_key)
+    recent_errs = "\n".join(f"  - {e}" for e in entry["errors"] if e)
+
+    if effective_count >= 5:
+        print(
+            f"CIRCUIT BREAKER: '{pattern_key}' failed {count}x (effective {effective_count:.1f}x).\n"
+            f"STOP. This approach is broken.\n"
+            f"{recent_errs}\n\n"
+            f"Domain hint: {domain_hint or 'none'}\n"
+            f"ESCALATE NOW — pick one:\n"
+            f"  /OMG:escalate codex \"Debug: {pattern_key} fails with {entry['errors'][-1][:80]}\"\n"
+            f"  /OMG:escalate gemini \"Review: approach for {pattern_key}\"\n"
+            f"  Ask the user for a completely different approach\n"
+            f"  Skip this step: mark [!] in checklist and move on",
+            file=sys.stderr
+        )
+        # NOTE: exit(0), not exit(2). Non-zero exits crash sibling hooks
+        # ("Sibling tool call errored"). The warning is in stderr.
+        sys.exit(0)
+
+    elif effective_count >= 3:
+        print(
+            f"CIRCUIT BREAKER WARNING: '{pattern_key}' failed {count}x (effective {effective_count:.1f}x).\n"
+            f"Last error: {entry['errors'][-1][:150] if entry['errors'] else 'unknown'}\n\n"
+            f"Domain hint: {domain_hint or 'none'}\n"
+            f"STOP auto-retrying. Try:\n"
+            f"  1. Fundamentally different approach\n"
+            f"  2. /OMG:escalate codex \"Why does {pattern_key} keep failing?\"\n"
+            f"  3. Ask the user",
+            file=sys.stderr
+        )
+        sys.exit(0)
+
+else:
+    # On success, clear this pattern AND similar variants
+    # Helper to normalize a tracker key by re-normalizing the command part
+    def _normalize_tracker_key(pk):
+        """Normalize a tracker key by applying the same rules as pattern_key generation."""
+        if not pk.startswith("Bash:"):
+            return pk
+        
+        cmd = pk[5:]  # Remove "Bash:" prefix
+        # Apply the same normalization as in pattern_key generation
+        cmd_clean = cmd
+        cmd_clean = cmd_clean.replace("npx ", "").replace("pnpm ", "npm ").replace("yarn ", "npm ").replace("bunx ", "")
+        if cmd_clean.startswith("python3 -m "):
+            cmd_clean = cmd_clean.replace("python3 -m ", "", 1)
+        elif cmd_clean.startswith("python -m "):
+            cmd_clean = cmd_clean.replace("python -m ", "", 1)
+        words = cmd_clean.split()[:3]
+        words = [w for w in words if not w.startswith("-") and w not in ("run", "exec")][:2]
+        normalized_cmd = ' '.join(words) if words else cmd[:30]
+        return f"Bash:{normalized_cmd}"
+    
+    normalized_pattern_key = _normalize_tracker_key(pattern_key)
+    changed = False
+    keys_to_remove = []
+    
+    for k in tracker:
+        # Normalize both keys for comparison
+        normalized_k = _normalize_tracker_key(k)
+        if normalized_k == normalized_pattern_key:
+            keys_to_remove.append(k)
+    
+    for k in keys_to_remove:
+        del tracker[k]
+        changed = True
+    
+    if changed:
+        try:
+            with open(tracker_path, "w") as f:
+                json.dump(tracker, f, indent=2)
+        except Exception:
+            pass
+        _recovery_path = os.path.join(ledger_dir, 'recovery.jsonl')
+        try:
+            import json as _json
+            _rec = _json.dumps({
+                'pattern': pattern_key,
+                'recovered_at': now.isoformat(),
+                'cleared_count': len(keys_to_remove),
+            })
+            with open(_recovery_path, 'a') as _rf:
+                _rf.write(_rec + '\n')
+            try:
+                with open(_recovery_path, 'r') as _rf:
+                    _lines = _rf.readlines()
+                if len(_lines) > 200:
+                    with open(_recovery_path, 'w') as _rf:
+                        _rf.writelines(_lines[-200:])
+            except OSError:
+                pass
+        except Exception:
+            pass
+
+sys.exit(0)