geo/apps/desktop-client/src/main/storage-cleaner.ts

import { readdir, rm, stat } from 'node:fs/promises'
import { join } from 'node:path'

import type { Session } from 'electron/main'
import { app, session as electronSession } from 'electron/main'

import {
  getDesktopAppSettings,
  getLastStorageCleanupAt,
  markStorageCleanupRun,
} from './app-settings'
import {
  forgetPersistedSessionPartition,
  listPersistedSessionPartitions,
  listSessionHandles,
} from './session-registry'

export type DesktopStorageEntryKind = 'app-cache' | 'bound-account' | 'unbound-session'

export interface DesktopStorageEntry {
  id: string
  label: string
  kind: DesktopStorageEntryKind
  partition: string | null
  accountId: string | null
  sizeBytes: number
  cleanableBytes: number
  lastModifiedAt: number | null
  cleanable: boolean
  active: boolean
}

export interface DesktopStorageSnapshot {
  generatedAt: number
  totalBytes: number
  cleanableBytes: number
  appCacheBytes: number
  boundAccountBytes: number
  boundAccountCleanableBytes: number
  unboundSessionBytes: number
  unboundSessionCount: number
  boundAccountCount: number
  entries: DesktopStorageEntry[]
  autoCleanStorageWeekly: boolean
  lastCleanupAt: number | null
  nextCleanupAt: number | null
}

export interface DesktopStorageCleanupResult {
  cleanedAt: number
  cleanedBytes: number
  removedEntries: number
  before: DesktopStorageSnapshot
  after: DesktopStorageSnapshot
}

interface DirectoryFootprint {
  bytes: number
  lastModifiedAt: number | null
}

interface PartitionDirectory {
  name: string
  partition: string
  path: string
}

const weeklyCleanupIntervalMs = 7 * 24 * 60 * 60 * 1000
const cleanupSchedulerIntervalMs = 6 * 60 * 60 * 1000
const cleanableCacheRelativePaths = [
  'Cache',
  'Code Cache',
  'GPUCache',
  'DawnGraphiteCache',
  'DawnWebGPUCache',
  'Shared Dictionary/cache',
]

let storageCleanupTimer: ReturnType<typeof setInterval> | null = null
let cleanupInFlight: Promise<DesktopStorageCleanupResult> | null = null

function userDataPath(): string {
  return app.getPath('userData')
}

function partitionsRootPath(): string {
  return join(userDataPath(), 'Partitions')
}

function partitionDirectoryName(partition: string): string | null {
  if (!partition.startsWith('persist:')) {
    return null
  }
  const name = partition.slice('persist:'.length)
  return name ? name : null
}

function accountIdFromPartition(partition: string): string | null {
  const name = partitionDirectoryName(partition)
  if (!name?.startsWith('acc-')) {
    return null
  }
  const accountId = name.slice('acc-'.length)
  return accountId ? accountId : null
}

function partitionPath(partition: string): string | null {
  const name = partitionDirectoryName(partition)
  return name ? join(partitionsRootPath(), name) : null
}

function maxModifiedAt(left: number | null, right: number | null): number | null {
  if (left === null) {
    return right
  }
  if (right === null) {
    return left
  }
  return Math.max(left, right)
}

async function directoryFootprint(target: string): Promise<DirectoryFootprint> {
  let info
  try {
    info = await stat(target)
  } catch {
    return { bytes: 0, lastModifiedAt: null }
  }

  if (!info.isDirectory()) {
    return {
      bytes: info.size,
      lastModifiedAt: info.mtimeMs,
    }
  }

  let bytes = 0
  let lastModifiedAt: number | null = info.mtimeMs
  let entries
  try {
    entries = await readdir(target, { withFileTypes: true })
  } catch {
    return { bytes, lastModifiedAt }
  }

  for (const entry of entries) {
    const child = join(target, entry.name)
    if (entry.isDirectory()) {
      const nested = await directoryFootprint(child)
      bytes += nested.bytes
      lastModifiedAt = maxModifiedAt(lastModifiedAt, nested.lastModifiedAt)
      continue
    }

    try {
      const childInfo = await stat(child)
      bytes += childInfo.size
      lastModifiedAt = maxModifiedAt(lastModifiedAt, childInfo.mtimeMs)
    } catch {
      // Ignore files that disappear while Chromium is writing cache entries.
    }
  }

  return { bytes, lastModifiedAt }
}

