claudetools/.claude/scripts/recover_grok_session.py

#!/usr/bin/env python3
"""
recover_grok_session.py -- reconstruct a ClaudeTools session log from a Grok session.

Grok stores per-workspace sessions under:
  ~/.grok/sessions/<slug>/                 (slug = URL-encoded path, e.g. D%3A%5Cclaudetools)
  ~/.grok/sessions/<slug>/<uuid>/          (one dir per actual session)
    - chat_history.jsonl, events.jsonl, updates.jsonl
    - terminal/ (call-*.log files with exact command + output for run_terminal_cmd etc.)
    - summary.json, system_prompt.txt, prompt_context.json, etc.

This is the Grok counterpart to recover_session.py (which targets Claude Code's
~/.claude/projects/<slug>/<uuid>.jsonl transcripts).

The goal is the same: produce a markdown session log in the format expected by
.claude/commands/save.md (## User, ## Summary, Key Decisions, Commands & Outputs,
File Changes, etc.) so that /save, wiki-compile, memory, etc. work uniformly
regardless of which driver (Claude Code or Grok) was used for the session.

Status: skeleton / MVP. Extracts terminal commands + outputs (the most valuable
verbatim evidence) and basic session metadata. Full chat reconstruction + Ollama
prose drafting (like the Claude recover) can be added by sharing code with
recover_session.py.

Usage (from repo root):
  python .claude/scripts/recover_grok_session.py --list
  python .claude/scripts/recover_grok_session.py --latest --print
  python .claude/scripts/recover_grok_session.py --uuid 019e8b67-f97e-7b33-9c45-ec34b342d3eb --auto

The /recover command will be extended to dispatch to this for Grok sessions.
"""

from __future__ import annotations
import argparse
import json
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

# --- Resolution of Grok home and current workspace slug (mirrors how Grok does it) ---

def get_grok_home() -> Path:
    if os.name == "nt":
        return Path(os.environ.get("USERPROFILE", "~")).expanduser() / ".grok"
    return Path(os.environ.get("HOME", "~")).expanduser() / ".grok"

def get_workspace_slug(cwd: Optional[Path] = None) -> str:
    """Mimic Grok's slug for the sessions/<slug> directory.
    Grok uses URL-style encoding of the absolute path ( : -> %3A, \\ -> %5C, / -> %2F ).
    """
    root = (cwd or Path.cwd()).resolve()
    # On Windows the path has drive letter + backslashes.
    s = str(root)
    # Percent-encode unsafe for dir name (Grok appears to use % encoding).
    # Simple approach that matched observed dirs: replace special with %HH
    encoded = re.sub(r'([^\w\-./])', lambda m: f"%{ord(m.group(1)):02X}", s)
    # Grok observed: D%3A%5Cclaudetools (backslashes and colon encoded, no extra for / sometimes)
    # But on the FS it was D%3A%5Cclaudetools
    return encoded.replace("/", "%2F")

def find_session_dirs(workspace_root: Optional[Path] = None) -> List[Tuple[str, Path]]:
    """Return list of (uuid, dir_path) for sessions under the current workspace slug."""
    grok_home = get_grok_home()
    slug = get_workspace_slug(workspace_root)
    base = grok_home / "sessions" / slug
    if not base.exists():
        return []
    results = []
    for p in base.iterdir():
        if p.is_dir() and re.match(r'^[0-9a-f]{8}-', p.name):  # rough UUID check
            results.append((p.name, p))
    results.sort(key=lambda t: t[1].stat().st_mtime, reverse=True)
    return results

def get_latest_session(workspace_root: Optional[Path] = None) -> Optional[Tuple[str, Path]]:
    sessions = find_session_dirs(workspace_root)
    return sessions[0] if sessions else None

# --- Extraction ---

def read_jsonl(path: Path) -> List[Dict[str, Any]]:
    if not path.exists():
        return []
    out = []
    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if line:
            try:
                out.append(json.loads(line))
            except Exception:
                pass
    return out

def extract_terminal_commands(session_dir: Path) -> List[Dict[str, str]]:
    """Pull exact commands + outputs from the terminal/ call logs.
    These are the highest-fidelity evidence for the session log "Commands & Outputs" section.
    """
    term_dir = session_dir / "terminal"
    if not term_dir.exists():
        return []
    results = []
    for f in sorted(term_dir.glob("call-*.log")):
        try:
            content = f.read_text(encoding="utf-8", errors="replace")
            # The tool already logs the command in the filename sometimes, or inside.
            # From observed usage the file contains the full command + stdout/stderr.
            results.append({
                "file": f.name,
                "content": content[:20000]  # cap per call for practicality
            })
        except Exception:
            pass
    return results

