session-picker.py

#!/usr/bin/env python3
"""AgentBox session picker — shown by ttyd on each browser connection.

Displays active tmux sessions (with copilot process status) and recent
Copilot CLI sessions (from workspace.yaml).  Auto-attaches when there is
exactly one detached tmux session.
"""

import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

# ── Colours (ANSI) ───────────────────────────────────────────────────────────
CYAN = "\033[96m"
DIM = "\033[2m"
BOLD = "\033[1m"
RESET = "\033[0m"
GREEN = "\033[92m"
YELLOW = "\033[93m"

COPILOT_SESSION_DIR = Path.home() / ".copilot" / "session-state"


# ── Helpers ──────────────────────────────────────────────────────────────────

def run(cmd: list[str], fallback: str = "") -> str:
    try:
        return subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return fallback


def time_ago(iso_str: str) -> str:
    """Convert ISO timestamp to human-readable relative time."""
    try:
        dt = datetime.fromisoformat(iso_str.replace("Z", "+00:00"))
        delta = datetime.now(timezone.utc) - dt
        secs = int(delta.total_seconds())
        if secs < 60:
            return "just now"
        if secs < 3600:
            return f"{secs // 60}m ago"
        if secs < 86400:
            return f"{secs // 3600}h ago"
        return f"{secs // 86400}d ago"
    except (ValueError, TypeError):
        return ""


def parse_workspace_yaml(path: Path) -> dict:
    """Minimal YAML parser for workspace.yaml (key: value lines only)."""
    data = {}
    try:
        for line in path.read_text().splitlines():
            if ":" in line:
                key, _, val = line.partition(":")
                data[key.strip()] = val.strip()
    except OSError:
        pass
    return data


# ── tmux session listing ─────────────────────────────────────────────────────

def get_tmux_sessions() -> list[dict]:
    """Return list of tmux sessions with metadata."""
    raw = run(["tmux", "list-sessions", "-F",
               "#{session_name}\t#{session_windows}\t#{session_created}\t#{session_attached}"])
    if not raw:
        return []

    sessions = []
    for line in raw.splitlines():
        parts = line.split("\t")
        if len(parts) < 4:
            continue
        name, windows, created_ts, attached = parts
        # Check if copilot is running in any pane of this session
        pane_cmds = run(["tmux", "list-panes", "-t", name, "-F", "#{pane_current_command}"])
        copilot_running = "copilot" in pane_cmds.lower() if pane_cmds else False

        try:
            created = datetime.fromtimestamp(int(created_ts), tz=timezone.utc)
            age = time_ago(created.isoformat())
        except (ValueError, OSError):
            age = ""

        sessions.append({
            "name": name,
            "windows": windows,
            "age": age,
            "attached": attached != "0",
            "copilot": copilot_running,
        })
    return sessions


# ── Copilot CLI session listing ──────────────────────────────────────────────

def get_copilot_sessions(limit: int = 5) -> list[dict]:
    """Return recent Copilot CLI sessions from workspace.yaml files."""
    if not COPILOT_SESSION_DIR.is_dir():
        return []

    sessions = []
    for d in COPILOT_SESSION_DIR.iterdir():
        ws = d / "workspace.yaml"
        if not ws.exists():
            continue
        data = parse_workspace_yaml(ws)
        if not data.get("id"):
            continue
        sessions.append({
            "id": data["id"],
            "summary": data.get("summary", ""),
            "repo": data.get("repository", ""),
            "branch": data.get("branch", ""),
            "updated": data.get("updated_at", data.get("created_at", "")),
        })

    # Sort by most recently updated
    sessions.sort(key=lambda s: s["updated"], reverse=True)
    return sessions[:limit]


# ── Display ──────────────────────────────────────────────────────────────────

