browser-cli/browser_cli/commands/session.py

import click
from browser_cli.client import active_browser_targets, remote_browser_targets, send_command, BrowserNotConnected
from rich.console import Console

console = Console()


def _handle(command, args=None, profile=None):
    try:
        return send_command(command, args or {}, profile=profile)
    except BrowserNotConnected as e:
        console.print(f"[red]Error:[/red] {e}")
        raise SystemExit(1)
    except RuntimeError as e:
        console.print(f"[red]Browser error:[/red] {e}")
        raise SystemExit(1)


def _handle_multi(command, args=None, profile=None, remote=None):
    try:
        if remote:
            return send_command(command, args or {}, profile=profile, remote=remote)
        return send_command(command, args or {}, profile=profile)
    except (BrowserNotConnected, RuntimeError):
        return None


def _multi_browser_targets():
    root = click.get_current_context().find_root()
    if root.obj.get("browser_explicit"):
        return []
    remote = root.obj.get("remote")
    key = root.obj.get("key")
    if remote:
        targets = remote_browser_targets(remote, key=key)
    else:
        targets = active_browser_targets(key=key)
    if len(targets) <= 1 and not any(target.remote for target in targets):
        return []
    return targets


@click.group("session")
def session_group():
    """Save and restore browser sessions."""


@session_group.command("save")
@click.argument("name")
def session_save(name):
    """Save all current tabs as session NAME."""
    result = _handle("session.save", {"name": name})
    count = result.get("tabs", 0) if isinstance(result, dict) else 0
    console.print(f"[green]Session '{name}' saved[/green] ({count} tabs)")


@session_group.command("load")
@click.argument("name")
def session_load(name):
    """Restore session NAME (opens all saved tabs)."""
    result = _handle("session.load", {"name": name})
    count = result.get("tabs", 0) if isinstance(result, dict) else 0
    console.print(f"[green]Session '{name}' loaded[/green] ({count} tabs opened)")


@session_group.command("diff")
@click.argument("name_a")
@click.argument("name_b")
def session_diff(name_a, name_b):
    """Show tabs added/removed between two saved sessions."""
    diff = _handle("session.diff", {"nameA": name_a, "nameB": name_b})
    if not diff:
        console.print("[yellow]No diff data returned[/yellow]")
        return

    added = diff.get("added") or []
    removed = diff.get("removed") or []

    if added:
        console.print(f"[green]Added in '{name_b}':[/green]")
        for url in added:
            console.print(f"  + {url}")

    if removed:
        console.print(f"[red]Removed in '{name_b}':[/red]")
        for url in removed:
            console.print(f"  - {url}")

    if not added and not removed:
        console.print("[green]Sessions are identical[/green]")


@session_group.command("list")
def session_list():
    """List all saved sessions."""
    from rich.table import Table
    targets = _multi_browser_targets()
    show_browser = bool(targets)
    if targets:
        sessions = []
        for target in targets:
            result = _handle_multi("session.list", profile=target.profile, remote=target.remote)
            if result is None:
                continue
            sessions.extend({**session, "browser": target.display_name} for session in result)
        if not sessions:
            console.print("[red]Error:[/red] Cannot resolve a browser socket automatically.")
            raise SystemExit(1)
    else:
        sessions = _handle("session.list")
    if not sessions:
        console.print("[yellow]No saved sessions[/yellow]")
        return
    table = Table(show_header=True, header_style="bold cyan")
    if show_browser:
        table.add_column("Browser")
    table.add_column("Name")
    table.add_column("Tabs", width=6)
    table.add_column("Saved at")
    for s in sessions:
        from datetime import datetime
        saved = datetime.fromtimestamp(s["savedAt"] / 1000).strftime("%Y-%m-%d %H:%M") if s.get("savedAt") else ""
        row = [s.get("browser", "")] if show_browser else []
        row.extend([s["name"], str(s["tabs"]), saved])
        table.add_row(*row)
    console.print(table)


@session_group.command("remove")
@click.argument("name")
def session_remove(name):
    """Delete a saved session."""
    _handle("session.remove", {"name": name})
    console.print(f"[green]Session '{name}' removed[/green]")


@session_group.command("auto-save")
@click.argument("state", type=click.Choice(["on", "off"]))
def session_auto_save(state):
    """Enable or disable automatic session saving."""
    enabled = state == "on"
    _handle("session.auto_save", {"enabled": enabled})
    console.print(f"[green]Auto-save {state}[/green]")