@trac3er/oh-my-god 2.0.4 → 2.0.7

package/build/lib/hooks/_cost_ledger.py

@@ -0,0 +1,176 @@
+"""
+Cost Ledger Storage — JSONL persistence for token/cost tracking.
+
+Provides append_cost_entry, read_cost_summary, and rotate_cost_ledger.
+Follows the same fcntl locking + 5MB rotation pattern as tool-ledger.py.
+
+Entry schema:
+    {"ts": ISO8601, "tool": str, "tokens_in": int, "tokens_out": int,
+     "cost_usd": float, "model": str, "session_id": str}
+
+Pure stdlib — no external deps.
+"""
+import fcntl
+import json
+import os
+import shutil
+
+# ── Constants ──
+_LEDGER_SUBDIR = os.path.join(".omg", "state", "ledger")
+_LEDGER_FILENAME = "cost-ledger.jsonl"
+MAX_BYTES = 5 * 1024 * 1024  # 5MB rotation threshold
+
+
+def _ledger_path(project_dir: str) -> str:
+    """Return the absolute path to the cost ledger JSONL file."""
+    return os.path.join(project_dir, _LEDGER_SUBDIR, _LEDGER_FILENAME)
+
+
+def append_cost_entry(project_dir: str, entry: dict) -> None:
+    """Append a cost entry to the cost ledger JSONL file.
+
+    Creates .omg/state/ledger/ if missing. Uses fcntl file locking
+    with fallback to unlocked write (crash isolation invariant).
+
+    Args:
+        project_dir: Project root directory.
+        entry: Dict with keys ts, tool, tokens_in, tokens_out,
+               cost_usd, model, session_id.
+    """
+    ledger_dir = os.path.join(project_dir, _LEDGER_SUBDIR)
+    os.makedirs(ledger_dir, exist_ok=True)
+
+    path = _ledger_path(project_dir)
+    line = json.dumps(entry, separators=(",", ":")) + "\n"
+
+    try:
+        fd = open(path, "a")
+        fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        fd.write(line)
+        fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
+        fd.close()
+    except (ImportError, BlockingIOError):
+        # Fallback: write without locking
+        try:
+            with open(path, "a") as f:
+                f.write(line)
+        except Exception:
+            pass
+    except Exception:
+        pass  # Non-blocking: crash isolation invariant
+
+
+def read_cost_summary(project_dir: str, time_range=None) -> dict:
+    """Read and aggregate cost entries from the ledger.
+
+    Args:
+        project_dir: Project root directory.
+        time_range: Optional (start, end) ISO8601 strings for filtering.
+                    Not yet implemented — reserved for future use.
+
+    Returns:
+        Dict with keys:
+            total_tokens (int): Sum of tokens_in + tokens_out across all entries.
+            total_cost_usd (float): Sum of cost_usd across all entries.
+            by_tool (dict): {tool_name: {"tokens": int, "cost_usd": float, "count": int}}
+            by_session (dict): {session_id: {"tokens": int, "cost_usd": float, "count": int}}
+            entry_count (int): Number of valid entries processed.
+    """
+    empty = {
+        "total_tokens": 0,
+        "total_cost_usd": 0.0,
+        "by_tool": {},
+        "by_session": {},
+        "entry_count": 0,
+    }
+
+    path = _ledger_path(project_dir)
+    if not os.path.exists(path):
+        return empty
+
+    total_tokens = 0
+    total_cost = 0.0
+    by_tool: dict = {}
+    by_session: dict = {}
+    entry_count = 0
+
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    entry = json.loads(line)
+                except (json.JSONDecodeError, ValueError):
+                    # Skip malformed lines gracefully
+                    continue
+
+                tokens_in = entry.get("tokens_in", 0)
+                tokens_out = entry.get("tokens_out", 0)
+                cost_usd = entry.get("cost_usd", 0.0)
+                tool = entry.get("tool", "unknown")
+                session_id = entry.get("session_id", "unknown")
+                line_tokens = tokens_in + tokens_out
+
+                total_tokens += line_tokens
+                total_cost += cost_usd
+                entry_count += 1
+
+                # Aggregate by tool
+                if tool not in by_tool:
+                    by_tool[tool] = {"tokens": 0, "cost_usd": 0.0, "count": 0}
+                by_tool[tool]["tokens"] += line_tokens
+                by_tool[tool]["cost_usd"] += cost_usd
+                by_tool[tool]["count"] += 1
+
+                # Aggregate by session
+                if session_id not in by_session:
+                    by_session[session_id] = {"tokens": 0, "cost_usd": 0.0, "count": 0}
+                by_session[session_id]["tokens"] += line_tokens
+                by_session[session_id]["cost_usd"] += cost_usd
+                by_session[session_id]["count"] += 1
+
+    except Exception:
+        pass  # Crash isolation: return what we have
+
+    return {
+        "total_tokens": total_tokens,
+        "total_cost_usd": total_cost,
+        "by_tool": by_tool,
+        "by_session": by_session,
+        "entry_count": entry_count,
+    }
+
+
+def rotate_cost_ledger(project_dir: str) -> None:
+    """Rotate cost ledger when it exceeds 5MB.
+
+    Follows the same pattern as tool-ledger.py:
+    - Size-only heuristic (avoids O(n) line-count scan)
+    - Keeps only one archive with .1 suffix
+    - Removes old archive before moving current file
+
+    Args:
+        project_dir: Project root directory.
+    """
+    path = _ledger_path(project_dir)
+
+    try:
+        if not os.path.exists(path):
+            return
+
+        size = os.path.getsize(path)
+        if size <= MAX_BYTES:
+            return
+
+        archive = path + ".1"
+        # Keep only one archive
+        if os.path.exists(archive):
+            try:
+                os.remove(archive)
+            except OSError:
+                pass
+        shutil.move(path, archive)
+    except Exception:
+        pass  # Non-blocking: crash isolation invariant