def show_picker(tmux_sessions: list[dict], copilot_sessions: list[dict]):
    print()
    print(f"  {CYAN}{BOLD}═══ AgentBox Sessions ═══{RESET}")
    print()

    if tmux_sessions:
        print(f"  {BOLD}Terminal sessions:{RESET}")
        for i, s in enumerate(tmux_sessions, 1):
            status = ""
            if s["copilot"]:
                status = f" {GREEN}🤖 copilot running{RESET}"
            elif s["attached"]:
                status = f" {DIM}(attached){RESET}"

            print(f"    [{CYAN}{i}{RESET}] {s['name']}"
                  f"  {DIM}│ {s['windows']} win │ {s['age']}{RESET}{status}")
        print()

    if copilot_sessions:
        print(f"  {BOLD}Recent Copilot sessions (resumable):{RESET}")
        for i, s in enumerate(copilot_sessions, 1):
            label = s["summary"] or s["id"][:12]
            context = ""
            if s["repo"]:
                context = f" ({s['repo']}"
                if s["branch"]:
                    context += f", {s['branch']}"
                context += ")"
            age = time_ago(s["updated"])
            print(f"    [{CYAN}c{i}{RESET}] {label}{DIM}{context} — {age}{RESET}")
        print()

    print(f"  {BOLD}Actions:{RESET}")
    print(f"    [{CYAN}n{RESET}] New terminal session")
    print(f"    [{CYAN}c{RESET}] New session + {YELLOW}copilot --continue{RESET}")
    print(f"    [{CYAN}r{RESET}] New session + {YELLOW}copilot --resume{RESET} (picker)")
    print(f"    [{CYAN}q{RESET}] Quit")
    print()


def prompt_session_name() -> str:
    """Ask for optional session name, default to auto-generated."""
    try:
        name = input(f"  Session name [{DIM}Enter for auto{RESET}]: ").strip()
    except (EOFError, KeyboardInterrupt):
        name = ""
    if not name:
        # Generate a short name based on count of existing sessions
        existing = run(["tmux", "list-sessions", "-F", "#{session_name}"])
        count = len(existing.splitlines()) + 1 if existing else 1
        name = f"session-{count}"
    return name


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    tmux_sessions = get_tmux_sessions()
    copilot_sessions = get_copilot_sessions()

    # Auto-attach: exactly one detached tmux session → go straight in
    detached = [s for s in tmux_sessions if not s["attached"]]
    if len(detached) == 1 and len(tmux_sessions) == 1:
        print(f"  {DIM}Reconnecting to {detached[0]['name']}…{RESET}")
        os.execvp("tmux", ["tmux", "attach-session", "-t", detached[0]["name"]])
        return  # unreachable

    show_picker(tmux_sessions, copilot_sessions)

    try:
        choice = input(f"  Select: ").strip().lower()
    except (EOFError, KeyboardInterrupt):
        print()
        return

    # ── Handle choice ────────────────────────────────────────────────────
    if choice == "q":
        return

    if choice == "n":
        name = prompt_session_name()
        os.execvp("tmux", ["tmux", "new-session", "-s", name])

    elif choice == "c":
        name = prompt_session_name()
        os.execvp("tmux", ["tmux", "new-session", "-s", name, "copilot", "--continue"])

    elif choice == "r":
        name = prompt_session_name()
        os.execvp("tmux", ["tmux", "new-session", "-s", name, "copilot", "--resume"])

    elif choice.startswith("c") and choice[1:].isdigit():
        idx = int(choice[1:]) - 1
        if 0 <= idx < len(copilot_sessions):
            sid = copilot_sessions[idx]["id"]
            name = prompt_session_name()
            os.execvp("tmux", ["tmux", "new-session", "-s", name,
                                "copilot", "--resume", sid])
        else:
            print(f"  Invalid selection.")
            main()

    elif choice.isdigit():
        idx = int(choice) - 1
        if 0 <= idx < len(tmux_sessions):
            os.execvp("tmux", ["tmux", "attach-session", "-t",
                                tmux_sessions[idx]["name"]])
        else:
            print(f"  Invalid selection.")
            main()

    else:
        print(f"  Invalid selection.")
        main()


if __name__ == "__main__":
    main()