dotfiles/dot_config/kitty/day_session.py

"""Daily kitty sessions, archived per week.

    ~/.local/share/kitty-sessions/<monday YYYY-MM-DD>/<day YYYY-MM-DD>.kitty-session
    ~/.local/share/kitty-sessions/<monday YYYY-MM-DD>/<day YYYY-MM-DD>.scrollback/<window key>.ansi

One file, three entry points (kitty runs it under a different __name__ each):
- geninclude (__main__): creates today's session file and points
  startup_session at it whenever no running kitty owns today's session (an
  flock held for the owner's lifetime). Every launch after kitty was closed
  reopens the day; a second concurrent instance starts plain instead of
  duplicating today's tabs.
- watcher (__kitty_watcher__): SAVE_DELAY seconds after a window opens,
  closes, moves, gains focus or finishes a command (and every SAVE_EVERY
  seconds, for output from long-running programs), rewrites the session this
  instance started with: tabs/splits/cwds, each shell window's scrollback, and
  the program running in it. On restore the shell `cat`s the scrollback, then
  re-runs the program. Days opened later via the picker are left as read-only
  snapshots.
- kitten (map ... kitten day_session.py): picker over every archived day,
  newest first; the choice is opened with goto_session.
"""
import fcntl
import glob
import json
import os
import re
import shlex
import tempfile
from datetime import date, timedelta

from kittens.tui.handler import result_handler

ROOT = os.path.expanduser('~/.local/share/kitty-sessions')
SAVE_DELAY = 1  # seconds; batches bursts of events, and the most change lost when kitty dies
SAVE_EVERY = 30  # seconds; catches output that arrives without any window event
SB_VAR = 'day_scrollback'  # user var naming a window's scrollback dump; kitty saves/restores it via launch --var
UNSERIALIZE = 'kitty-unserialize-data='  # kitty.constants.unserialize_launch_flag
# Never replay a command that was mid-run when the session was saved if any word
# of it is one of these: re-running them on the next launch can destroy data.
NO_REPLAY = {
    'rm', 'rmdir', 'mv', 'dd', 'shred', 'truncate', 'wipefs', 'fdisk', 'parted',
    'sudo', 'doas', 'su', 'run0', 'pkexec', 'kill', 'pkill', 'killall',
    'reboot', 'shutdown', 'poweroff',
}
SGR = re.compile(r'\x1b\[[0-9;:]*m')
PROMPT_MARK = '\x1b]133;A'  # shell integration's start-of-prompt mark
_schedule_save = None  # set by on_load when this instance owns a day session


def today_path() -> str:
    d = date.today()
    return os.path.join(ROOT, str(d - timedelta(days=d.weekday())), f'{d}.kitty-session')


def own_day(path: str) -> bool:
    # Lock a side file, not the session itself: atomic_save replaces that inode.
    # The fd stays open (and non-inheritable) so the lock dies with this kitty.
    lock = os.path.join(os.environ.get('XDG_RUNTIME_DIR') or tempfile.gettempdir(),
                        f'kitty-day-{os.path.basename(path)}.lock')
    fd = os.open(lock, os.O_RDONLY | os.O_CREAT, 0o600)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        os.close(fd)
        return False
    return True


def replay_safe(cmd: str) -> bool:
    try:
        words = shlex.split(cmd)
    except ValueError:
        return False
    return not any((b := os.path.basename(w)) in NO_REPLAY or b.startswith('mkfs') for w in words)


def scrollback(window) -> str:
    from kitty.window import as_text

    s = window.screen
    # While a full-screen program (nvim) is up, the shell's output lives in the other buffer.
    text = as_text(s, as_ansi=True, add_history=True, alternate_screen=s.is_using_alternate_linebuf())
    if window.at_prompt and (i := text.rfind(PROMPT_MARK)) >= 0:
        text = text[:i]  # the idle prompt; the restored shell draws a fresh one
    lines = text.split('\n')
    while lines and not SGR.sub('', lines[-1]).strip():  # the unused rows below the cursor
        lines.pop()
    return '\n'.join(lines) + '\x1b[m\n' if lines else ''


