browser-cli/extension/src/commands/session.ts

import { webExtApi as api } from '../browser-api';
import { getLargeOperationThrottle, getProfileAlias, getSessionTabs, getSessions, groupTabs, normalizeGroupColor, runLargeOperation, throwIfJobCancelled, updateJobProgress, updateTabGroup, yieldForLargeOperation } from '../core';
import { CommandGroup } from '../classes/CommandGroup';
import { AutoSaveManager } from './autosave';
import { captureCurrentSession } from './session-snapshot';
import type { CommandEntry } from '../classes/CommandGroup';
import type { Json, LazySessionMap, SessionSaveArgs, SessionRemoveArgs, SessionDiffArgs, SessionImportArgs, SessionLoadArgs, SessionAutoSaveArgs, ClientsRenameProfileArgs, StoredSession } from '../types';

function lazyPlaceholderUrl(url: string) {
  const escaped = String(url).replace(/[&<>"']/g, ch => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;", "\"": "&quot;", "'": "&#39;" }[ch]));
  const html = `<!doctype html><title>Lazy tab</title><body style="font-family:sans-serif;padding:2rem"><h1>Lazy tab</h1><p>This tab will load when selected.</p><p><a href="${escaped}">${escaped}</a></p></body>`;
  return `data:text/html;charset=utf-8,${encodeURIComponent(html)}`;
}

export class SessionCommands extends CommandGroup {
  readonly namespace = "session";
  readonly commands: Record<string, CommandEntry> = {
    "session.save":      (a: SessionSaveArgs) => this.sessionSave(a),
    "session.load":      { background: true, run: (a: SessionLoadArgs) => this.sessionLoad(a) },
    "session.list":      () => this.sessionList(),
    "session.remove":    (a: SessionRemoveArgs) => this.sessionRemove(a),
    "session.diff":      (a: SessionDiffArgs) => this.sessionDiff(a),
    "session.export":    (a: SessionSaveArgs) => this.sessionExport(a),
    "session.import":    (a: SessionImportArgs) => this.sessionImport(a),
    "session.auto_save": (a: SessionAutoSaveArgs) => this.autoSaveManager.setEnabled(Boolean(a.enabled)),
    "clients.list":           () => this.clientsList(),
    "clients.rename_profile": (a: ClientsRenameProfileArgs) => this.clientsRenameProfile(a),
  };

  private readonly autoSaveManager = new AutoSaveManager();

  private async sessionSave({ name }: SessionSaveArgs) {
    const { session, tabCount } = await captureCurrentSession();
    const sessions = await getSessions();
    sessions[name] = session;
    await api.storage.local.set({ sessions });
    return { name, tabs: tabCount };
  }

  // Public: invoked from index.ts on api.tabs.onActivated.
  async activateLazyTab(tabId: number | string) {
    const { lazySessionTabs } = await api.storage.local.get<{ lazySessionTabs?: LazySessionMap }>("lazySessionTabs");
    const entry = lazySessionTabs?.[tabId];
    if (!entry?.url) return false;
    delete lazySessionTabs[tabId];
    await api.storage.local.set({ lazySessionTabs });
    await api.tabs.update(Number(tabId), { url: entry.url });
    return true;
  }

