geo/apps/desktop-client/src/main/session-registry.ts

import { randomUUID } from "node:crypto";
import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { dirname, join } from "node:path";

import { app, session } from "electron/main";
import type { Session } from "electron/main";

import { STANDARD_ACCEPT_LANGUAGES, STANDARD_USER_AGENT } from "./user-agent";

export interface SessionHandle {
  accountId: string;
  partition: string;
  session: Session;
}

export interface SessionHandleSnapshot {
  accountId: string;
  partition: string;
}

const registry = new Map<string, SessionHandle>();
const sessionsWithUA = new WeakSet<Session>();
let persistedPartitionsCache: Record<string, string> | null = null;

function persistedPartitionsPath(): string {
  return join(app.getPath("userData"), "desktop-session-partitions.json");
}

function readPersistedPartitions(): Record<string, string> {
  if (persistedPartitionsCache) {
    return persistedPartitionsCache;
  }

  try {
    persistedPartitionsCache = JSON.parse(
      readFileSync(persistedPartitionsPath(), "utf8"),
    ) as Record<string, string>;
  } catch {
    persistedPartitionsCache = {};
  }

  return persistedPartitionsCache;
}

function writePersistedPartitions(next: Record<string, string>): void {
  persistedPartitionsCache = next;
  const target = persistedPartitionsPath();
  mkdirSync(dirname(target), { recursive: true });
  writeFileSync(target, JSON.stringify(next, null, 2), "utf8");
}

function persistedPartitionFor(accountId: string): string | null {
  return readPersistedPartitions()[accountId] ?? null;
}

export function getPersistedPartition(accountId: string): string | null {
  return persistedPartitionFor(accountId);
}

function rememberPersistedPartition(accountId: string, partition: string): void {
  const current = readPersistedPartitions();
  if (current[accountId] === partition) {
    return;
  }

  writePersistedPartitions({
    ...current,
    [accountId]: partition,
  });
}

function forgetPersistedPartition(accountId: string): string | null {
  const current = readPersistedPartitions();
  const partition = current[accountId] ?? null;
  if (!partition) {
    return null;
  }

  const next = { ...current };
  delete next[accountId];
  writePersistedPartitions(next);
  return partition;
}

function partitionFor(accountId: string): string {
  return `persist:acc-${accountId}`;
}

function pendingPartitionFor(seed: string): string {
  return `persist:pending-${seed}-${randomUUID()}`;
}

function applyStandardUserAgent(target: Session): Session {
  if (!sessionsWithUA.has(target)) {
    target.setUserAgent(STANDARD_USER_AGENT, STANDARD_ACCEPT_LANGUAGES);
    sessionsWithUA.add(target);
  }
  return target;
}

function prepareSession(target: Session): Session {
  applyStandardUserAgent(target);
  target.clearHostResolverCache().catch((error) => {
    console.warn("[desktop-session] clearHostResolverCache failed", error);
  });
  return target;
}

export async function initSessionRegistry(): Promise<void> {
  registry.clear();
}

export function createSessionHandle(accountId?: string): SessionHandle {
  const resolvedAccountID = accountId ?? randomUUID();
  const existing = registry.get(resolvedAccountID);
  if (existing) {
    return existing;
  }

  const partition = persistedPartitionFor(resolvedAccountID) ?? partitionFor(resolvedAccountID);
  const handle: SessionHandle = {
    accountId: resolvedAccountID,
    partition,
    session: prepareSession(session.fromPartition(partition)),
  };
  registry.set(resolvedAccountID, handle);
  return handle;
}

export function createSessionHandleForPartition(accountId: string, partition: string): SessionHandle {
  const existing = registry.get(accountId);
  if (existing && existing.partition === partition) {
    return existing;
  }

  const handle: SessionHandle = {
    accountId,
    partition,
    session: prepareSession(session.fromPartition(partition)),
  };
  registry.set(accountId, handle);
  rememberPersistedPartition(accountId, partition);
  return handle;
}

export function createPendingSessionHandle(seed = "bind"): SessionHandle {
  const accountId = `pending:${seed}:${randomUUID()}`;
  const partition = pendingPartitionFor(seed);
  const handle: SessionHandle = {
    accountId,
    partition,
    session: prepareSession(session.fromPartition(partition)),
  };
  registry.set(accountId, handle);
  return handle;
}

export function attachSessionHandle(accountId: string, handle: SessionHandle): SessionHandle {
  const existing = registry.get(accountId);
  if (existing && existing !== handle) {
    registry.delete(existing.accountId);
  }

  if (handle.accountId !== accountId) {
    registry.delete(handle.accountId);
  }

  const next: SessionHandle = {
    accountId,
    partition: handle.partition,
    session: handle.session,
  };
  registry.set(accountId, next);
  rememberPersistedPartition(accountId, handle.partition);
  return next;
}

export function forgetSessionHandle(accountId: string): void {
  registry.delete(accountId);
}

export async function clearSessionHandle(accountId: string): Promise<void> {
  const active = registry.get(accountId);
  const persistedPartition = forgetPersistedPartition(accountId);
  const partition = active?.partition ?? persistedPartition;
  registry.delete(accountId);

  if (!partition) {
    return;
  }

  const target = prepareSession(session.fromPartition(partition));

  try {
    await target.clearStorageData();
  } catch (error) {
    console.warn("[desktop-session] clearStorageData failed", { accountId, partition, error });
  }

  try {
    await target.clearCache();
  } catch (error) {
    console.warn("[desktop-session] clearCache failed", { accountId, partition, error });
  }

  try {
    await target.cookies.flushStore();
  } catch (error) {
    console.warn("[desktop-session] flushStore failed", { accountId, partition, error });
  }
}

export function getSessionHandle(accountId: string): SessionHandle | undefined {
  return registry.get(accountId);
}

export function listSessionHandles(): SessionHandle[] {
  return [...registry.values()];
}

export function listSessionHandleSnapshots(): SessionHandleSnapshot[] {
  return [...registry.values()].map(({ accountId, partition }) => ({
    accountId,
    partition,
  }));
}