def on_load(boss, data) -> None:
    from uuid import uuid4

    from kitty.config import atomic_save
    from kitty.fast_data_types import add_timer, get_options
    from kitty.session import parse_save_as_options_spec_args, session_arg_to_name

    global _schedule_save
    path = get_options().startup_session or ''
    if not path.startswith(ROOT + os.sep):
        return
    name = session_arg_to_name(path)
    sb_dir = path.rpartition('.')[0] + '.scrollback'
    ser_opts = parse_save_as_options_spec_args(
        ['--use-foreground-process', f'--match=session:^{re.escape(name)}$'])[0]
    last = ''
    dumps: dict[str, str] = {}
    pending = False

    def launch_parts(line: str) -> tuple[list[str], dict] | None:
        args = shlex.split(line) if line.startswith('launch ') else []
        if len(args) < 2 or not args[1].startswith(UNSERIALIZE):
            return None
        return args, json.loads(args[1][len(UNSERIALIZE):])

    # Per-window shell startup command (scrollback cat + program), by SB_VAR key.
    # Seeded from the file so restored windows keep theirs until their shell is up.
    startup: dict[str, str] = {}
    with open(path) as f:
        for args, data in filter(None, map(launch_parts, f.read().splitlines())):
            for a in args:
                if a.startswith(f'--var={SB_VAR}=') and data.get('cmd_at_shell_startup'):
                    startup[a.removeprefix(f'--var={SB_VAR}=')] = data['cmd_at_shell_startup']

    def startup_cmd(line: str) -> str:
        # Rewrite kitty's per-window restore data: cat the scrollback first, then
        # re-run the foreground program unless it is on the NO_REPLAY list.
        parts = launch_parts(line)
        if parts is None:
            return line
        args, data = parts
        w = boss.window_id_map.get(data['id'])
        key = w.user_vars.get(SB_VAR, '') if w else ''
        fg = data.pop('cmd_at_shell_startup', None)
        # Until the shell draws a prompt, kitty's "foreground process" is shell
        # init (mise/direnv hooks) or the replayed program: keep the previous value.
        if w and w.screen.last_reported_cwd:
            steps = []
            dump = os.path.join(sb_dir, f'{key}.ansi')
            if os.path.exists(dump):
                steps.append(f'cat {shlex.quote(dump)}')
            if fg:
                fg = fg if isinstance(fg, str) else shlex.join(fg)
                if replay_safe(fg):
                    steps.append(fg)
            startup[key] = '; '.join(steps)
        if startup.get(key):
            data['cmd_at_shell_startup'] = startup[key]
        args[1] = UNSERIALIZE + json.dumps(data)
        return shlex.join(args)

    def save(timer_id: int) -> None:
        nonlocal last, pending
        pending = False
        windows = [w for w in boss.all_windows if w.created_in_session_name == name]
        # Nothing left in the session means kitty is closing: keep the last good save.
        if not windows:
            return
        os.makedirs(sb_dir, exist_ok=True)
        keep = set()
        for w in windows:
            key = w.user_vars.get(SB_VAR)
            if not key:
                key = uuid4().hex
                w.set_user_var(SB_VAR, key)
            keep.add(f'{key}.ansi')
            # Only once the shell has drawn a prompt (and reported its cwd): a
            # restored window cats its old dump before that, so writing earlier
            # could replace the history before it was replayed.
            if w.screen.last_reported_cwd:
                text = scrollback(w)
                if dumps.get(key) != text:
                    atomic_save(text.encode(), os.path.join(sb_dir, f'{key}.ansi'))
                    dumps[key] = text
        for f in set(os.listdir(sb_dir)) - keep:  # windows closed since the last save
            os.remove(os.path.join(sb_dir, f))
        text = '\n'.join(map(startup_cmd, boss.serialize_state_as_session(path, ser_opts)))
        if text != last:
            atomic_save(text.encode(), path)
            last = text

    def schedule() -> None:
        nonlocal pending
        if not pending:
            pending = True
            add_timer(save, SAVE_DELAY, False)

    _schedule_save = schedule
    add_timer(lambda timer_id: schedule(), SAVE_EVERY, True)


def _changed(boss, window, data) -> None:
    if _schedule_save is not None:
        _schedule_save()


# on_resize covers windows/tabs being created, closed or re-split; on_cmd_startstop
# catches `cd` and finished commands. Title changes are skipped: spinners would
# re-save every second.
on_resize = on_close = on_focus_change = on_cmd_startstop = _changed


def main(args: list[str]) -> None:
    pass


@result_handler(no_ui=True)
def handle_result(args: list[str], answer: str, target_window_id: int, boss) -> None:
    from kitty.session import goto_session

    days = sorted(glob.glob(os.path.join(ROOT, '*', '*.kitty-session')), key=os.path.basename, reverse=True)
    if not days:
        boss.show_error('No day sessions', f'Nothing archived under {ROOT} yet.')
        return

    def label(path: str) -> str:
        d = date.fromisoformat(os.path.basename(path).partition('.')[0])
        return f'{d}  {d:%a}  week of {os.path.basename(os.path.dirname(path))}'

    boss.choose_entry('Open a day', ((p, label(p)) for p in days), lambda p: p and goto_session(boss, (p,)))


if __name__ == '__main__':
    path = today_path()
    if own_day(path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        try:
            open(path, 'x').close()  # empty session = one tab with the default shell
        except FileExistsError:
            pass
        print(f'startup_session {path}')