@tmustier/pi-usage-extension 0.3.2 → 0.5.0

package/index.ts

@@ -5,89 +5,17 @@
  * - Tab cycles: Today → This Week → Last Week → All Time
  * - Arrow keys navigate providers
  * - Enter expands/collapses to show models
+ *
+ * Data collection and caching live in ./data.ts.
  */
 
 import type { ExtensionAPI, ExtensionCommandContext, Theme } from "@earendil-works/pi-coding-agent";
 import { DynamicBorder } from "@earendil-works/pi-coding-agent";
 import { CancellableLoader, Container, Spacer, matchesKey, visibleWidth, truncateToWidth, wrapTextWithAnsi } from "@earendil-works/pi-tui";
-import { readdir, readFile } from "node:fs/promises";
-import { join } from "node:path";
-import { homedir } from "node:os";
-
-// =============================================================================
-// Types
-// =============================================================================
-
-interface TokenStats {
-	total: number;
-	input: number;
-	output: number;
-	cacheRead: number;
-	cacheWrite: number;
-}
-
-interface BaseStats {
-	messages: number;
-	cost: number;
-	tokens: TokenStats;
-}
-
-interface ModelStats extends BaseStats {
-	sessions: Set<string>;
-}
 
-interface ProviderStats extends BaseStats {
-	sessions: Set<string>;
-	models: Map<string, ModelStats>;
-}
-
-interface TotalStats extends BaseStats {
-	sessions: number;
-}
-
-interface Insight {
-	percent: number; // 0-100
-	headline: string;
-	advice: string;
-}
-
-interface PeriodInsights {
-	insights: Insight[];
-}
-
-interface RawMessage {
-	sessionId: string;
-	timestamp: number;
-	cost: number;
-	input: number;
-	cacheRead: number;
-	cacheWrite: number;
-}
+import { collectUsageData, TAB_ORDER } from "./data";
+import type { BaseStats, TabName, UsageData } from "./data";
 
-interface PeriodRawData {
-	messages: RawMessage[];
-	sessionCosts: Map<string, number>;
-}
-
-interface GlobalSessionSpan {
-	startMs: number;
-	endMs: number;
-}
-
-interface TimeFilteredStats {
-	providers: Map<string, ProviderStats>;
-	totals: TotalStats;
-	insights: PeriodInsights;
-}
-
-interface UsageData {
-	today: TimeFilteredStats;
-	thisWeek: TimeFilteredStats;
-	lastWeek: TimeFilteredStats;
-	allTime: TimeFilteredStats;
-}
-
-type TabName = "today" | "thisWeek" | "lastWeek" | "allTime";
 type ViewMode = "table" | "insights";
 
 // =============================================================================
@@ -182,482 +110,6 @@ const TABLE_LAYOUTS: TableLayoutCandidate[] = [
 	{ columns: [COST_COLUMN], minNameWidth: 8, compact: true },
 ];
 
