browser-cli/browser_cli/commands/session.py

import json
from pathlib import Path
import click
from browser_cli.commands import client_from_ctx, gentle_mode_option, handle_errors
from rich.console import Console

console = Console()

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

@session_group.command("save")
@click.argument("name")
@handle_errors
def session_save(name):
    """Save all current tabs as session NAME."""
    result = client_from_ctx().session.save(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")
@gentle_mode_option("Throttle mode for large restores.")
@click.option("--discard-background-tabs", is_flag=True, help="Discard restored background tabs after opening to reduce load.")
@click.option("--lazy", is_flag=True, help="Create lightweight placeholder tabs after --eager-tabs; placeholders load when selected.")
@click.option("--eager-tabs", type=int, default=10, show_default=True, help="Number of real tabs to open before lazy placeholders.")
@click.option("--background", "background_job", is_flag=True, help="Start restore as a background job and return immediately.")
@handle_errors
def session_load(name, gentle_mode, discard_background_tabs, lazy, eager_tabs, background_job):
    """Restore session NAME (opens all saved tabs)."""
    b = client_from_ctx()
    if background_job:
        result = b.session.load_background(
            name, gentle_mode=gentle_mode, discard_background_tabs=discard_background_tabs,
            lazy=lazy, eager_tabs=eager_tabs,
        )
        if isinstance(result, dict) and result.get("jobId"):
            console.print(f"[green]Session restore started[/green] job={result['jobId']}")
            return
    else:
        result = b.session.load(
            name, gentle_mode=gentle_mode, discard_background_tabs=discard_background_tabs,
            lazy=lazy, eager_tabs=eager_tabs,
        )
    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("export")
@click.argument("name", required=False)
@click.option("-o", "output", type=click.Path(dir_okay=False, path_type=Path), default=None, help="Write JSON to file instead of stdout")
@handle_errors
def session_export(name, output):
    """Export one saved session, or all sessions as JSON."""
    data = client_from_ctx().session.export(name)
    text = json.dumps(data, indent=2, sort_keys=True)
    if output:
        output.write_text(text + "\n", encoding="utf-8")
        console.print(f"[green]Exported session data to {output}[/green]")
    else:
        click.echo(text)

@session_group.command("import")
@click.argument("name")
@click.argument("file", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--overwrite", is_flag=True, help="Replace an existing saved session")
@handle_errors
def session_import(name, file, overwrite):
    """Import a saved session JSON file."""
    payload = json.loads(file.read_text(encoding="utf-8"))
    session = payload.get("session", payload) if isinstance(payload, dict) else payload
    result = client_from_ctx().session.import_(name, session, overwrite=overwrite)
    count = result.get("tabs", 0) if isinstance(result, dict) else 0
    console.print(f"[green]Imported session '{name}'[/green] ({count} tabs)")

@session_group.command("diff")
@click.argument("name_a")
@click.argument("name_b")
@handle_errors
def session_diff(name_a, name_b):
    """Show tabs added/removed between two saved sessions."""
    diff = client_from_ctx().session.diff(name_a, 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")
@handle_errors
def session_list():
    """List all saved sessions."""
    from datetime import datetime
    from rich.table import Table
    sessions = client_from_ctx().session.list()
    if not sessions:
        console.print("[yellow]No saved sessions[/yellow]")
        return
    show_browser = any("browser" in s for s in sessions)
    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:
        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")
@handle_errors
def session_remove(name):
    """Delete a saved session."""
    client_from_ctx().session.remove(name)
    console.print(f"[green]Session '{name}' removed[/green]")

@session_group.command("job-status")
@click.argument("job_id")
@handle_errors
def session_job_status(job_id):
    """Show status for a background session job."""
    result = client_from_ctx().perf.job_status(job_id)
    status = result.get("status", "unknown")
    console.print(f"[bold]{job_id}[/bold]: {status}")
    if result.get("error"):
        console.print(f"[red]{result['error']}[/red]")
    elif result.get("result"):
        console.print(result["result"])

@session_group.command("job-cancel")
@click.argument("job_id")
@handle_errors
def session_job_cancel(job_id):
    """Cancel a running background job."""
    client_from_ctx().perf.job_cancel(job_id)
    console.print(f"[green]Cancel requested for {job_id}[/green]")

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