add pir: fuzzy-find and resume pi sessions for the current dir

`pi -c` blindly grabs the newest session, and the alternative is digging
through ~/.pi/agent/sessions/--home-user-dev-foo--/ for a UUID by hand.
This lists every session for the cwd in an fzf picker with a live transcript
preview and execs `pi --session <file>` on the pick.

  pir              pick a session for this dir
  pir -a           pick across every project dir
  pir <query...>   start fzf with an initial query
  pir -g TEXT      only sessions whose full transcript contains TEXT

The -g pass exists because the fzf query matches the list line -- title and
id -- not deep history, so "the session where I fixed the mouse bug" is
unfindable without it. It shells out to rg when available and falls back to
grep.

In the picker: enter resumes, ctrl-f forks instead (execs --fork, leaving the
original untouched), ctrl-o prints the path and exits without launching,
ctrl-a toggles current-dir vs all-dirs, ctrl-s toggles recent vs longest.
The toggles are fzf --bind=reload of this same script, which is what --list
and --preview are for -- they are internal modes invoked by fzf, not user
entry points.

Session directory resolution mirrors pi's own getDefaultSessionDirPath()
encoding, and honours PI_APP_NAME / <APP>_CODING_AGENT_DIR, so a renamed or
relocated agent dir still resolves. Forked sessions are detected from the
header's parentSession and marked in the list.

Requires fzf; checked up front with a clear error rather than failing inside
a pipeline.

1 files changed, 537 insertions, 0 deletions