  private async sessionLoad({ name, gentleMode, discardBackgroundTabs = false, lazy = false, eagerTabs = 10, __job }: SessionLoadArgs = {}) {
    return runLargeOperation("session.load", async () => {
      const sessions = await getSessions();
      const session = sessions[name];
      if (!session) throw new Error(`Session '${name}' not found`);

      const sessionTabs = getSessionTabs(session).sort((a, b) => Number(Boolean(b.pinned)) - Number(Boolean(a.pinned)));
      const createdTabs = [];
      const throttle = await getLargeOperationThrottle(sessionTabs.length, gentleMode);
      const createBatchSize = Math.max(1, Math.min(10, throttle.batchSize));
      const eagerLimit = lazy ? Math.max(0, Number(eagerTabs) || 0) : sessionTabs.length;
      const { lazySessionTabs } = await api.storage.local.get<{ lazySessionTabs?: LazySessionMap }>("lazySessionTabs");
      const lazyMap: LazySessionMap = lazySessionTabs || {};
      updateJobProgress(__job, { phase: "opening", current: 0, total: sessionTabs.length });

      for (const [idx, entry] of sessionTabs.entries()) {
        throwIfJobCancelled(__job);
        const shouldLazy = lazy && idx >= eagerLimit;
        const tab = await api.tabs.create({ url: shouldLazy ? lazyPlaceholderUrl(entry.url) : entry.url, active: false, pinned: Boolean(entry.pinned) });
        createdTabs.push({ tabId: tab.id, entry });
        if (shouldLazy) {
          lazyMap[String(tab.id)] = { url: entry.url, createdAt: Date.now() };
        } else if (discardBackgroundTabs && !entry.pinned && api.tabs.discard) {
          try { await api.tabs.discard(tab.id); } catch (_) {}
        }
        updateJobProgress(__job, { phase: shouldLazy ? "creating lazy placeholders" : "opening", current: createdTabs.length, total: sessionTabs.length });
        await yieldForLargeOperation(createdTabs.length, createBatchSize, Math.max(50, throttle.pauseMs));
      }
      if (lazy) await api.storage.local.set({ lazySessionTabs: lazyMap });

      const groups = new Map();
      for (const { tabId, entry } of createdTabs) {
        if (!entry.group) continue;
        const key = entry.group.key || `${entry.group.title || "group"}:${groups.size}`;
        if (!groups.has(key)) {
          groups.set(key, { meta: entry.group, tabIds: [] });
        }
        groups.get(key).tabIds.push(tabId);
      }

      let restoredGroups = 0;
      updateJobProgress(__job, { phase: "restoring groups", current: 0, total: groups.size });
      for (const { meta, tabIds } of groups.values()) {
        throwIfJobCancelled(__job);
        const restoredGroupId = await groupTabs({ tabIds });
        await updateTabGroup(restoredGroupId, {
          title: meta.title || "",
          color: normalizeGroupColor(meta.color),
          collapsed: Boolean(meta.collapsed),
        });
        restoredGroups++;
        updateJobProgress(__job, { phase: "restoring groups", current: restoredGroups, total: groups.size });
        await yieldForLargeOperation(restoredGroups, 5, Math.max(50, throttle.pauseMs));
      }

      return { name, tabs: sessionTabs.length, gentle: throttle.gentle, audible: throttle.audible, discarded: Boolean(discardBackgroundTabs), lazy: Boolean(lazy), eagerTabs: eagerLimit };
    });
  }

  private async sessionList() {
    const sessions = await getSessions();
    return Object.entries(sessions).map(([name, s]) => ({
      name,
      tabs: getSessionTabs(s).length,
      savedAt: s.savedAt || null,
    }));
  }

  private async sessionRemove({ name }: SessionRemoveArgs) {
    const sessions = await getSessions();
    if (!(name in sessions)) throw new Error(`Session '${name}' not found`);
    delete sessions[name];
    await api.storage.local.set({ sessions });
    return { name };
  }

  private async sessionDiff({ nameA, nameB }: SessionDiffArgs) {
    const sessions = await getSessions();
    const a = new Set(getSessionTabs(sessions[nameA]).map(tab => tab.url));
    const b = new Set(getSessionTabs(sessions[nameB]).map(tab => tab.url));
    return {
      added:   [...b].filter(u => !a.has(u)),
      removed: [...a].filter(u => !b.has(u)),
    };
  }

  private async sessionExport({ name }: SessionSaveArgs) {
    const sessions = await getSessions();
    if (!name) return { sessions };
    const session = sessions[name];
    if (!session) throw new Error(`Session '${name}' not found`);
    return { name, session };
  }

  private isSession(value: Json | undefined): boolean {
    if (!value || typeof value !== "object" || Array.isArray(value)) return false;
    const candidate = value as object as StoredSession;
    return Array.isArray(candidate.tabs) || Array.isArray(candidate.urls);
  }

  private async sessionImport({ name, session, overwrite = false }: SessionImportArgs) {
    if (!name) throw new Error("Session name is required");
    if (!this.isSession(session)) throw new Error("Session payload must contain tabs or urls");
    const sessions = await getSessions();
    if (!overwrite && sessions[name]) throw new Error(`Session '${name}' already exists`);
    const stored = session as object as StoredSession;
    sessions[name] = { ...stored, savedAt: Number(stored.savedAt) || Date.now() };
    await api.storage.local.set({ sessions });
    return { name, tabs: getSessionTabs(sessions[name]).length };
  }

  private async clientsList() {
    const manifest = api.runtime.getManifest();
    const alias = await getProfileAlias();
    const browserInfo = api.runtime.getBrowserInfo ? await api.runtime.getBrowserInfo() : null;
    const userAgent = navigator.userAgent;
    return [{
      name: browserInfo?.name || (userAgent.includes("Firefox/") ? "Firefox" : "Chrome"),
      version: browserInfo?.version || userAgent.match(/(?:Chrome|Firefox)\/([\d.]+)/)?.[1] || "unknown",
      platform: navigator.platform,
      extensionVersion: manifest.version,
      profile: alias,
    }];
  }

  private async clientsRenameProfile({ alias }: ClientsRenameProfileArgs) {
    await api.storage.local.set({ profileAlias: alias });
    return { alias };
  }
}