-// =============================================================================
-// Data Collection
-// =============================================================================
-
-function getSessionsDir(): string {
-	// Replicate Pi's logic: respect PI_CODING_AGENT_DIR env var
-	const agentDir = process.env.PI_CODING_AGENT_DIR || join(homedir(), ".pi", "agent");
-	return join(agentDir, "sessions");
-}
-
-async function collectSessionFilesRecursively(dir: string, files: string[], signal?: AbortSignal): Promise<void> {
-	try {
-		const entries = await readdir(dir, { withFileTypes: true });
-		for (const entry of entries) {
-			if (signal?.aborted) return;
-			const entryPath = join(dir, entry.name);
-			if (entry.isDirectory()) {
-				await collectSessionFilesRecursively(entryPath, files, signal);
-			} else if (entry.isFile() && entry.name.endsWith(".jsonl")) {
-				files.push(entryPath);
-			}
-		}
-	} catch {
-		// Skip directories we can't read
-	}
-}
-
-async function getAllSessionFiles(signal?: AbortSignal): Promise<string[]> {
-	const files: string[] = [];
-	await collectSessionFilesRecursively(getSessionsDir(), files, signal);
-	files.sort();
-	return files;
-}
-
-interface SessionMessage {
-	provider: string;
-	model: string;
-	cost: number;
-	input: number;
-	output: number;
-	cacheRead: number;
-	cacheWrite: number;
-	timestamp: number;
-}
-
-interface ParsedSessionFile {
-	sessionId: string;
-	messages: SessionMessage[];
-}
-
-async function parseSessionFile(
-	filePath: string,
-	seenHashes: Set<string>,
-	signal?: AbortSignal
-): Promise<ParsedSessionFile | null> {
-	try {
-		const content = await readFile(filePath, "utf8");
-		if (signal?.aborted) return null;
-		const lines = content.trim().split("\n");
-		const messages: SessionMessage[] = [];
-		let sessionId = "";
-
-		for (let i = 0; i < lines.length; i++) {
-			if (signal?.aborted) return null;
-			if (i % 500 === 0) {
-				await new Promise<void>((resolve) => setImmediate(resolve));
-			}
-			const line = lines[i]!;
-			if (!line.trim()) continue;
-			try {
-				const entry = JSON.parse(line);
-
-				if (entry.type === "session") {
-					sessionId = entry.id;
-				} else if (entry.type === "message" && entry.message?.role === "assistant") {
-					const msg = entry.message;
-					if (msg.usage && msg.provider && msg.model) {
-						const input = msg.usage.input || 0;
-						const output = msg.usage.output || 0;
-						const cacheRead = msg.usage.cacheRead || 0;
-						const cacheWrite = msg.usage.cacheWrite || 0;
-						const fallbackTs = entry.timestamp ? new Date(entry.timestamp).getTime() : 0;
-						const timestamp = msg.timestamp || (Number.isNaN(fallbackTs) ? 0 : fallbackTs);
-
-						// Deduplicate copied history across branched session files.
-						// Keep the existing ccusage-style hash so current totals remain comparable.
-						const totalTokens = input + output + cacheRead + cacheWrite;
-						const hash = `${timestamp}:${totalTokens}`;
-						if (seenHashes.has(hash)) continue;
-						seenHashes.add(hash);
-
-						messages.push({
-							provider: msg.provider,
-							model: msg.model,
-							cost: msg.usage.cost?.total || 0,
-							input,
-							output,
-							cacheRead,
-							cacheWrite,
-							timestamp,
-						});
-					}
-				}
-			} catch {
-				// Skip malformed lines
-			}
-		}
-
-		return sessionId ? { sessionId, messages } : null;
-	} catch {
-		return null;
-	}
-}
-
-// Helper to accumulate stats into a target
-function accumulateStats(
-	target: BaseStats,
-	cost: number,
-	tokens: { total: number; input: number; output: number; cacheRead: number; cacheWrite: number }
-): void {
-	target.messages++;
-	target.cost += cost;
-	target.tokens.total += tokens.total;
-	target.tokens.input += tokens.input;
-	target.tokens.output += tokens.output;
-	target.tokens.cacheRead += tokens.cacheRead;
-	target.tokens.cacheWrite += tokens.cacheWrite;
-}
-
-function emptyTokens(): TokenStats {
-	return { total: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
-}
-
-function emptyModelStats(): ModelStats {
-	return { sessions: new Set(), messages: 0, cost: 0, tokens: emptyTokens() };
-}
-
-function emptyProviderStats(): ProviderStats {
-	return { sessions: new Set(), messages: 0, cost: 0, tokens: emptyTokens(), models: new Map() };
-}
-
-function emptyTimeFilteredStats(): TimeFilteredStats {
-	return {
-		providers: new Map(),
-		totals: { sessions: 0, messages: 0, cost: 0, tokens: emptyTokens() },
-		insights: { insights: [] },
-	};
-}
-
-function emptyPeriodRawData(): PeriodRawData {
-	return { messages: [], sessionCosts: new Map() };
-}
-
-function emptyUsageData(): UsageData {
-	return {
-		today: emptyTimeFilteredStats(),
-		thisWeek: emptyTimeFilteredStats(),
-		lastWeek: emptyTimeFilteredStats(),
-		allTime: emptyTimeFilteredStats(),
-	};
-}
-
-function getPeriodsForTimestamp(timestamp: number, todayMs: number, weekStartMs: number, lastWeekStartMs: number): TabName[] {
-	const periods: TabName[] = ["allTime"];
-	if (timestamp >= todayMs) periods.push("today");
-	if (timestamp >= weekStartMs) {
-		periods.push("thisWeek");
-	} else if (timestamp >= lastWeekStartMs) {
-		periods.push("lastWeek");
-	}
-	return periods;
-}
-
-function addMessagesToUsageData(
-	data: UsageData,
-	sessionId: string,
-	messages: SessionMessage[],
-	todayMs: number,
-	weekStartMs: number,
-	lastWeekStartMs: number,
-	rawByPeriod: Record<TabName, PeriodRawData>,
-	globalSessionSpans: Map<string, GlobalSessionSpan>
-): void {
-	const sessionContributed = { today: false, thisWeek: false, lastWeek: false, allTime: false };
-
-	for (const msg of messages) {
-		// Track real per-session lifetime across every message we see, regardless of
-		// which period the message falls into. Used later for the "8h+ session" insight.
-		if (msg.timestamp > 0) {
-			const span = globalSessionSpans.get(sessionId);
-			if (!span) {
-				globalSessionSpans.set(sessionId, { startMs: msg.timestamp, endMs: msg.timestamp });
-			} else {
-				if (msg.timestamp < span.startMs) span.startMs = msg.timestamp;
-				if (msg.timestamp > span.endMs) span.endMs = msg.timestamp;
-			}
-		}
-
-		const periods = getPeriodsForTimestamp(msg.timestamp, todayMs, weekStartMs, lastWeekStartMs);
-		const tokens = {
-			// Count fresh tokens processed this turn.
-			// Include cacheWrite because those prompt tokens were newly written and billed.
-			// Exclude cacheRead because repeated cache hits would otherwise dominate totals.
-			total: msg.input + msg.output + msg.cacheWrite,
-			input: msg.input,
-			output: msg.output,
-			cacheRead: msg.cacheRead,
-			cacheWrite: msg.cacheWrite,
-		};
-
-		for (const period of periods) {
-			const stats = data[period];
-
-			let providerStats = stats.providers.get(msg.provider);
-			if (!providerStats) {
-				providerStats = emptyProviderStats();
-				stats.providers.set(msg.provider, providerStats);
-			}
-
-			let modelStats = providerStats.models.get(msg.model);
-			if (!modelStats) {
-				modelStats = emptyModelStats();
-				providerStats.models.set(msg.model, modelStats);
-			}
-
-			modelStats.sessions.add(sessionId);
-			accumulateStats(modelStats, msg.cost, tokens);
-
-			providerStats.sessions.add(sessionId);
-			accumulateStats(providerStats, msg.cost, tokens);
-
-			accumulateStats(stats.totals, msg.cost, tokens);
-			sessionContributed[period] = true;
-
-			const raw = rawByPeriod[period];
-			raw.messages.push({
-				sessionId,
-				timestamp: msg.timestamp,
-				cost: msg.cost,
-				input: msg.input,
-				cacheRead: msg.cacheRead,
-				cacheWrite: msg.cacheWrite,
-			});
-			raw.sessionCosts.set(sessionId, (raw.sessionCosts.get(sessionId) ?? 0) + msg.cost);
-		}
-	}
-
-	if (sessionContributed.today) data.today.totals.sessions++;
-	if (sessionContributed.thisWeek) data.thisWeek.totals.sessions++;
-	if (sessionContributed.lastWeek) data.lastWeek.totals.sessions++;
-	if (sessionContributed.allTime) data.allTime.totals.sessions++;
-}
-
-async function collectUsageData(signal?: AbortSignal): Promise<UsageData | null> {
-	const startOfToday = new Date();
-	startOfToday.setHours(0, 0, 0, 0);
-	const todayMs = startOfToday.getTime();
-
-	// Start of current week (Monday 00:00)
-	const startOfWeek = new Date();
-	const dayOfWeek = startOfWeek.getDay(); // 0 = Sunday, 1 = Monday, ...
-	const daysSinceMonday = dayOfWeek === 0 ? 6 : dayOfWeek - 1;
-	startOfWeek.setDate(startOfWeek.getDate() - daysSinceMonday);
-	startOfWeek.setHours(0, 0, 0, 0);
-	const weekStartMs = startOfWeek.getTime();
-
-	// Start of last week (previous Monday 00:00)
-	const startOfLastWeek = new Date(startOfWeek);
-	startOfLastWeek.setDate(startOfLastWeek.getDate() - 7);
-	const lastWeekStartMs = startOfLastWeek.getTime();
-
-	const data = emptyUsageData();
-	const rawByPeriod: Record<TabName, PeriodRawData> = {
-		today: emptyPeriodRawData(),
-		thisWeek: emptyPeriodRawData(),
-		lastWeek: emptyPeriodRawData(),
-		allTime: emptyPeriodRawData(),
-	};
-	const globalSessionSpans = new Map<string, GlobalSessionSpan>();
-
-	const sessionFiles = await getAllSessionFiles(signal);
-	if (signal?.aborted) return null;
-	const seenHashes = new Set<string>();
-
-	for (const filePath of sessionFiles) {
-		if (signal?.aborted) return null;
-		const parsed = await parseSessionFile(filePath, seenHashes, signal);
-		if (signal?.aborted) return null;
-		if (!parsed) continue;
-
-		addMessagesToUsageData(
-			data,
-			parsed.sessionId,
-			parsed.messages,
-			todayMs,
-			weekStartMs,
-			lastWeekStartMs,
-			rawByPeriod,
-			globalSessionSpans
-		);
-
-		await new Promise<void>((resolve) => setImmediate(resolve));
-	}
-
-	// Classify sessions that are globally long-running once, then reuse across periods.
-	const longSessionIds = new Set<string>();
-	for (const [id, span] of globalSessionSpans) {
-		if (span.endMs - span.startMs >= LONG_SESSION_MS) longSessionIds.add(id);
-	}
-
-	for (const period of TAB_ORDER) {
-		data[period].insights = computeInsights(rawByPeriod[period], longSessionIds);
-	}
-
-	return data;
-}
-
-// =============================================================================
-// Insights
-// =============================================================================
-
-const PARALLEL_WINDOW_MS = 2 * 60_000; // exact ±N milliseconds around each message
-const PARALLEL_SESSION_THRESHOLD = 4;
-const LARGE_CONTEXT_THRESHOLD = 150_000;
-const LARGE_CACHE_MISS_THRESHOLD = 100_000;
-const LONG_SESSION_MS = 8 * 60 * 60 * 1000;
-const TOP_SESSION_COUNT = 5;
-const MIN_MESSAGES_FOR_PARALLEL_INSIGHT = 10;
-const MIN_PERCENT_TO_SHOW = 1;
-
-/**
- * Insights are weighted by recorded API cost. Periods with zero total cost produce
- * an empty `insights` list — the UI renders a distinct empty-state for that case.
- * Long-running-session classification is passed in from a global pass so that a
- * session's real lifetime is used rather than the slice visible inside this period.
- */
-function computeInsights(raw: PeriodRawData, longSessionIds: Set<string>): PeriodInsights {
-	if (raw.messages.length === 0) {
-		return { insights: [] };
-	}
-
-	const total = raw.messages.reduce((sum, m) => sum + m.cost, 0);
-	if (total <= 0) {
-		return { insights: [] };
-	}
-
-	const candidates: Insight[] = [];
-
-	// 1. Parallel sessions — ≥ N unique sessions active within an exact ±W ms window.
-	const parallelWeight = computeParallelCostWeight(raw.messages);
-	if (parallelWeight !== null) {
-		candidates.push({
-			percent: (parallelWeight / total) * 100,
-			headline: `of your cost was while ${PARALLEL_SESSION_THRESHOLD}+ sessions ran in parallel`,
-			advice:
-				"All sessions share one rate limit. If you don't need them all at once, queueing uses capacity more evenly.",
-		});
-	}
-
-	// 2. Large context — input + cacheRead + cacheWrite > threshold.
-	const largeContextWeight = raw.messages
-		.filter((m) => m.input + m.cacheRead + m.cacheWrite > LARGE_CONTEXT_THRESHOLD)
-		.reduce((sum, m) => sum + m.cost, 0);
-	candidates.push({
-		percent: (largeContextWeight / total) * 100,
-		headline: `of your cost was at >${formatThresholdTokens(LARGE_CONTEXT_THRESHOLD)} context`,
-		advice:
-			"Longer sessions are more expensive even when cached. /compact mid-task, /clear when switching to new tasks.",
-	});
-
-	// 3. Large uncached prompt — fresh (non-cached) input > threshold, per the v0.2.0 formula.
-	const uncachedWeight = raw.messages
-		.filter((m) => m.input + m.cacheWrite > LARGE_CACHE_MISS_THRESHOLD)
-		.reduce((sum, m) => sum + m.cost, 0);
-	candidates.push({
-		percent: (uncachedWeight / total) * 100,
-		headline: `of your cost came from >${formatThresholdTokens(LARGE_CACHE_MISS_THRESHOLD)}-token uncached prompts`,
-		advice:
-			"Uncached input is expensive, and often happens when sending a message to a session that has gone idle. /compact before stepping away keeps the cold-start small.",
-	});
-
-	// 4. Long-running sessions — classification comes from the global pass so we use
-	//    true session lifetime, not just the span visible inside this period slice.
-	const longWeight = raw.messages
-		.filter((m) => longSessionIds.has(m.sessionId))
-		.reduce((sum, m) => sum + m.cost, 0);
-	if (longWeight > 0) {
-		candidates.push({
-			percent: (longWeight / total) * 100,
-			headline: `of your cost came from sessions active for ${LONG_SESSION_MS / 3_600_000}+ hours`,
-			advice:
-				"These are often background/loop sessions. Continuous usage can add up quickly so make sure it is intentional.",
-		});
-	}
-
-	// 5. Top-N session concentration.
-	if (raw.sessionCosts.size > TOP_SESSION_COUNT) {
-		const sortedSessions = Array.from(raw.sessionCosts.values()).sort((a, b) => b - a);
-		const topN = Math.min(TOP_SESSION_COUNT, sortedSessions.length);
-		const topWeight = sortedSessions.slice(0, topN).reduce((sum, c) => sum + c, 0);
-		candidates.push({
-			percent: (topWeight / total) * 100,
-			headline: `of your cost came from your top ${topN} session${topN === 1 ? "" : "s"}`,
-			advice:
-				"A small number of sessions drives most of your spend. The table view can help pinpoint which ones.",
-		});
-	}
-
-	const insights = candidates.filter((i) => i.percent >= MIN_PERCENT_TO_SHOW).sort((a, b) => b.percent - a.percent);
-	return { insights };
-}
-
-/**
- * Two-pointer sweep of messages sorted by timestamp. For each message, count the
- * number of distinct session IDs whose messages fall within an exact ± window.
- * Returns the total cost attributable to moments when ≥ threshold sessions were
- * active, or null if the period has too few sessions/messages to call it.
- *
- * Messages with missing/invalid timestamps (parsed as 0) are filtered out first —
- * otherwise they would collapse into a single synthetic instant and inflate the
- * parallel count on older or incomplete logs.
- */
-function computeParallelCostWeight(messages: RawMessage[]): number | null {
-	const timed = messages.filter((m) => m.timestamp > 0);
-	if (timed.length < MIN_MESSAGES_FOR_PARALLEL_INSIGHT) return null;
-	const distinctSessions = new Set(timed.map((m) => m.sessionId));
-	if (distinctSessions.size < PARALLEL_SESSION_THRESHOLD) return null;
-
-	const sorted = timed.slice().sort((a, b) => a.timestamp - b.timestamp);
-	const sidCount = new Map<string, number>();
-	let uniqueCount = 0;
-	let left = 0;
-	let right = 0;
-	let parallelCost = 0;
-
-	for (let i = 0; i < sorted.length; i++) {
-		const current = sorted[i]!;
-		const high = current.timestamp + PARALLEL_WINDOW_MS;
-		const low = current.timestamp - PARALLEL_WINDOW_MS;
-
-		while (right < sorted.length && sorted[right]!.timestamp <= high) {
-			const sid = sorted[right]!.sessionId;
-			const next = (sidCount.get(sid) ?? 0) + 1;
-			sidCount.set(sid, next);
-			if (next === 1) uniqueCount++;
-			right++;
-		}
-		while (left < right && sorted[left]!.timestamp < low) {
-			const sid = sorted[left]!.sessionId;
-			const next = (sidCount.get(sid) ?? 0) - 1;
-			if (next === 0) {
-				sidCount.delete(sid);
-				uniqueCount--;
-			} else {
-				sidCount.set(sid, next);
-			}
-			left++;
-		}
-
-		if (uniqueCount >= PARALLEL_SESSION_THRESHOLD) parallelCost += current.cost;
-	}
-
-	return parallelCost;
-}
-
-function formatThresholdTokens(n: number): string {
-	if (n >= 1_000_000) return `${n / 1_000_000}M`;
-	if (n >= 1_000) return `${n / 1_000}k`;
-	return String(n);
-}
-
-function formatInsightPercent(p: number): string {
-	if (p >= 10) return `${Math.round(p)}%`;
-	return `${Math.round(p * 10) / 10}%`;
-}
-
 // =============================================================================
 // Formatting Helpers
 // =============================================================================
@@ -685,6 +137,11 @@ function formatNumber(n: number): string {
 	return n.toLocaleString();
 }
 
+function formatInsightPercent(p: number): string {
+	if (p >= 10) return `${Math.round(p)}%`;
+	return `${Math.round(p * 10) / 10}%`;
+}
+
 function padLeft(s: string, len: number): string {
 	const vis = visibleWidth(s);
 	if (vis >= len) return s;
@@ -753,11 +210,10 @@ const TAB_LABELS: Record<TabName, string> = {
 	today: "Today",
 	thisWeek: "This Week",
 	lastWeek: "Last Week",
+	last30Days: "Last 30 Days",
 	allTime: "All Time",
 };
 
-const TAB_ORDER: TabName[] = ["today", "thisWeek", "lastWeek", "allTime"];
-
 class UsageComponent {
 	private activeTab: TabName = "allTime";
 	private viewMode: ViewMode = "table";
@@ -1090,7 +546,7 @@ export default function (pi: ExtensionAPI) {
 
 				loader.onAbort = () => finish(null);
 
-				collectUsageData(loader.signal)
+				collectUsageData({ signal: loader.signal })
 					.then(finish)
 					.catch(() => finish(null));
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tmustier/pi-usage-extension",
-  "version": "0.3.2",
+  "version": "0.5.0",
   "description": "Usage statistics dashboard for Pi sessions.",
   "license": "MIT",
   "author": "Thomas Mustier",
@@ -25,7 +25,7 @@
     "@earendil-works/pi-tui": "*"
   },
   "devDependencies": {
-    "@earendil-works/pi-coding-agent": "^0.74.0",
-    "@earendil-works/pi-tui": "^0.74.0"
+    "@earendil-works/pi-coding-agent": "^0.80.3",
+    "@earendil-works/pi-tui": "^0.80.3"
   }
 }