diff --git a/pi-ext/scripts/pi-resume b/pi-ext/scripts/pi-resume
new file mode 100755
index 0000000..d327362
--- /dev/null
+++ b/pi-ext/scripts/pi-resume
@@ -0,0 +1,537 @@
+#!/usr/bin/env python3
+"""
+pi-resume (aka `pir`) - fuzzy-find and resume old pi sessions for the current dir.
+
+Instead of `pi -c` (blindly grabs the newest) or digging through
+~/.pi/agent/sessions/--home-user-dev-foo--/ for a UUID by hand, this lists every
+session for the cwd in an fzf picker with a live transcript preview, then execs
+`pi --session <file>` on the one you pick.
+
+Usage:
+  pir                 pick a session for the current dir, resume it
+  pir -a/--all        pick across every project dir
+  pir -q QUERY        start fzf with an initial query
+  pir <query...>      same, bare words are treated as the query
+  pir -g TEXT         only sessions whose full transcript contains TEXT
+                      (the fzf query alone matches title/id, not deep history)
+
+In fzf:
+  enter     resume the session
+  ctrl-f    fork it instead (leaves the original untouched)
+  ctrl-o    print the session file path and exit (no launch)
+  ctrl-a    toggle current-dir <-> all-dirs
+  ctrl-s    toggle sort: recent <-> longest
+  ctrl-/    toggle preview
+  ?         toggle preview window size
+
+Internal modes (used by fzf itself): --list, --preview <file>
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import shlex
+import shutil
+import subprocess
+import sys
+import time
+from datetime import datetime, timezone
+
+# ---------------------------------------------------------------- paths
+
+APP = os.environ.get("PI_APP_NAME", "pi")
+
+
+def agent_dir() -> str:
+    env = os.environ.get(f"{APP.upper()}_CODING_AGENT_DIR")
+    if env:
+        return os.path.expanduser(env)
+    return os.path.join(os.path.expanduser("~"), f".{APP}", "agent")
+
+
+def sessions_root() -> str:
+    return os.path.join(agent_dir(), "sessions")
+
+
+def encode_cwd(cwd: str) -> str:
+    """Mirror pi's getDefaultSessionDirPath() encoding exactly."""
+    resolved = os.path.realpath(os.path.expanduser(cwd))
+    safe = re.sub(r"[/\\:]", "-", re.sub(r"^[/\\]", "", resolved))
+    return f"--{safe}--"
+
+
+def session_dir_for(cwd: str) -> str:
+    env = os.environ.get(f"{APP.upper()}_CODING_AGENT_SESSION_DIR")
+    if env:
+        return os.path.expanduser(env)
+    return os.path.join(sessions_root(), encode_cwd(cwd))
+
+
+CACHE = os.path.join(
+    os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")),
+    "pi-resume",
+    "index.json",
+)
+
+# ---------------------------------------------------------------- scanning
+
+MSG_MARK = b'"type":"message"'
+INFO_MARK = b'"type":"session_info"'
+HEAD_BYTES = 512 * 1024
+
+
+def _text_of(msg: dict) -> str:
+    """Flatten a message's content into plain text (ignore thinking/tool noise)."""
+    content = msg.get("content")
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return ""
+    out = []
+    for part in content:
+        if isinstance(part, dict) and part.get("type") == "text":
+            out.append(part.get("text") or "")
+    return " ".join(out)
+
+
+def _clean(s: str, limit: int = 400) -> str:
+    s = re.sub(r"\s+", " ", (s or "")).strip()
+    return s[:limit]
+
+
+def scan(path: str) -> dict | None:
+    """Cheap metadata scan: header, msg count, name, first user message."""
+    try:
+        st = os.stat(path)
+    except OSError:
+        return None
+
+    msgs = 0
+    has_info = False
+    head = bytearray()
+    tail = b""
+    keep = max(len(MSG_MARK), len(INFO_MARK)) - 1
+
+    try:
+        with open(path, "rb") as f:
+            while True:
+                chunk = f.read(1 << 20)
+                if not chunk:
+                    break
+                if len(head) < HEAD_BYTES:
+                    head += chunk[: HEAD_BYTES - len(head)]
+                buf = tail + chunk
+                msgs += buf.count(MSG_MARK)
+                if not has_info and INFO_MARK in buf:
+                    has_info = True
+                tail = buf[-keep:] if keep else b""
+    except OSError:
+        return None
+
+    header, first_user, model = {}, "", ""
+    for line in head.split(b"\n"):
+        if not line.strip():
+            continue
+        try:
+            rec = json.loads(line)
+        except Exception:
+            continue  # truncated final line of the head window
+        t = rec.get("type")
+        if t == "session" and not header:
+            header = rec
+        elif t == "model_change" and not model:
+            model = rec.get("modelId") or ""
+        elif t == "message" and not first_user:
+            m = rec.get("message") or {}
+            if m.get("role") == "user":
+                txt = _clean(_text_of(m))
+                if txt:
+                    first_user = txt
+        if header and first_user and model:
+            break
+
+    name = ""
+    if has_info:
+        # Rare; a full line pass is fine here. Last one wins.
+        try:
+            with open(path, "rb") as f:
+                for line in f:
+                    if INFO_MARK not in line:
+                        continue
+                    try:
+                        rec = json.loads(line)
+                    except Exception:
+                        continue
+                    if rec.get("type") == "session_info" and rec.get("name"):
+                        name = _clean(rec["name"], 120)
+        except OSError:
+            pass
+
+    return {
+        "path": path,
+        "id": header.get("id") or os.path.basename(path).split("_", 1)[-1][:-6],
+        "cwd": header.get("cwd") or "",
+        "forked": bool(header.get("parentSession")),
+        "created": header.get("timestamp") or "",
+        "mtime": st.st_mtime,
+        "size": st.st_size,
+        "msgs": msgs,
+        "name": name,
+        "model": model,
+        "first": first_user,
+    }
+
+
+def load_cache() -> dict:
+    try:
+        with open(CACHE) as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def save_cache(cache: dict) -> None:
+    try:
+        os.makedirs(os.path.dirname(CACHE), exist_ok=True)
+        tmp = CACHE + f".{os.getpid()}.tmp"
+        with open(tmp, "w") as f:
+            json.dump(cache, f)
+        os.replace(tmp, CACHE)
+    except Exception:
+        pass
+
+
+def grep_filter(files: list[str], pattern: str) -> set[str]:
+    """Full-text prefilter over raw session files. Returns matching paths."""
+    if not files:
+        return set()
+    tool = "rg" if shutil.which("rg") else "grep"
+    cmd = [tool, "-l", "-i", "-F", "--", pattern] if tool == "rg" else [
+        "grep", "-l", "-i", "-F", "-e", pattern, "--"
+    ]
+    hits: set[str] = set()
+    # chunk to stay under ARG_MAX
+    for i in range(0, len(files), 400):
+        batch = files[i : i + 400]
+        try:
+            p = subprocess.run(cmd + batch, capture_output=True, text=True)
+        except OSError:
+            return set(files)
+        hits.update(l.strip() for l in p.stdout.splitlines() if l.strip())
+    return hits
+
+
+def collect(all_dirs: bool, grep: str = "") -> list[dict]:
+    if all_dirs:
+        root = sessions_root()
+        dirs = [os.path.join(root, d) for d in os.listdir(root)] if os.path.isdir(root) else []
+    else:
+        dirs = [session_dir_for(os.getcwd())]
+
+    files = []
+    for d in dirs:
+        if not os.path.isdir(d):
+            continue
+        for fn in os.listdir(d):
+            if fn.endswith(".jsonl"):
+                files.append(os.path.join(d, fn))
+
+    cache = load_cache()
+    out, dirty = [], False
+    for p in files:
+        try:
+            st = os.stat(p)
+        except OSError:
+            continue
+        hit = cache.get(p)
+        if hit and hit.get("mtime") == st.st_mtime and hit.get("size") == st.st_size:
+            out.append(hit)
+            continue
+        rec = scan(p)
+        if rec:
+            cache[p] = rec
+            out.append(rec)
+            dirty = True
+
+    live = set(files)
+    for stale in [k for k in cache if k not in live]:
+        del cache[stale]
+        dirty = True
+    if dirty:
+        save_cache(cache)
+
+    out = [r for r in out if r.get("msgs")]
+    if grep:
+        keep = grep_filter([r["path"] for r in out], grep)
+        out = [r for r in out if r["path"] in keep]
+    return out
+
+
+# ---------------------------------------------------------------- display
+
+C_DIM = "\x1b[38;5;242m"
+C_AGE = "\x1b[38;5;110m"
+C_NUM = "\x1b[38;5;108m"
+C_NAME = "\x1b[38;5;215m"
+C_TXT = "\x1b[0m"
+C_HDR = "\x1b[1;38;5;215m"
+C_USER = "\x1b[1;38;5;110m"
+C_ASST = "\x1b[38;5;250m"
+C_TOOL = "\x1b[38;5;242m"
+R = "\x1b[0m"
+
+
+def age(ts: float) -> str:
+    d = max(0.0, time.time() - ts)
+    if d < 3600:
+        return f"{int(d // 60)}m"
+    if d < 86400:
+        return f"{int(d // 3600)}h"
+    if d < 86400 * 7:
+        return f"{int(d // 86400)}d"
+    return datetime.fromtimestamp(ts).strftime("%b %d")
+
+
+def age_phrase(ts: float) -> str:
+    """Human phrasing: '3h ago' for recent, absolute date for old."""
+    if time.time() - ts < 86400 * 7:
+        return f"{age(ts)} ago"
+    return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M")
+
+
+def fmt_line(r: dict, show_cwd: bool) -> str:
+    title = r["name"] or r["first"] or "(no messages)"
+    label = f"{C_NAME}{title}{R}" if r["name"] else f"{C_TXT}{title}{R}"
+    bits = [
+        f"{C_AGE}{age(r['mtime']):>6}{R}",
+        f"{C_NUM}{r['msgs']:>4} msg{R}",
+    ]
+    if show_cwd:
+        home = os.path.expanduser("~")
+        cwd = r["cwd"] or "?"
+        cwd = "~" + cwd[len(home):] if cwd.startswith(home) else cwd
+        bits.append(f"{C_DIM}{cwd:<28}{R}")
+    if r["forked"]:
+        bits.append(f"{C_DIM}fork{R}")
+    return "  ".join(bits) + "  " + label
+
+
+def emit_list(all_dirs: bool, sort: str, grep: str = "") -> None:
+    rows = collect(all_dirs, grep)
+    if sort == "longest":
+        rows.sort(key=lambda r: r["msgs"], reverse=True)
+    else:
+        rows.sort(key=lambda r: r["mtime"], reverse=True)
+    buf = []
+    for r in rows:
+        # field 1 = path (hidden from display via --with-nth), field 2 = pretty
+        # searchable haystack is appended dim so fzf can match on it too
+        hay = f"{r['id']} {r['first']}"
+        buf.append(f"{r['path']}\t{fmt_line(r, all_dirs)}\t{C_DIM}{_clean(hay, 300)}{R}")
+    sys.stdout.write("\n".join(buf) + ("\n" if buf else ""))
+
+
+# ---------------------------------------------------------------- preview
+
+TAIL_BYTES = 1_500_000
+
+
+def preview(path: str) -> None:
+    try:
+        st = os.stat(path)
+    except OSError:
+        print("gone")
+        return
+
+    truncated = False
+    with open(path, "rb") as f:
+        if st.st_size <= 2_000_000:
+            data = f.read()
+        else:
+            head = f.read(256 * 1024)
+            f.seek(-TAIL_BYTES, os.SEEK_END)
+            tailbuf = f.read()
+            # drop partial lines at both seams
+            head = head[: head.rfind(b"\n") + 1]
+            tailbuf = tailbuf[tailbuf.find(b"\n") + 1 :]
+            data = head + b"\n@@TRUNC@@\n" + tailbuf
+            truncated = True
+
+    meta, out = {}, []
+    for raw in data.split(b"\n"):
+        if not raw.strip():
+            continue
+        if raw == b"@@TRUNC@@":
+            out.append(f"{C_TOOL}{'-' * 20} [ ... middle of session omitted ... ] {'-' * 20}{R}")
+            continue
+        try:
+            rec = json.loads(raw)
+        except Exception:
+            continue
+        t = rec.get("type")
+        if t == "session":
+            meta = rec
+        elif t == "session_info" and rec.get("name"):
+            meta["name"] = rec["name"]
+        elif t == "compaction":
+            out.append(f"{C_TOOL}--- compacted ---{R}")
+        elif t == "message":
+            m = rec.get("message") or {}
+            role = m.get("role")
+            txt = _clean(_text_of(m), 100000)
+            if role == "user":
+                if not txt:
+                    continue
+                out.append("")
+                out.append(f"{C_USER}> {txt[:1500]}{R}")
+            elif role == "assistant":
+                tools = [
+                    p.get("name")
+                    for p in (m.get("content") or [])
+                    if isinstance(p, dict) and p.get("type") == "toolCall"
+                ]
+                if txt:
+                    out.append(f"{C_ASST}  {txt[:600]}{R}")
+                if tools:
+                    seen = ", ".join(dict.fromkeys(t for t in tools if t))
+                    out.append(f"{C_TOOL}  [{seen}]{R}")
+
+    created = meta.get("timestamp", "")
+    if created:
+        try:
+            created = (
+                datetime.fromisoformat(created.replace("Z", "+00:00"))
+                .astimezone()
+                .strftime("%Y-%m-%d %H:%M")
+            )
+        except Exception:
+            pass
+
+    print(f"{C_HDR}{meta.get('name') or os.path.basename(path)}{R}")
+    print(f"{C_DIM}id      {meta.get('id', '?')}{R}")
+    print(f"{C_DIM}cwd     {meta.get('cwd', '?')}{R}")
+    print(
+        f"{C_DIM}started {created}   modified {age_phrase(st.st_mtime)}"
+        f"   {st.st_size // 1024} KB{R}"
+    )
+    if truncated:
+        print(f"{C_DIM}(large session: showing head + tail){R}")
+    print()
+    print("\n".join(out))
+
+
+# ---------------------------------------------------------------- main
+
+def die(msg: str, code: int = 1):
+    print(f"pir: {msg}", file=sys.stderr)
+    sys.exit(code)
+
+
+def main() -> None:
+    argv = sys.argv[1:]
+    self_path = os.path.abspath(__file__)
+
+    # internal modes invoked by fzf
+    if argv and argv[0] == "--preview":
+        preview(argv[1])
+        return
+    if argv and argv[0] == "--list":
+        g = ""
+        if "--grep" in argv:
+            gi = argv.index("--grep")
+            g = argv[gi + 1] if gi + 1 < len(argv) else ""
+        emit_list("--all" in argv, "longest" if "--longest" in argv else "recent", g)
+        return
+
+    all_dirs, query, grep, words = False, "", "", []
+    i = 0
+    while i < len(argv):
+        a = argv[i]
+        if a in ("-a", "--all"):
+            all_dirs = True
+        elif a in ("-q", "--query"):
+            i += 1
+            query = argv[i] if i < len(argv) else ""
+        elif a in ("-g", "--grep"):
+            i += 1
+            grep = argv[i] if i < len(argv) else ""
+        elif a in ("-h", "--help"):
+            print(__doc__.strip())
+            return
+        else:
+            words.append(a)
+        i += 1
+    if words and not query:
+        query = " ".join(words)
+
+    if not shutil.which("fzf"):
+        die("fzf not found in PATH")
+    pi_bin = os.environ.get("PI_BIN", APP)
+    if not shutil.which(pi_bin):
+        die(f"'{pi_bin}' not found in PATH")
+
+    rows = collect(all_dirs, grep)
+    if not rows:
+        where = "any project" if all_dirs else os.getcwd()
+        extra = f" matching {grep!r}" if grep else ""
+        die(f"no sessions found for {where}{extra}", 1)
+
+    gargs = ["--grep", grep] if grep else []
+    listing = subprocess.run(
+        [sys.executable, self_path, "--list"] + (["--all"] if all_dirs else []) + gargs,
+        capture_output=True,
+        text=True,
+    ).stdout
+
+    gsuffix = f" --grep {shlex.quote(grep)}" if grep else ""
+    base = f"{shlex.quote(sys.executable)} {shlex.quote(self_path)} --list"
+    fzf_args = [
+        "fzf",
+        "--ansi",
+        "--delimiter=\t",
+        "--with-nth=2..",
+        "--layout=reverse",
+        "--info=inline",
+        "--border",
+        "--height=100%",
+        f"--query={query}",
+        "--prompt=session> ",
+        "--header="
+        "enter resume | ctrl-f fork | ctrl-o print path | ctrl-a all dirs | ctrl-s sort",
+        f"--preview={shlex.quote(sys.executable)} {shlex.quote(self_path)} --preview {{1}}",
+        "--preview-window=right,55%,wrap",
+        "--bind=ctrl-/:toggle-preview",
+        "--bind=?:change-preview-window(right,80%|right,55%)",
+        f"--bind=ctrl-a:reload({base} --all{gsuffix})",
+        f"--bind=ctrl-s:reload({base}{' --all' if all_dirs else ''}{gsuffix} --longest)",
+        "--expect=ctrl-f,ctrl-o",
+    ]
+
+    proc = subprocess.run(fzf_args, input=listing, capture_output=True, text=True)
+    if proc.returncode not in (0, 1) or not proc.stdout.strip():
+        sys.exit(130)  # cancelled
+
+    out = proc.stdout.split("\n")
+    key = out[0].strip()
+    pick = next((l for l in out[1:] if l.strip()), "")
+    if not pick:
+        sys.exit(130)
+    path = pick.split("\t", 1)[0]
+
+    if key == "ctrl-o":
+        print(path)
+        return
+
+    flag = "--fork" if key == "ctrl-f" else "--session"
+    os.execvp(pi_bin, [pi_bin, flag, path])
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        sys.exit(130)
+    except BrokenPipeError:
+        sys.exit(0)