def extract_basic_metadata(session_dir: Path) -> Dict[str, Any]:
    meta: Dict[str, Any] = {"uuid": session_dir.name}
    # Try summary.json
    sj = session_dir / "summary.json"
    if sj.exists():
        try:
            meta["summary"] = json.loads(sj.read_text(encoding="utf-8", errors="replace"))
        except Exception:
            pass
    # opened time from dir mtime or active_sessions if present
    meta["dir_mtime"] = datetime.fromtimestamp(session_dir.stat().st_mtime, tz=timezone.utc).isoformat()
    return meta

def build_session_log(uuid: str, session_dir: Path, workspace_root: Path) -> str:
    """Produce markdown in the approximate shape expected by /save and wiki tools."""
    lines: List[str] = []
    lines.append(f"# Grok Session Log — {uuid}")
    lines.append("")
    lines.append(f"**Workspace:** {workspace_root}")
    lines.append(f"**Grok session dir:** {session_dir}")
    lines.append(f"**Recovered at:** {datetime.now(timezone.utc).isoformat()}")
    lines.append("**Driver:** Grok")
    lines.append("")

    meta = extract_basic_metadata(session_dir)
    lines.append("## Metadata")
    lines.append(json.dumps(meta, indent=2, default=str))
    lines.append("")

    # Terminal calls (commands + output)
    calls = extract_terminal_commands(session_dir)
    lines.append("## Commands & Outputs (from terminal call logs)")
    if calls:
        for c in calls:
            lines.append(f"### {c['file']}")
            lines.append("```")
            lines.append(c["content"].rstrip())
            lines.append("```")
            lines.append("")
    else:
        lines.append("(no terminal call logs found)")
        lines.append("")

    # TODO (future): parse chat_history / events / updates for user/assistant turns,
    # file edits (search_replace results), read_file contents, decisions, etc.
    # Then feed prose sections (Summary, Key Decisions, Pending Tasks) to Ollama
    # exactly like recover_session.py does.

    lines.append("## Notes")
    lines.append("This is an MVP recovery from Grok native session artifacts.")
    lines.append("Full fidelity (conversation turns, structured tool calls, file changes) ")
    lines.append("and Ollama-assisted summarization can be added by extending this script ")
    lines.append("or sharing logic with recover_session.py.")
    lines.append("Run with --auto to write a session-log/ file in the conventional location.")
    lines.append("")

    return "\n".join(lines)

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--uuid", help="Specific Grok session UUID to recover")
    ap.add_argument("--latest", action="store_true", help="Recover the most recent session for this workspace")
    ap.add_argument("--list", action="store_true", help="List available Grok sessions for this workspace slug")
    ap.add_argument("--print", action="store_true", help="Print the reconstructed log to stdout")
    ap.add_argument("--auto", action="store_true", help="Write a session log under session-logs/ (or projects/... ) following ClaudeTools conventions")
    ap.add_argument("--workspace", type=Path, default=None, help="Override workspace root for slug calculation")
    args = ap.parse_args()

    ws_root = (args.workspace or Path.cwd()).resolve()

    if args.list:
        for uuid, d in find_session_dirs(ws_root):
            print(f"{uuid}  {d}  (mtime {datetime.fromtimestamp(d.stat().st_mtime)})")
        return 0

    target: Optional[Tuple[str, Path]] = None
    if args.uuid:
        base = get_grok_home() / "sessions" / get_workspace_slug(ws_root)
        p = base / args.uuid
        if p.exists():
            target = (args.uuid, p)
        else:
            print(f"[ERROR] UUID dir not found: {p}", file=sys.stderr)
            return 1
    elif args.latest:
        target = get_latest_session(ws_root)
        if not target:
            print("[ERROR] No Grok sessions found for this workspace", file=sys.stderr)
            return 1
    else:
        print("Specify --uuid, --latest, or --list", file=sys.stderr)
        return 2

    uuid, sdir = target
    log = build_session_log(uuid, sdir, ws_root)

    if args.print or not args.auto:
        print(log)

    if args.auto:
        # Minimal auto placement: root session-logs/ for now (real /save logic is more sophisticated).
        ts = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        out_dir = ws_root / "session-logs"
        out_dir.mkdir(parents=True, exist_ok=True)
        out_path = out_dir / f"grok-session-{ts}-{uuid[:8]}.md"
        out_path.write_text(log, encoding="utf-8")
        print(f"[OK] Wrote {out_path}")
        # In real use we would also update coord, create a checkpoint entry etc.

    return 0

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