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. `kitty --session <past day>` (the kitty-days tv
  cable) opens that day as a read-only snapshot instead.
- 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.
- kitten `day_session.py flush` (sent over remote control by kitty-day-save
  before shutdown/reboot): saves right away, then holds further saves for
  FREEZE seconds so windows dying during shutdown can't shrink the saved day.
"""
import fcntl
import json
import os
import re
import shlex
import sys
import tempfile
from datetime import date, timedelta
from time import monotonic

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
FREEZE = 60  # seconds without saves after a flush; long enough for any shutdown, short if it's cancelled
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)

    frozen_until = 0.0

    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 or monotonic() < frozen_until:
            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)

    def flush() -> None:
        nonlocal frozen_until
        frozen_until = 0.0
        save(0)
        frozen_until = monotonic() + FREEZE

    _schedule_save = schedule
    boss.day_session_flush = flush  # the flush kitten runs in its own module copy, so it finds this via boss
    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:
    if 'flush' in args and (flush := getattr(boss, 'day_session_flush', None)):
        flush()


def session_arg() -> str | None:
    """`--session` from kitty's own command line (geninclude runs inside kitty)."""
    argv = sys.argv[1:]
    for i, a in enumerate(argv):
        if a.startswith('--session='):
            return a.partition('=')[2]
        if a == '--session' and i + 1 < len(argv):
            return argv[i + 1]
    return None


if __name__ == '__main__':
    path = today_path()
    other = session_arg()
    # An explicit --session for another day must not claim today's lock.
    if (other is None or os.path.abspath(os.path.expanduser(other)) == path) and 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}')