package/build/lib/hooks/_learnings.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""Learnings storage utilities for OMG compound learning."""
+import os
+import glob
+import re
+
+
+def read_file_safe(path, max_bytes=4096):
+    """Safely read a file, returning empty string on error."""
+    try:
+        with open(path, 'r') as f:
+            return f.read(max_bytes)
+    except (OSError, IOError):
+        return ''
+
+
+def aggregate_learnings(project_dir: str, max_patterns: int = 10) -> str:
+    """Read all learning files, aggregate top patterns into summary.
+
+    Returns formatted string with top tool patterns, max 500 chars.
+    """
+    learn_dir = os.path.join(project_dir, '.omg', 'state', 'learnings')
+    if not os.path.isdir(learn_dir):
+        return ''
+
+    all_tools = {}  # tool -> total count across sessions
+    all_files = {}  # file -> total count across sessions
+
+    for fname in os.listdir(learn_dir):
+        if not fname.endswith('.md'):
+            continue
+        content = read_file_safe(os.path.join(learn_dir, fname))
+        in_tools = False
+        in_files = False
+        for line in content.split('\n'):
+            if line.startswith('## Most Used Tools'):
+                in_tools = True
+                in_files = False
+                continue
+            if line.startswith('## Most Modified Files'):
+                in_tools = False
+                in_files = True
+                continue
+            if line.startswith('##'):
+                in_tools = False
+                in_files = False
+                continue
+            # Parse '- toolname: Nx' format
+            match = re.match(r'^-\s+(.+?):\s+(\d+)x\s*$', line.strip())
+            if match:
+                name = match.group(1).strip()
+                count = int(match.group(2))
+                if in_tools:
+                    all_tools[name] = all_tools.get(name, 0) + count
+                elif in_files:
+                    all_files[name] = all_files.get(name, 0) + count
+
+    return format_critical_patterns(all_tools, all_files, max_patterns)
+
+
+def format_critical_patterns(tools: dict, files: dict, max_patterns: int = 10) -> str:
+    """Format tool and file patterns into critical-patterns summary.
+
+    Returns string ≤500 chars.
+    """
+    if not tools and not files:
+        return ''
+
+    lines = ['# Critical Patterns']
+
+    if tools:
+        lines.append('## Top Tools')
+        for tool, count in sorted(tools.items(), key=lambda x: -x[1])[:max_patterns]:
+            lines.append(f'- {tool}: {count}x total')
+
+    if files:
+        lines.append('## Top Files')
+        for fpath, count in sorted(files.items(), key=lambda x: -x[1])[:max_patterns]:
+            basename = os.path.basename(fpath)
+            lines.append(f'- {basename}: {count}x total')
+
+    result = '\n'.join(lines)
+    return result[:500]  # Cap at 500 chars
+
+
+def rotate_learnings(project_dir: str, max_files: int = 30) -> int:
+    """Delete oldest learning files if count exceeds max_files.
+
+    Returns number of files deleted.
+    """
+    learn_dir = os.path.join(project_dir, '.omg', 'state', 'learnings')
+    if not os.path.isdir(learn_dir):
+        return 0
+
+    files = sorted(glob.glob(os.path.join(learn_dir, '*.md')))
+    excess = len(files) - max_files
+    if excess <= 0:
+        return 0
+
+    for f in files[:excess]:
+        try:
+            os.remove(f)
+        except OSError:
+            pass
+    return excess
+
+
+def save_critical_patterns(project_dir: str) -> str:
+    """Generate and save critical-patterns.md to .omg/knowledge/.
+
+    Returns the path of the written file, or empty string on failure.
+    """
+    content = aggregate_learnings(project_dir)
+    if not content:
+        return ''
+
+    knowledge_dir = os.path.join(project_dir, '.omg', 'knowledge')
+    os.makedirs(knowledge_dir, exist_ok=True)
+    path = os.path.join(knowledge_dir, 'critical-patterns.md')
+
+    try:
+        with open(path, 'w') as f:
+            f.write(content)
+        return path
+    except OSError:
+        return ''

package/build/lib/hooks/_memory.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+import glob
+import os
+from datetime import datetime
+
+
+def save_memory(project_dir: str, session_id: str, content: str) -> str:
+    memory_dir = os.path.join(project_dir, ".omg", "state", "memory")
+    os.makedirs(memory_dir, exist_ok=True)
+    date_str = datetime.now().strftime("%Y-%m-%d")
+    session_short = session_id[:8] if len(session_id) > 8 else session_id
+    filename = f"{date_str}-{session_short}.md"
+    filepath = os.path.join(memory_dir, filename)
+    content = content[:500]
+    if os.path.exists(filepath):
+        with open(filepath, "a") as file_obj:
+            _ = file_obj.write("\n" + content)
+    else:
+        with open(filepath, "w") as file_obj:
+            _ = file_obj.write(content)
+    return filepath
+
+
+def get_recent_memories(
+    project_dir: str, max_files: int = 5, max_chars_total: int = 300
+) -> str:
+    memory_dir = os.path.join(project_dir, ".omg", "state", "memory")
+    if not os.path.exists(memory_dir):
+        return ""
+    files = sorted(glob.glob(os.path.join(memory_dir, "*.md")), reverse=True)
+    files = files[:max_files]
+    result: list[str] = []
+    total = 0
+    separator = "\n---\n"
+    for file_path in files:
+        try:
+            with open(file_path) as file_obj:
+                content = file_obj.read()
+            separator_len = len(separator) if result else 0
+            remaining = max_chars_total - total - separator_len
+            if remaining <= 0:
+                break
+            if len(content) > remaining:
+                content = content[:remaining]
+            if not content:
+                break
+            if result:
+                total += separator_len
+            result.append(content)
+            total += len(content)
+            if total >= max_chars_total:
+                break
+        except OSError:
+            continue
+    return separator.join(result)
+
+
+def rotate_memories(project_dir: str, max_files: int = 50) -> int:
+    memory_dir = os.path.join(project_dir, ".omg", "state", "memory")
+    if not os.path.exists(memory_dir):
+        return 0
+    files = sorted(glob.glob(os.path.join(memory_dir, "*.md")))
+    excess = len(files) - max_files
+    if excess <= 0:
+        return 0
+    for file_path in files[:excess]:
+        try:
+            os.remove(file_path)
+        except OSError:
+            pass
+    return excess
+
+
+
+def search_memories(project_dir: str, query_keywords: list, max_results: int = 3, max_chars: int = 200) -> str:
+    """Search memory files by keyword relevance. Returns formatted excerpt string."""
+    memory_dir = os.path.join(project_dir, '.omg', 'state', 'memory')
+    if not os.path.isdir(memory_dir):
+        return ''
+    results = []
+    for fname in sorted(os.listdir(memory_dir), reverse=True):
+        if not fname.endswith('.md'):
+            continue
+        fpath = os.path.join(memory_dir, fname)
+        try:
+            with open(fpath, 'r', encoding='utf-8', errors='ignore') as f:
+                content = f.read(2048)
+        except OSError:
+            continue
+        score = sum(1 for kw in query_keywords if kw.lower() in content.lower())
+        if score > 0:
+            results.append((score, fname, content))
+    results.sort(key=lambda x: -x[0])
+    summary_parts = []
+    chars_used = 0
+    for score, fname, content in results[:max_results]:
+        lines = [l.strip() for l in content.split('\n') if l.strip() and not l.startswith('#')]
+        excerpt = ' '.join(lines[:3])[:100]
+        if chars_used + len(excerpt) > max_chars:
+            break
+        summary_parts.append(f'[{fname}] {excerpt}')
+        chars_used += len(excerpt)
+    return '\n'.join(summary_parts)

package/build/lib/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 ""