async function sumDirectoryFootprints(paths: string[]): Promise<DirectoryFootprint> {
  let bytes = 0
  let lastModifiedAt: number | null = null
  for (const target of paths) {
    const footprint = await directoryFootprint(target)
    bytes += footprint.bytes
    lastModifiedAt = maxModifiedAt(lastModifiedAt, footprint.lastModifiedAt)
  }
  return { bytes, lastModifiedAt }
}

function cleanableCachePaths(root: string): string[] {
  return cleanableCacheRelativePaths.map((relativePath) => join(root, relativePath))
}

async function listPartitionDirectories(): Promise<PartitionDirectory[]> {
  let entries
  try {
    entries = await readdir(partitionsRootPath(), { withFileTypes: true })
  } catch {
    return []
  }

  return entries
    .filter((entry) => entry.isDirectory())
    .map((entry) => ({
      name: entry.name,
      partition: `persist:${entry.name}`,
      path: join(partitionsRootPath(), entry.name),
    }))
}

function activePartitionSet(): Set<string> {
  return new Set(listSessionHandles().map((handle) => handle.partition))
}

function persistedPartitionMap(): Map<string, string> {
  const result = new Map<string, string>()
  for (const item of listPersistedSessionPartitions()) {
    result.set(item.partition, item.accountId)
  }
  return result
}

function nextCleanupAt(lastCleanupAt: number | null, enabled: boolean): number | null {
  if (!enabled) {
    return null
  }
  return (lastCleanupAt ?? Date.now()) + weeklyCleanupIntervalMs
}

export async function getDesktopStorageSnapshot(): Promise<DesktopStorageSnapshot> {
  const activePartitions = activePartitionSet()
  const persistedPartitions = persistedPartitionMap()
  const partitionDirs = await listPartitionDirectories()
  const entries: DesktopStorageEntry[] = []

  const appCache = await sumDirectoryFootprints(cleanableCachePaths(userDataPath()))
  entries.push({
    id: 'app-cache',
    label: '应用临时缓存',
    kind: 'app-cache',
    partition: null,
    accountId: null,
    sizeBytes: appCache.bytes,
    cleanableBytes: appCache.bytes,
    lastModifiedAt: appCache.lastModifiedAt,
    cleanable: appCache.bytes > 0,
    active: true,
  })

  const boundPartitions = new Set<string>()
  for (const dir of partitionDirs) {
    const accountId =
      persistedPartitions.get(dir.partition) ?? accountIdFromPartition(dir.partition)
    const isBound = Boolean(accountId)
    if (isBound) {
      boundPartitions.add(dir.partition)
      continue
    }

    const active = activePartitions.has(dir.partition)
    if (active) {
      continue
    }

    const size = await directoryFootprint(dir.path)
    entries.push({
      id: dir.partition,
      label: '未绑定缓存',
      kind: 'unbound-session',
      partition: dir.partition,
      accountId: null,
      sizeBytes: size.bytes,
      cleanableBytes: size.bytes,
      lastModifiedAt: size.lastModifiedAt,
      cleanable: size.bytes > 0,
      active: false,
    })
  }

  for (const [partition, accountId] of persistedPartitions) {
    if (entries.some((entry) => entry.partition === partition)) {
      continue
    }
    if (accountId) {
      boundPartitions.add(partition)
    }
  }

  const totalBytes = entries.reduce((total, entry) => total + entry.sizeBytes, 0)
  const cleanableBytes = entries.reduce((total, entry) => total + entry.cleanableBytes, 0)
  const appCacheBytes = entries
    .filter((entry) => entry.kind === 'app-cache')
    .reduce((total, entry) => total + entry.sizeBytes, 0)
  const boundAccountBytes = 0
  const boundAccountCleanableBytes = 0
  const unboundSessionEntries = entries.filter((entry) => entry.kind === 'unbound-session')
  const unboundSessionBytes = unboundSessionEntries.reduce(
    (total, entry) => total + entry.sizeBytes,
    0,
  )
  const lastCleanupAt = getLastStorageCleanupAt()
  const autoCleanStorageWeekly = getDesktopAppSettings().autoCleanStorageWeekly

  return {
    generatedAt: Date.now(),
    totalBytes,
    cleanableBytes,
    appCacheBytes,
    boundAccountBytes,
    boundAccountCleanableBytes,
    unboundSessionBytes,
    unboundSessionCount: unboundSessionEntries.length,
    boundAccountCount: boundPartitions.size,
    entries: entries.sort((left, right) => right.sizeBytes - left.sizeBytes),
    autoCleanStorageWeekly,
    lastCleanupAt,
    nextCleanupAt: nextCleanupAt(lastCleanupAt, autoCleanStorageWeekly),
  }
}

