codexHalo/electron/usage-reader.js

import { EventEmitter } from 'node:events'
import { closeSync, existsSync, fstatSync, openSync, readSync, readdirSync, statSync, watch } from 'node:fs'
import path from 'node:path'

export const SHORT_WINDOW_MINUTES = 300
export const WEEK_WINDOW_MINUTES = 10080

const MAX_FILES_TO_SCAN = 160
const TAIL_BYTES_PER_FILE = 2 * 1024 * 1024

function clamp(value, min, max) {
  return Math.min(max, Math.max(min, value))
}

export function parseUsageLine(line, sourceFile = '') {
  if (!line.includes('"rate_limits"') || !line.includes('"token_count"')) return []

  try {
    const root = JSON.parse(line)
    if (root?.type !== 'event_msg' || root?.payload?.type !== 'token_count') return []

    const observedAt = Number.isNaN(Date.parse(root.timestamp))
      ? Date.now()
      : Date.parse(root.timestamp)

    return ['primary', 'secondary'].flatMap((key) => {
      const window = root.payload.rate_limits?.[key]
      const windowMinutes = Number(window?.window_minutes)
      const usedPercent = Number(window?.used_percent)
      const resetsAtSeconds = Number(window?.resets_at)

      if (![SHORT_WINDOW_MINUTES, WEEK_WINDOW_MINUTES].includes(windowMinutes)) return []
      if (![usedPercent, resetsAtSeconds].every(Number.isFinite)) return []

      return [{
        windowMinutes,
        usedPercent: clamp(usedPercent, 0, 100),
        remainingPercent: clamp(100 - usedPercent, 0, 100),
        resetsAt: resetsAtSeconds * 1000,
        observedAt,
        sourceFile,
      }]
    })
  } catch {
    return []
  }
}

function collectRolloutFiles(rootPath) {
  const results = []
  const pending = [rootPath]

  while (pending.length > 0) {
    const directory = pending.pop()
    let entries

    try {
      entries = readdirSync(directory, { withFileTypes: true })
    } catch {
      continue
    }

    for (const entry of entries) {
      const fullPath = path.join(directory, entry.name)
      if (entry.isDirectory()) pending.push(fullPath)
      else if (entry.isFile() && /^rollout-.*\.jsonl$/i.test(entry.name)) results.push(fullPath)
    }
  }

  return results
    .map((filePath) => {
      try {
        return { filePath, modifiedAt: statSync(filePath).mtimeMs }
      } catch {
        return null
      }
    })
    .filter(Boolean)
    .sort((a, b) => b.modifiedAt - a.modifiedAt)
    .slice(0, MAX_FILES_TO_SCAN)
}

function readFileTail(filePath) {
  let descriptor

  try {
    descriptor = openSync(filePath, 'r')
    const size = fstatSync(descriptor).size
    const start = Math.max(0, size - TAIL_BYTES_PER_FILE)
    const length = size - start
    const buffer = Buffer.alloc(length)
    readSync(descriptor, buffer, 0, length, start)

    const lines = buffer.toString('utf8').split(/\r?\n/).filter(Boolean)
    if (start > 0) lines.shift()
    return lines
  } catch {
    return []
  } finally {
    if (descriptor !== undefined) closeSync(descriptor)
  }
}

export function readLatestUsage(sessionsPath) {
  if (!existsSync(sessionsPath)) {
    return {
      shortWindow: null,
      weekWindow: null,
      status: '未找到 Codex 会话目录',
      updatedAt: Date.now(),
    }
  }

  let shortWindow = null
  let weekWindow = null

  for (const { filePath, modifiedAt } of collectRolloutFiles(sessionsPath)) {
    for (const line of readFileTail(filePath)) {
      for (const snapshot of parseUsageLine(line, filePath)) {
        if (snapshot.windowMinutes === SHORT_WINDOW_MINUTES &&
            (!shortWindow || snapshot.observedAt > shortWindow.observedAt)) {
          shortWindow = snapshot
        }

        if (snapshot.windowMinutes === WEEK_WINDOW_MINUTES &&
            (!weekWindow || snapshot.observedAt > weekWindow.observedAt)) {
          weekWindow = snapshot
        }
      }
    }

    if (shortWindow && weekWindow &&
        modifiedAt < shortWindow.observedAt && modifiedAt < weekWindow.observedAt) break
  }

  const now = Date.now()
  if (shortWindow?.resetsAt <= now) shortWindow = null
  if (weekWindow?.resetsAt <= now) weekWindow = null

  return {
    shortWindow,
    weekWindow,
    status: shortWindow || weekWindow ? '已从本地日志同步' : '本地日志中暂无额度快照',
    updatedAt: Date.now(),
  }
}

export class UsageReader extends EventEmitter {
  constructor(sessionsPath) {
    super()
    this.sessionsPath = sessionsPath
    this.fileWatcher = null
    this.refreshTimer = null
    this.safetyTimer = null
  }

  readLatest() {
    return readLatestUsage(this.sessionsPath)
  }

  refresh() {
    const snapshot = this.readLatest()
    this.emit('snapshot', snapshot)
    return snapshot
  }

  requestRefresh() {
    clearTimeout(this.refreshTimer)
    this.refreshTimer = setTimeout(() => this.refresh(), 250)
  }

  start() {
    if (existsSync(this.sessionsPath)) {
      try {
        this.fileWatcher = watch(this.sessionsPath, { recursive: true }, (_event, fileName) => {
          if (!fileName || /(^|[\\/])rollout-.*\.jsonl$/i.test(fileName)) this.requestRefresh()
        })
        this.fileWatcher.on('error', () => this.requestRefresh())
      } catch {
        this.fileWatcher = null
      }
    }

    this.safetyTimer = setInterval(() => this.requestRefresh(), 15_000)
  }

  stop() {
    clearTimeout(this.refreshTimer)
    clearInterval(this.safetyTimer)
    this.fileWatcher?.close()
    this.fileWatcher = null
  }
}