zn-ai/electron/utils/token-usage.ts

import { readdir, readFile, stat } from 'fs/promises';
import { join } from 'path';
import logManager from '@electron/service/logger';
import { getUserDataDir } from './paths';
import {
  extractSessionIdFromTranscriptFileName,
  parseUsageEntriesFromJsonl,
  type TokenUsageHistoryEntry,
} from './token-usage-core';

export {
  extractSessionIdFromTranscriptFileName,
  parseUsageEntriesFromJsonl,
  type TokenUsageHistoryEntry,
} from './token-usage-core';

const TRANSCRIPT_ROOT_DIR_NAMES = ['models', 'agents'] as const;

async function listAgentIdsWithSessionDirs(rootDirName: string): Promise<string[]> {
  const rootDir = join(getUserDataDir(), rootDirName);
  const agentIds = new Set<string>();

  try {
    const agentEntries = await readdir(rootDir, { withFileTypes: true });
    for (const entry of agentEntries) {
      if (entry.isDirectory()) {
        const normalized = entry.name.trim();
        if (normalized) {
          agentIds.add(normalized);
        }
      }
    }
  } catch {
    // Ignore disk discovery failures and return whatever we already found.
  }

  return [...agentIds];
}

async function listRecentSessionFiles(): Promise<Array<{ filePath: string; sessionId: string; agentId: string; mtimeMs: number }>> {
  const filesBySession = new Map<string, { filePath: string; sessionId: string; agentId: string; mtimeMs: number }>();

  try {
    for (const rootDirName of TRANSCRIPT_ROOT_DIR_NAMES) {
      const agentEntries = await listAgentIdsWithSessionDirs(rootDirName);

      for (const agentId of agentEntries) {
        const sessionsDir = join(getUserDataDir(), rootDirName, agentId, 'sessions');
        try {
          const sessionEntries = await readdir(sessionsDir);

          for (const fileName of sessionEntries) {
            const sessionId = extractSessionIdFromTranscriptFileName(fileName);
            if (!sessionId) continue;
            const filePath = join(sessionsDir, fileName);
            try {
              const fileStat = await stat(filePath);
              const sessionKey = `${agentId}:${sessionId}`;
              const existing = filesBySession.get(sessionKey);

              if (!existing || fileStat.mtimeMs > existing.mtimeMs) {
                filesBySession.set(sessionKey, {
                  filePath,
                  sessionId,
                  agentId,
                  mtimeMs: fileStat.mtimeMs,
                });
              }
            } catch {
              continue;
            }
          }
        } catch {
          continue;
        }
      }
    }

    const files = [...filesBySession.values()];
    files.sort((a, b) => b.mtimeMs - a.mtimeMs);
    return files;
  } catch {
    return [];
  }
}

export async function getRecentTokenUsageHistory(limit?: number): Promise<TokenUsageHistoryEntry[]> {
  const files = await listRecentSessionFiles();
  const results: TokenUsageHistoryEntry[] = [];
  const maxEntries = typeof limit === 'number' && Number.isFinite(limit)
    ? Math.max(Math.floor(limit), 0)
    : Number.POSITIVE_INFINITY;

  for (const file of files) {
    if (results.length >= maxEntries) break;
    try {
      const content = await readFile(file.filePath, 'utf8');
      const entries = parseUsageEntriesFromJsonl(content, {
        sessionId: file.sessionId,
        agentId: file.agentId,
      }, Number.isFinite(maxEntries) ? maxEntries - results.length : undefined);
      results.push(...entries);
    } catch (error) {
      logManager.error(`Failed to read token usage transcript ${file.filePath}:`, error);
    }
  }

  results.sort((a, b) => Date.parse(b.timestamp) - Date.parse(a.timestamp));
  return Number.isFinite(maxEntries) ? results.slice(0, maxEntries) : results;
}