async function clearBrowserCache(target: Session): Promise<void> {
  await target.clearCache().catch((error) => {
    console.warn('[desktop-storage] clearCache failed', error)
  })
  await target.clearCodeCaches({ urls: [] }).catch((error) => {
    console.warn('[desktop-storage] clearCodeCaches failed', error)
  })
  await target.clearSharedDictionaryCache().catch((error) => {
    console.warn('[desktop-storage] clearSharedDictionaryCache failed', error)
  })
  await target.clearHostResolverCache().catch((error) => {
    console.warn('[desktop-storage] clearHostResolverCache failed', error)
  })
}

async function removeCleanableCacheDirectories(root: string): Promise<void> {
  for (const target of cleanableCachePaths(root)) {
    await rm(target, { recursive: true, force: true }).catch((error) => {
      console.warn('[desktop-storage] remove cache directory failed', { target, error })
    })
  }
}

async function removeUnboundPartition(partition: string): Promise<boolean> {
  const root = partitionPath(partition)
  if (!root) {
    return false
  }

  const target = electronSession.fromPartition(partition)
  await target.closeAllConnections().catch((error) => {
    console.warn('[desktop-storage] close unbound session connections failed', { partition, error })
  })
  await target.clearData().catch((error) => {
    console.warn('[desktop-storage] clear unbound session data failed', { partition, error })
  })
  await target.clearStorageData().catch((error) => {
    console.warn('[desktop-storage] clear unbound storage data failed', { partition, error })
  })
  await clearBrowserCache(target)
  await target.cookies.flushStore().catch((error) => {
    console.warn('[desktop-storage] flush unbound session cookies failed', { partition, error })
  })
  await rm(root, { recursive: true, force: true })
  return true
}

async function cleanDesktopStorageOnce(): Promise<DesktopStorageCleanupResult> {
  const before = await getDesktopStorageSnapshot()
  const activePartitions = activePartitionSet()
  const persistedPartitions = persistedPartitionMap()

  await clearBrowserCache(electronSession.defaultSession)
  await removeCleanableCacheDirectories(userDataPath())

  let removedEntries = 0
  for (const entry of before.entries) {
    if (
      !entry.partition ||
      entry.kind !== 'unbound-session' ||
      activePartitions.has(entry.partition)
    ) {
      continue
    }

    try {
      if (await removeUnboundPartition(entry.partition)) {
        removedEntries += 1
      }
    } catch (error) {
      console.warn('[desktop-storage] remove unbound partition failed', {
        partition: entry.partition,
        error,
      })
    }
  }

  for (const [partition, accountId] of persistedPartitions) {
    const root = partitionPath(partition)
    const footprint = root ? await directoryFootprint(root) : { bytes: 0, lastModifiedAt: null }
    if (footprint.bytes > 0 || footprint.lastModifiedAt) {
      continue
    }
    forgetPersistedSessionPartition(accountId, partition)
  }

  markStorageCleanupRun()
  const after = await getDesktopStorageSnapshot()

  return {
    cleanedAt: after.generatedAt,
    cleanedBytes: Math.max(0, before.totalBytes - after.totalBytes),
    removedEntries,
    before,
    after,
  }
}

export function cleanDesktopStorage(): Promise<DesktopStorageCleanupResult> {
  if (!cleanupInFlight) {
    cleanupInFlight = cleanDesktopStorageOnce().finally(() => {
      cleanupInFlight = null
    })
  }
  return cleanupInFlight
}

export async function runDueStorageCleanup(): Promise<DesktopStorageCleanupResult | null> {
  if (!getDesktopAppSettings().autoCleanStorageWeekly) {
    return null
  }

  const lastCleanupAt = getLastStorageCleanupAt()
  if (lastCleanupAt && Date.now() - lastCleanupAt < weeklyCleanupIntervalMs) {
    return null
  }

  return cleanDesktopStorage()
}

export function startStorageCleanupScheduler(): void {
  if (storageCleanupTimer) {
    clearInterval(storageCleanupTimer)
  }

  storageCleanupTimer = setInterval(() => {
    void runDueStorageCleanup().catch((error) => {
      console.warn('[desktop-storage] scheduled cleanup failed', error)
    })
  }, cleanupSchedulerIntervalMs)

  void runDueStorageCleanup().catch((error) => {
    console.warn('[desktop-storage] startup cleanup failed', error)
  })
}