diff --git a/src/main/cost-tracker.test.ts b/src/main/cost-tracker.test.ts new file mode 100644 index 0000000..265e05c --- /dev/null +++ b/src/main/cost-tracker.test.ts @@ -0,0 +1,468 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import type { LogAdapter, TokenUsage } from './log-adapters' + +// --------------------------------------------------------------------------- +// Mocks +// --------------------------------------------------------------------------- + +type ExecFileCb = (err: Error | null, stdout: string, stderr: string) => void + +const mockExecFile = vi.fn() + +vi.mock('child_process', () => ({ + execFile: (...args: unknown[]) => mockExecFile(...args), +})) + +vi.mock('./logger', () => ({ + createLogger: () => ({ + info: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + error: vi.fn(), + }), +})) + +// Dynamic import so mocks are wired before module evaluation +const { createCostTracker } = await import('./cost-tracker') + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +function makeMockWindow(destroyed = false): import('electron').BrowserWindow { + const send = vi.fn() + return { + isDestroyed: () => destroyed, + webContents: { send }, + } as unknown as import('electron').BrowserWindow +} + +/** Create a minimal adapter for testing */ +function makeTestAdapter(overrides: Partial = {}): LogAdapter { + return { + agent: 'claude-code', + getLogDirs: () => ['~/.claude/projects/test/sessions/'], + getFilePattern: () => '*.jsonl', + matchSession: () => true, + parseUsage: (_line: string, acc: TokenUsage): TokenUsage | null => { + return { + inputTokens: acc.inputTokens + 100, + outputTokens: acc.outputTokens + 50, + cacheReadTokens: acc.cacheReadTokens, + cacheWriteTokens: acc.cacheWriteTokens, + totalCostUsd: acc.totalCostUsd + 0.01, + } + }, + ...overrides, + } +} + +const BIND_OPTS = { + agent: 'claude-code', + projectPath: '/home/rooty/project', + cwd: '/home/rooty/project', + spawnAt: Date.now(), +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +beforeEach(() => { + vi.useFakeTimers() + mockExecFile.mockReset() +}) + +afterEach(() => { + vi.useRealTimers() +}) + +// ─── bindSession ──────────────────────────────────────────────────── + +describe('bindSession', () => { + it('is a no-op for unsupported agents', () => { + const win = makeMockWindow() + const adapter = makeTestAdapter({ agent: 'claude-code' }) + const tracker = createCostTracker(win, [adapter]) + + tracker.bindSession('s1', { ...BIND_OPTS, agent: 'goose' }) + + // No WSL calls should have been made + expect(mockExecFile).not.toHaveBeenCalled() + + tracker.destroy() + }) + + it('starts discovery polling for a supported agent', async () => { + const win = makeMockWindow() + const adapter = makeTestAdapter() + const tracker = createCostTracker(win, [adapter]) + + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + cb(null, '', '') + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Advance past first discovery poll — async version flushes microtasks + await vi.advanceTimersByTimeAsync(2000) + + expect(mockExecFile).toHaveBeenCalled() + + tracker.destroy() + }) +}) + +// ─── unbindSession ────────────────────────────────────────────────── + +describe('unbindSession', () => { + it('clears session and stops timers', async () => { + const win = makeMockWindow() + const adapter = makeTestAdapter() + const tracker = createCostTracker(win, [adapter]) + + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + cb(null, '', '') + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + tracker.unbindSession('s1') + + // Advancing timers should NOT trigger any more WSL calls + mockExecFile.mockReset() + await vi.advanceTimersByTimeAsync(10_000) + expect(mockExecFile).not.toHaveBeenCalled() + + tracker.destroy() + }) + + it('is a no-op for unknown sessionId', () => { + const win = makeMockWindow() + const tracker = createCostTracker(win, []) + + // Should not throw + tracker.unbindSession('nonexistent') + + tracker.destroy() + }) +}) + +// ─── destroy ──────────────────────────────────────────────────────── + +describe('destroy', () => { + it('clears all sessions and stops all timers', async () => { + const win = makeMockWindow() + const adapter = makeTestAdapter() + const tracker = createCostTracker(win, [adapter]) + + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + cb(null, '', '') + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + tracker.bindSession('s2', { ...BIND_OPTS, agent: 'claude-code' }) + tracker.destroy() + + // Advancing timers should NOT trigger any more WSL calls + mockExecFile.mockReset() + await vi.advanceTimersByTimeAsync(60_000) + expect(mockExecFile).not.toHaveBeenCalled() + }) +}) + +// ─── File discovery ───────────────────────────────────────────────── + +describe('file discovery', () => { + it('finds a log file and begins tailing', async () => { + const win = makeMockWindow() + const adapter = makeTestAdapter({ + matchSession: () => true, + }) + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + // Discovery: find returns a file + cb(null, '/home/rooty/.claude/projects/test/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + // Read first 3 lines for matchSession + cb(null, '{"cwd":"/home/rooty/project"}\n{"type":"init"}\n{"type":"start"}\n', '') + } else { + // Tailing poll: stat size + new content + cb(null, '50\n{"message":{"usage":{"input_tokens":100}}}\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Discovery poll at 2s (find + head resolve via microtasks) + await vi.advanceTimersByTimeAsync(2000) + + // Tailing poll at 3s after discovery + await vi.advanceTimersByTimeAsync(3000) + + expect(callCount).toBeGreaterThanOrEqual(3) + + tracker.destroy() + }) + + it('stops discovery after 30s with no match', async () => { + const win = makeMockWindow() + const adapter = makeTestAdapter({ + matchSession: () => false, + }) + const tracker = createCostTracker(win, [adapter]) + + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + // find always returns a file, but matchSession always returns false + cb(null, '/home/rooty/.claude/projects/test/sessions/abc.jsonl\n', '') + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Advance 32s — well past the 30s discovery timeout + await vi.advanceTimersByTimeAsync(32_000) + + // After discovery timeout, no more calls should happen + mockExecFile.mockReset() + await vi.advanceTimersByTimeAsync(10_000) + expect(mockExecFile).not.toHaveBeenCalled() + + tracker.destroy() + }) +}) + +// ─── File tailing ─────────────────────────────────────────────────── + +describe('file tailing', () => { + it('parses complete lines and sends cost:update IPC', async () => { + const win = makeMockWindow() + const send = (win.webContents as unknown as { send: ReturnType }).send + const parseUsage = vi.fn((_line: string, acc: TokenUsage): TokenUsage | null => { + return { + inputTokens: acc.inputTokens + 100, + outputTokens: acc.outputTokens + 50, + cacheReadTokens: acc.cacheReadTokens, + cacheWriteTokens: acc.cacheWriteTokens, + totalCostUsd: acc.totalCostUsd + 0.01, + } + }) + const adapter = makeTestAdapter({ parseUsage }) + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + // Discovery: find returns a file + cb(null, '/home/rooty/.claude/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + // head: session match + cb(null, '{"cwd":"/home/rooty/project"}\n', '') + } else if (callCount === 3) { + // First tail poll: stat + content with two complete lines + cb(null, '80\n{"line":"one"}\n{"line":"two"}\n', '') + } else { + // Subsequent polls: no new data + cb(null, '80\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Discovery poll (find + head) + await vi.advanceTimersByTimeAsync(2000) + + // Tail poll + await vi.advanceTimersByTimeAsync(3000) + + // parseUsage should have been called for each complete line + expect(parseUsage).toHaveBeenCalledTimes(2) + + // cost:update IPC should have been sent + expect(send).toHaveBeenCalledWith( + 'cost:update', + expect.objectContaining({ + sessionId: 's1', + usage: expect.objectContaining({ + inputTokens: 200, + outputTokens: 100, + }), + }), + ) + + tracker.destroy() + }) + + it('buffers partial lines until the next poll completes them', async () => { + const win = makeMockWindow() + const parseUsage = vi.fn((_line: string, acc: TokenUsage): TokenUsage | null => { + return { + inputTokens: acc.inputTokens + 100, + outputTokens: acc.outputTokens + 50, + cacheReadTokens: acc.cacheReadTokens, + cacheWriteTokens: acc.cacheWriteTokens, + totalCostUsd: acc.totalCostUsd + 0.01, + } + }) + const adapter = makeTestAdapter({ parseUsage }) + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + cb(null, '/home/rooty/.claude/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + cb(null, '{"cwd":"/home/rooty/project"}\n', '') + } else if (callCount === 3) { + // First tail: one complete line + one partial (no trailing newline) + cb(null, '60\n{"line":"one"}\n{"line":"tw', '') + } else if (callCount === 4) { + // Second tail: completes the partial line + cb(null, '80\no"}\n', '') + } else { + cb(null, '80\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Discovery + await vi.advanceTimersByTimeAsync(2000) + + // First tail — only line one is complete + await vi.advanceTimersByTimeAsync(3000) + expect(parseUsage).toHaveBeenCalledTimes(1) + + // Second tail — completes the partial line + await vi.advanceTimersByTimeAsync(3000) + expect(parseUsage).toHaveBeenCalledTimes(2) + + tracker.destroy() + }) + + it('resets offset when file is truncated (stat size < offset)', async () => { + const win = makeMockWindow() + const parseUsage = vi.fn((_line: string, acc: TokenUsage): TokenUsage | null => { + return { + inputTokens: acc.inputTokens + 100, + outputTokens: acc.outputTokens + 50, + cacheReadTokens: acc.cacheReadTokens, + cacheWriteTokens: acc.cacheWriteTokens, + totalCostUsd: acc.totalCostUsd + 0.01, + } + }) + const adapter = makeTestAdapter({ parseUsage }) + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + cb(null, '/home/rooty/.claude/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + cb(null, '{"cwd":"/home/rooty/project"}\n', '') + } else if (callCount === 3) { + // First tail: 80 bytes of content + cb(null, '80\n{"line":"one"}\n', '') + } else if (callCount === 4) { + // File was truncated: stat says 20 bytes, which is less than our offset + // The tracker should reset and re-read from beginning + cb(null, '20\n{"line":"reset"}\n', '') + } else { + cb(null, '20\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + + // Discovery + first tail + await vi.advanceTimersByTimeAsync(2000) + await vi.advanceTimersByTimeAsync(3000) + expect(parseUsage).toHaveBeenCalledTimes(1) + + // Truncation tail — re-polls immediately after reset, then the re-read + await vi.advanceTimersByTimeAsync(3000) + expect(parseUsage).toHaveBeenCalledTimes(2) + + tracker.destroy() + }) + + it('does not send IPC when window is destroyed', async () => { + const win = makeMockWindow(true) // destroyed + const send = (win.webContents as unknown as { send: ReturnType }).send + const adapter = makeTestAdapter() + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + cb(null, '/home/rooty/.claude/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + cb(null, '{"cwd":"/home/rooty/project"}\n', '') + } else { + cb(null, '80\n{"line":"one"}\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + await vi.advanceTimersByTimeAsync(2000) + await vi.advanceTimersByTimeAsync(3000) + + expect(send).not.toHaveBeenCalled() + + tracker.destroy() + }) + + it('does not call parseUsage when line is empty', async () => { + const win = makeMockWindow() + const parseUsage = vi.fn() + const adapter = makeTestAdapter({ parseUsage }) + const tracker = createCostTracker(win, [adapter]) + + let callCount = 0 + mockExecFile.mockImplementation( + (_bin: string, _args: string[], _opts: unknown, cb: ExecFileCb) => { + callCount++ + if (callCount === 1) { + cb(null, '/home/rooty/.claude/sessions/abc.jsonl\n', '') + } else if (callCount === 2) { + cb(null, '{"cwd":"/home/rooty/project"}\n', '') + } else { + // Tail returns stat + empty lines only + cb(null, '10\n\n\n', '') + } + }, + ) + + tracker.bindSession('s1', BIND_OPTS) + await vi.advanceTimersByTimeAsync(2000) + await vi.advanceTimersByTimeAsync(3000) + + expect(parseUsage).not.toHaveBeenCalled() + + tracker.destroy() + }) +}) diff --git a/src/main/cost-tracker.ts b/src/main/cost-tracker.ts new file mode 100644 index 0000000..81318e8 --- /dev/null +++ b/src/main/cost-tracker.ts @@ -0,0 +1,361 @@ +/** + * CostTracker — watches agent JSONL log files and pushes usage updates. + * + * For each bound PTY session, the tracker discovers the agent's log file + * (by polling candidate directories), then tails it every 3 seconds, + * parsing token-usage data via the matching LogAdapter and forwarding + * cumulative totals to the renderer over IPC. + */ +import type { BrowserWindow } from 'electron' +import { execFile } from 'child_process' +import { toWslPath } from './wsl-utils' +import { createLogger } from './logger' +import type { LogAdapter, TokenUsage } from './log-adapters' +import { ZERO_USAGE } from './log-adapters' + +const log = createLogger('cost-tracker') + +// ── Constants ─────────────────────────────────────────────────────── + +/** How often to poll for the log file during discovery (ms). */ +const DISCOVERY_INTERVAL_MS = 2000 + +/** How long to keep looking for a log file before giving up (ms). */ +const DISCOVERY_TIMEOUT_MS = 30_000 + +/** How often to poll the log file for new content once bound (ms). */ +const TAIL_INTERVAL_MS = 3000 + +/** Timeout for individual WSL exec calls (ms). */ +const WSL_TIMEOUT_MS = 5000 + +// ── Types ─────────────────────────────────────────────────────────── + +interface BoundSession { + sessionId: string + adapter: LogAdapter + projectPath: string + cwd: string + spawnAt: number + discoveryStartedAt: number + filePath: string | null + offset: number + partialLine: string + usage: TokenUsage + pollTimer: ReturnType | null +} + +export interface CostTracker { + bindSession( + sessionId: string, + opts: { + agent: string + projectPath: string + cwd: string + spawnAt: number + }, + ): void + unbindSession(sessionId: string): void + destroy(): void +} + +// ── Helpers ───────────────────────────────────────────────────────── + +/** Single-quote a path for safe use inside bash -lc commands. */ +function sq(s: string): string { + return "'" + s.replace(/'/g, "'\\''") + "'" +} + +/** + * Run a command inside WSL bash and return stdout. + * Resolves with stdout on success; rejects on error. + */ +function wslExec(cmd: string): Promise { + return new Promise((resolve, reject) => { + execFile('wsl.exe', ['bash', '-lc', cmd], { timeout: WSL_TIMEOUT_MS }, (err, stdout) => { + if (err) reject(err) + else resolve(stdout) + }) + }) +} + +// ── Factory ───────────────────────────────────────────────────────── + +export function createCostTracker(mainWindow: BrowserWindow, adapters: LogAdapter[]): CostTracker { + const sessions = new Map() + + // ── Discovery ─────────────────────────────────────────────────── + + function startDiscovery(session: BoundSession): void { + const dirs = session.adapter.getLogDirs(session.projectPath) + const pattern = session.adapter.getFilePattern() + + function discoveryPoll(): void { + // Session may have been unbound while we were waiting + if (!sessions.has(session.sessionId)) return + + // Timeout guard — uses the session-level start time so re-entries + // from tryMatchCandidates share the same deadline + if (Date.now() - session.discoveryStartedAt > DISCOVERY_TIMEOUT_MS) { + log.warn('Discovery timed out for session', { + sessionId: session.sessionId, + dirs, + pattern, + }) + return + } + + // Expand ~ in each dir and search for matching files modified after spawnAt + const findParts = dirs.map( + (d) => + `find ${sq(d.replace('~', '$HOME'))} -name ${sq(pattern)} -newermt @${Math.floor((session.spawnAt - 2000) / 1000)} 2>/dev/null`, + ) + const findCmd = findParts.join('; ') + + wslExec(findCmd) + .then((stdout) => { + if (!sessions.has(session.sessionId)) return + + const candidates = stdout + .split('\n') + .map((l) => l.trim()) + .filter(Boolean) + + if (candidates.length === 0) { + // Nothing yet — schedule next discovery poll + session.pollTimer = setTimeout(discoveryPoll, DISCOVERY_INTERVAL_MS) + return + } + + // Try each candidate: read first 3 lines and check matchSession + return tryMatchCandidates(session, candidates, 0) + }) + .catch((err) => { + if (!sessions.has(session.sessionId)) return + log.debug('Discovery find failed (will retry)', { + sessionId: session.sessionId, + err: String(err), + }) + session.pollTimer = setTimeout(discoveryPoll, DISCOVERY_INTERVAL_MS) + }) + } + + // Start first discovery poll + session.pollTimer = setTimeout(discoveryPoll, DISCOVERY_INTERVAL_MS) + } + + function tryMatchCandidates( + session: BoundSession, + candidates: string[], + index: number, + ): Promise { + if (!sessions.has(session.sessionId)) return Promise.resolve() + if (index >= candidates.length) { + // No match in this batch — schedule another discovery poll + session.pollTimer = setTimeout(() => { + startDiscovery(session) + }, DISCOVERY_INTERVAL_MS) + return Promise.resolve() + } + + const candidate = candidates[index] + if (!candidate) { + return tryMatchCandidates(session, candidates, index + 1) + } + + return wslExec(`head -n 3 ${sq(candidate)}`) + .then((headOutput): Promise | void => { + if (!sessions.has(session.sessionId)) return + + const firstLines = headOutput.split('\n').filter(Boolean) + if (session.adapter.matchSession(firstLines, session.cwd, session.spawnAt)) { + // Match found — bind and start tailing + session.filePath = candidate + log.info('Discovered log file for session', { + sessionId: session.sessionId, + filePath: candidate, + }) + startTailing(session) + } else { + // Try next candidate + return tryMatchCandidates(session, candidates, index + 1) + } + }) + .catch(() => { + if (!sessions.has(session.sessionId)) return + // head failed for this candidate — try next + return tryMatchCandidates(session, candidates, index + 1) + }) + } + + // ── Tailing ───────────────────────────────────────────────────── + + function startTailing(session: BoundSession): void { + function tailPoll(): void { + if (!sessions.has(session.sessionId)) return + if (!session.filePath) return + + // Get file size and new content in one call. + // tail -c +N is 1-indexed: +1 reads from byte 0, so we use offset+1. + const cmd = `stat -c %s ${sq(session.filePath)} && tail -c +${session.offset + 1} ${sq(session.filePath)}` + + wslExec(cmd) + .then((stdout) => { + if (!sessions.has(session.sessionId)) return + + // First line is the file size from stat, rest is content + const newlineIdx = stdout.indexOf('\n') + if (newlineIdx === -1) { + session.pollTimer = setTimeout(tailPoll, TAIL_INTERVAL_MS) + return + } + + const statLine = stdout.slice(0, newlineIdx).trim() + const fileSize = parseInt(statLine, 10) + const content = stdout.slice(newlineIdx + 1) + + // Truncation detection: file shrank since last read + if (!isNaN(fileSize) && fileSize < session.offset) { + log.debug('File truncated, resetting offset', { + sessionId: session.sessionId, + oldOffset: session.offset, + newSize: fileSize, + }) + session.offset = 0 + session.partialLine = '' + // Re-poll immediately to read from the start + session.pollTimer = setTimeout(tailPoll, 0) + return + } + + // No new content + if (!content) { + session.pollTimer = setTimeout(tailPoll, TAIL_INTERVAL_MS) + return + } + + // Update offset by actual bytes read + const contentBytes = Buffer.byteLength(content, 'utf8') + session.offset += contentBytes + + // Split on newlines, keeping the last incomplete part in partialLine + const text = session.partialLine + content + const lines = text.split('\n') + session.partialLine = lines.pop() ?? '' + + let usageChanged = false + for (const line of lines) { + const trimmed = line.trim() + if (!trimmed) continue + + const result = session.adapter.parseUsage(trimmed, session.usage) + if (result !== null) { + session.usage = result + usageChanged = true + } + } + + if (usageChanged && !mainWindow.isDestroyed()) { + mainWindow.webContents.send('cost:update', { + sessionId: session.sessionId, + usage: { ...session.usage }, + }) + } + + session.pollTimer = setTimeout(tailPoll, TAIL_INTERVAL_MS) + }) + .catch((err) => { + if (!sessions.has(session.sessionId)) return + log.debug('Tail poll failed (will retry)', { + sessionId: session.sessionId, + err: String(err), + }) + session.pollTimer = setTimeout(tailPoll, TAIL_INTERVAL_MS) + }) + } + + session.pollTimer = setTimeout(tailPoll, TAIL_INTERVAL_MS) + } + + // ── Public API ────────────────────────────────────────────────── + + function bindSession( + sessionId: string, + opts: { + agent: string + projectPath: string + cwd: string + spawnAt: number + }, + ): void { + // Find matching adapter + const adapter = adapters.find((a) => a.agent === opts.agent) + if (!adapter) { + log.debug('No adapter for agent, skipping cost tracking', { agent: opts.agent }) + return + } + + // Don't re-bind if already tracking + if (sessions.has(sessionId)) { + log.debug('Session already bound, skipping', { sessionId }) + return + } + + // Convert Windows paths to WSL format for log directory lookup + const wslCwd = toWslPath(opts.cwd) + const wslProjectPath = toWslPath(opts.projectPath) + + const session: BoundSession = { + sessionId, + adapter, + projectPath: wslProjectPath, + cwd: wslCwd, + spawnAt: opts.spawnAt, + discoveryStartedAt: Date.now(), + filePath: null, + offset: 0, + partialLine: '', + usage: { ...ZERO_USAGE }, + pollTimer: null, + } + + sessions.set(sessionId, session) + log.info('Bound session for cost tracking', { + sessionId, + agent: opts.agent, + cwd: opts.cwd, + }) + + startDiscovery(session) + } + + function unbindSession(sessionId: string): void { + const session = sessions.get(sessionId) + if (!session) return + + if (session.pollTimer !== null) { + clearTimeout(session.pollTimer) + session.pollTimer = null + } + + log.info('Unbound session from cost tracking', { + sessionId, + usage: session.usage, + }) + + sessions.delete(sessionId) + } + + function destroy(): void { + for (const session of sessions.values()) { + if (session.pollTimer !== null) { + clearTimeout(session.pollTimer) + session.pollTimer = null + } + } + sessions.clear() + log.info('CostTracker destroyed') + } + + return { bindSession, unbindSession, destroy } +} diff --git a/src/main/index.ts b/src/main/index.ts index e862ef7..e726a34 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -1,4 +1,4 @@ -import { app, BrowserWindow, safeStorage, screen } from 'electron' +import { app, BrowserWindow, ipcMain, safeStorage, screen } from 'electron' import { join } from 'path' import { createPtyManager, type PtyManager } from './pty-manager' import { createProjectStore, type AppStore } from './project-store' @@ -10,6 +10,8 @@ import { createWorkflowEngine } from './workflow-engine' import type { WorkflowEngine } from './workflow-engine' import { createWorktreeManager, type WorktreeManager } from './worktree-manager' import { createWslGitPort } from './git-port' +import { createCostTracker, type CostTracker } from './cost-tracker' +import { createClaudeAdapter, createCodexAdapter } from './log-adapters' import { registerPtyHandlers, registerWindowHandlers, @@ -28,6 +30,7 @@ let ptyManager: PtyManager | null = null let workflowEngine: WorkflowEngine | null = null let appStore: AppStore | null = null let worktreeManager: WorktreeManager | null = null +let costTracker: CostTracker | null = null // --- Crash cleanup handlers (REL-4) --- process.on('uncaughtException', (err) => { @@ -216,6 +219,24 @@ app createWindow() log.info('Window created') + if (mainWindow) { + costTracker = createCostTracker(mainWindow, [createClaudeAdapter(), createCodexAdapter()]) + } + + ipcMain.handle( + 'cost:bind', + ( + _, + sessionId: string, + opts: { agent: string; projectPath: string; cwd: string; spawnAt: number }, + ) => { + costTracker?.bindSession(sessionId, opts) + }, + ) + ipcMain.handle('cost:unbind', (_, sessionId: string) => { + costTracker?.unbindSession(sessionId) + }) + // Warn renderer if encryption is unavailable (secrets stored as plaintext) if (!safeStorage.isEncryptionAvailable() && mainWindow) { log.warn('safeStorage encryption unavailable — secrets stored as plaintext') @@ -251,6 +272,7 @@ app app.on('before-quit', () => { log.info('App quitting') + costTracker?.destroy() workflowEngine?.stopAll() ptyManager?.killAll() closeLogger() diff --git a/src/main/log-adapters.test.ts b/src/main/log-adapters.test.ts new file mode 100644 index 0000000..005ec95 --- /dev/null +++ b/src/main/log-adapters.test.ts @@ -0,0 +1,393 @@ +import { describe, it, expect } from 'vitest' +import { + formatTokens, + formatCost, + createClaudeAdapter, + createCodexAdapter, + ZERO_USAGE, +} from './log-adapters' + +// --------------------------------------------------------------------------- +// formatTokens +// --------------------------------------------------------------------------- + +describe('formatTokens', () => { + it('returns raw number under 1000', () => { + expect(formatTokens(500)).toBe('500') + }) + + it('returns raw "0" for zero', () => { + expect(formatTokens(0)).toBe('0') + }) + + it('formats 1500 as "1.5k"', () => { + expect(formatTokens(1500)).toBe('1.5k') + }) + + it('formats 12345 as "12.3k"', () => { + expect(formatTokens(12345)).toBe('12.3k') + }) + + it('formats exactly 1000 as "1.0k"', () => { + expect(formatTokens(1000)).toBe('1.0k') + }) +}) + +// --------------------------------------------------------------------------- +// formatCost +// --------------------------------------------------------------------------- + +describe('formatCost', () => { + it('returns empty string for 0', () => { + expect(formatCost(0)).toBe('') + }) + + it('formats 0.42 as "$0.42"', () => { + expect(formatCost(0.42)).toBe('$0.42') + }) + + it('formats 1.5 as "$1.50"', () => { + expect(formatCost(1.5)).toBe('$1.50') + }) + + it('formats large value with 2 decimals', () => { + expect(formatCost(12.999)).toBe('$13.00') + }) +}) + +// --------------------------------------------------------------------------- +// ClaudeAdapter +// --------------------------------------------------------------------------- + +describe('ClaudeAdapter', () => { + const adapter = createClaudeAdapter() + + it('has agent "claude-code"', () => { + expect(adapter.agent).toBe('claude-code') + }) + + it('getLogDirs returns 2 dirs', () => { + const dirs = adapter.getLogDirs('/home/rooty/my-project') + expect(dirs).toHaveLength(2) + }) + + it('getLogDirs first entry contains path slug', () => { + const dirs = adapter.getLogDirs('/home/rooty/my-project') + const first = dirs[0] + if (!first) throw new Error('Expected first dir') + // slashes replaced by dashes + expect(first).toContain('-home-rooty-my-project') + }) + + it('getLogDirs second entry is the glob fallback', () => { + const dirs = adapter.getLogDirs('/home/rooty/my-project') + const second = dirs[1] + if (!second) throw new Error('Expected second dir') + expect(second).toContain('~/.claude/projects/') + // fallback dir should not contain the path slug + expect(second).not.toContain('rooty') + }) + + it('getFilePattern returns "*.jsonl"', () => { + expect(adapter.getFilePattern()).toBe('*.jsonl') + }) + + it('parseUsage extracts usage from final Claude JSONL entry', () => { + const line = JSON.stringify({ + type: 'assistant', + message: { + model: 'claude-opus-4-6', + stop_reason: 'end_turn', + usage: { + input_tokens: 3, + output_tokens: 67, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 28750, + }, + }, + cwd: '/home/rooty/project', + }) + const result = adapter.parseUsage(line, { ...ZERO_USAGE }) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + expect(result.inputTokens).toBe(3) + expect(result.outputTokens).toBe(67) + expect(result.cacheReadTokens).toBe(0) + expect(result.cacheWriteTokens).toBe(28750) + }) + + it('parseUsage computes cost from model pricing', () => { + const line = JSON.stringify({ + type: 'assistant', + message: { + model: 'claude-opus-4-6', + stop_reason: 'end_turn', + usage: { + input_tokens: 1000, + output_tokens: 500, + cache_read_input_tokens: 10000, + cache_creation_input_tokens: 20000, + }, + }, + }) + const result = adapter.parseUsage(line, { ...ZERO_USAGE }) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + // opus: $5/1M input, $25/1M output + // input: 1000 × 5.0 / 1M = 0.005 + // cache read: 10000 × 5.0 × 0.1 / 1M = 0.005 + // cache write: 20000 × 5.0 × 1.25 / 1M = 0.125 + // output: 500 × 25.0 / 1M = 0.0125 + // total = 0.005 + 0.005 + 0.125 + 0.0125 = 0.1475 + expect(result.totalCostUsd).toBeCloseTo(0.1475) + }) + + it('parseUsage skips streaming partials (stop_reason: null)', () => { + const partial = JSON.stringify({ + type: 'assistant', + message: { + model: 'claude-opus-4-6', + stop_reason: null, + usage: { + input_tokens: 3, + output_tokens: 30, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 28750, + }, + }, + }) + expect(adapter.parseUsage(partial, { ...ZERO_USAGE })).toBeNull() + }) + + it('parseUsage accumulates across multiple final entries', () => { + const line = JSON.stringify({ + type: 'assistant', + message: { + model: 'claude-sonnet-4-6', + stop_reason: 'end_turn', + usage: { + input_tokens: 500, + output_tokens: 100, + cache_read_input_tokens: 28750, + cache_creation_input_tokens: 0, + }, + }, + }) + const acc = { + inputTokens: 3, + outputTokens: 67, + cacheReadTokens: 0, + cacheWriteTokens: 28750, + totalCostUsd: 0.18, + } + const result = adapter.parseUsage(line, acc) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + expect(result.inputTokens).toBe(503) + expect(result.outputTokens).toBe(167) + expect(result.cacheReadTokens).toBe(28750) + expect(result.cacheWriteTokens).toBe(28750) + expect(result.totalCostUsd).toBeGreaterThan(0.18) + }) + + it('parseUsage returns null for non-usage lines', () => { + const line = JSON.stringify({ type: 'user', message: { role: 'user', content: 'hello' } }) + expect(adapter.parseUsage(line, { ...ZERO_USAGE })).toBeNull() + }) + + it('parseUsage returns null for malformed JSON', () => { + expect(adapter.parseUsage('not json{{', { ...ZERO_USAGE })).toBeNull() + }) + + it('matchSession returns true when cwd found in lines', () => { + const lines = [JSON.stringify({ type: 'summary', cwd: '/home/rooty/project', ts: 1000 })] + expect(adapter.matchSession(lines, '/home/rooty/project', 1000)).toBe(true) + }) + + it('matchSession returns false when cwd not found', () => { + const lines = [JSON.stringify({ type: 'summary', cwd: '/home/rooty/other', ts: 1000 })] + expect(adapter.matchSession(lines, '/home/rooty/project', 1000)).toBe(false) + }) + + it('matchSession returns false for empty lines', () => { + expect(adapter.matchSession([], '/home/rooty/project', 1000)).toBe(false) + }) +}) + +// --------------------------------------------------------------------------- +// CodexAdapter +// --------------------------------------------------------------------------- + +describe('CodexAdapter', () => { + const adapter = createCodexAdapter() + + it('has agent "codex"', () => { + expect(adapter.agent).toBe('codex') + }) + + it('getLogDirs returns today date dir in YYYY/MM/DD format', () => { + const dirs = adapter.getLogDirs('/home/rooty/any') + expect(dirs).toHaveLength(1) + const dir = dirs[0] + if (!dir) throw new Error('Expected dir') + expect(dir).toContain('~/.codex/sessions/') + // Should match YYYY/MM/DD pattern at the end + expect(dir).toMatch(/\d{4}\/\d{2}\/\d{2}$/) + }) + + it('getFilePattern returns "rollout-*.jsonl"', () => { + expect(adapter.getFilePattern()).toBe('rollout-*.jsonl') + }) + + it('parseUsage extracts cumulative tokens from token_count event', () => { + // First, feed a turn_context line to set the model + const contextLine = JSON.stringify({ + type: 'turn_context', + payload: { turn_id: 't1', model: 'gpt-4o', cwd: '/home/rooty/project' }, + }) + adapter.parseUsage(contextLine, { ...ZERO_USAGE }) + + const line = JSON.stringify({ + type: 'event_msg', + payload: { + type: 'token_count', + info: { + total_token_usage: { + input_tokens: 2000, + output_tokens: 500, + cached_input_tokens: 100, + total_tokens: 2500, + }, + }, + }, + }) + const result = adapter.parseUsage(line, { ...ZERO_USAGE }) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + // inputTokens normalized: 2000 raw - 100 cached = 1900 non-cached + expect(result.inputTokens).toBe(1900) + expect(result.outputTokens).toBe(500) + expect(result.cacheReadTokens).toBe(100) + }) + + it('parseUsage replaces (not adds) accumulator for cumulative tokens', () => { + const contextLine = JSON.stringify({ + type: 'turn_context', + payload: { turn_id: 't1', model: 'o3' }, + }) + adapter.parseUsage(contextLine, { ...ZERO_USAGE }) + + const line = JSON.stringify({ + type: 'event_msg', + payload: { + type: 'token_count', + info: { + total_token_usage: { + input_tokens: 1000, + output_tokens: 200, + cached_input_tokens: 0, + total_tokens: 1200, + }, + }, + }, + }) + const acc = { + inputTokens: 9999, + outputTokens: 9999, + cacheReadTokens: 0, + cacheWriteTokens: 0, + totalCostUsd: 99, + } + const result = adapter.parseUsage(line, acc) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + expect(result.inputTokens).toBe(1000) + expect(result.outputTokens).toBe(200) + }) + + it('parseUsage computes cost from gpt-4o pricing using raw input (including cached)', () => { + const contextLine = JSON.stringify({ + type: 'turn_context', + payload: { turn_id: 't1', model: 'gpt-4o' }, + }) + adapter.parseUsage(contextLine, { ...ZERO_USAGE }) + + const line = JSON.stringify({ + type: 'event_msg', + payload: { + type: 'token_count', + info: { + total_token_usage: { + input_tokens: 1_000_000, + output_tokens: 1_000_000, + cached_input_tokens: 200_000, + total_tokens: 2_000_000, + }, + }, + }, + }) + const result = adapter.parseUsage(line, { ...ZERO_USAGE }) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + // Cost uses raw input (1M) not non-cached (800K) + // gpt-4o: input $2.50/1M + output $10.00/1M = $12.50 + expect(result.totalCostUsd).toBeCloseTo(12.5) + // Display tokens exclude cached + expect(result.inputTokens).toBe(800_000) + expect(result.cacheReadTokens).toBe(200_000) + }) + + it('parseUsage sets totalCostUsd to 0 for unknown model', () => { + const contextLine = JSON.stringify({ + type: 'turn_context', + payload: { turn_id: 't1', model: 'unknown-model-xyz' }, + }) + adapter.parseUsage(contextLine, { ...ZERO_USAGE }) + + const line = JSON.stringify({ + type: 'event_msg', + payload: { + type: 'token_count', + info: { + total_token_usage: { + input_tokens: 500, + output_tokens: 100, + cached_input_tokens: 0, + total_tokens: 600, + }, + }, + }, + }) + const result = adapter.parseUsage(line, { ...ZERO_USAGE }) + expect(result).not.toBeNull() + if (!result) throw new Error('Expected result') + expect(result.totalCostUsd).toBe(0) + }) + + it('parseUsage returns null for token_count with info: null', () => { + const line = JSON.stringify({ + type: 'event_msg', + payload: { type: 'token_count', info: null, rate_limits: {} }, + }) + expect(adapter.parseUsage(line, { ...ZERO_USAGE })).toBeNull() + }) + + it('parseUsage returns null for non-token_count events', () => { + const line = JSON.stringify({ type: 'event_msg', payload: { type: 'output', text: 'hello' } }) + expect(adapter.parseUsage(line, { ...ZERO_USAGE })).toBeNull() + }) + + it('parseUsage returns null for malformed JSON', () => { + expect(adapter.parseUsage('{bad json', { ...ZERO_USAGE })).toBeNull() + }) + + it('matchSession returns true when cwd found in lines', () => { + const lines = [JSON.stringify({ type: 'event', cwd: '/home/rooty/project', ts: 1000 })] + expect(adapter.matchSession(lines, '/home/rooty/project', 1000)).toBe(true) + }) + + it('matchSession returns false when cwd not found', () => { + const lines = [JSON.stringify({ type: 'event', cwd: '/home/rooty/other', ts: 1000 })] + expect(adapter.matchSession(lines, '/home/rooty/project', 1000)).toBe(false) + }) +}) diff --git a/src/main/log-adapters.ts b/src/main/log-adapters.ts new file mode 100644 index 0000000..1a920c8 --- /dev/null +++ b/src/main/log-adapters.ts @@ -0,0 +1,277 @@ +/** + * Log adapters for per-agent JSONL cost/token parsing. + * + * Each adapter knows where an agent writes its session logs, how to identify + * which log file corresponds to the current PTY session, and how to extract + * token-usage data from individual log lines. + */ + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface TokenUsage { + /** Non-cached input tokens (excludes cache reads for both Claude and Codex). */ + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number +} + +export const ZERO_USAGE: Readonly = Object.freeze({ + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + totalCostUsd: 0, +}) + +export interface LogAdapter { + agent: string + /** Return candidate directories (tilde-prefixed) to search for log files. */ + getLogDirs(projectPath: string): string[] + /** Glob pattern for log files within the dir. */ + getFilePattern(): string + /** + * Given the first few lines of a log file, decide whether this file belongs + * to the session spawned at `spawnAt` for the project at `cwd`. + */ + matchSession(firstLines: string[], cwd: string, spawnAt: number): boolean + /** + * Parse a single JSONL line and return an updated TokenUsage, or null if + * the line carries no usage data. + */ + parseUsage(line: string, accumulator: TokenUsage): TokenUsage | null +} + +// --------------------------------------------------------------------------- +// Formatters +// --------------------------------------------------------------------------- + +/** + * Format a token count for display. + * - Under 1 000: raw integer string. + * - 1 000 and above: one-decimal `k` suffix (e.g. `"1.5k"`, `"12.3k"`). + */ +export function formatTokens(n: number): string { + if (n < 1000) return String(n) + return `${(n / 1000).toFixed(1)}k` +} + +/** + * Format a USD cost for display. + * - Zero returns `""` (nothing to show). + * - Non-zero returns `"$X.XX"` with exactly 2 decimal places. + */ +export function formatCost(usd: number): string { + if (usd === 0) return '' + return `$${usd.toFixed(2)}` +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +/** Return today's date as `YYYY/MM/DD`. */ +function todayDateDir(): string { + const d = new Date() + const pad = (n: number): string => String(n).padStart(2, '0') + return `${d.getFullYear()}/${pad(d.getMonth() + 1)}/${pad(d.getDate())}` +} + +/** Return true if any stringified line contains the target substring. */ +function anyLineContains(lines: string[], target: string): boolean { + return lines.some((l) => l.includes(target)) +} + +// --------------------------------------------------------------------------- +// ClaudeAdapter +// --------------------------------------------------------------------------- + +/** Per-model pricing for Claude cost estimation. */ +const CLAUDE_PRICING: Record = { + opus: { inputPer1M: 5.0, outputPer1M: 25.0 }, + sonnet: { inputPer1M: 3.0, outputPer1M: 15.0 }, + haiku: { inputPer1M: 1.0, outputPer1M: 5.0 }, +} + +/** Match a Claude model ID (e.g. "claude-opus-4-6") to a pricing tier. */ +function getClaudePricing(model: string): { inputPer1M: number; outputPer1M: number } | undefined { + for (const [tier, pricing] of Object.entries(CLAUDE_PRICING)) { + if (model.includes(tier)) return pricing + } + return undefined +} + +export function createClaudeAdapter(): LogAdapter { + return { + agent: 'claude-code', + + getLogDirs(projectPath: string): string[] { + // Replace every `/` with `-` to match Claude's path-slug convention. + const pathSlug = projectPath.replace(/\//g, '-') + return [`~/.claude/projects/${pathSlug}/sessions/`, `~/.claude/projects/`] + }, + + getFilePattern(): string { + return '*.jsonl' + }, + + matchSession(firstLines: string[], cwd: string, _spawnAt: number): boolean { + return anyLineContains(firstLines, cwd) + }, + + parseUsage(line: string, accumulator: TokenUsage): TokenUsage | null { + let parsed: unknown + try { + parsed = JSON.parse(line) + } catch { + return null + } + + if (typeof parsed !== 'object' || parsed === null) return null + + const obj = parsed as Record + const message = obj['message'] + if (typeof message !== 'object' || message === null) return null + + const msg = message as Record + + // Claude Code logs both streaming partials (stop_reason: null) and the + // final entry (stop_reason: "end_turn" / "tool_use" / etc.) for the same + // API call, with identical usage blocks. Only count the final entry. + const stopReason = msg['stop_reason'] + if (stopReason === null || stopReason === undefined) return null + + const usage = msg['usage'] + if (typeof usage !== 'object' || usage === null) return null + + const u = usage as Record + const inputTokens = typeof u['input_tokens'] === 'number' ? u['input_tokens'] : 0 + const outputTokens = typeof u['output_tokens'] === 'number' ? u['output_tokens'] : 0 + const cacheReadTokens = + typeof u['cache_read_input_tokens'] === 'number' ? u['cache_read_input_tokens'] : 0 + const cacheWriteTokens = + typeof u['cache_creation_input_tokens'] === 'number' ? u['cache_creation_input_tokens'] : 0 + + // Compute cost from model pricing (Claude JSONL has no costUSD field). + // Cache writes cost 1.25× base input; cache reads cost 0.1× base input. + const model = typeof msg['model'] === 'string' ? (msg['model'] as string) : '' + const pricing = getClaudePricing(model) + const turnCost = + pricing !== undefined + ? (inputTokens / 1_000_000) * pricing.inputPer1M + + (cacheReadTokens / 1_000_000) * pricing.inputPer1M * 0.1 + + (cacheWriteTokens / 1_000_000) * pricing.inputPer1M * 1.25 + + (outputTokens / 1_000_000) * pricing.outputPer1M + : 0 + + return { + inputTokens: accumulator.inputTokens + inputTokens, + outputTokens: accumulator.outputTokens + outputTokens, + cacheReadTokens: accumulator.cacheReadTokens + cacheReadTokens, + cacheWriteTokens: accumulator.cacheWriteTokens + cacheWriteTokens, + totalCostUsd: accumulator.totalCostUsd + turnCost, + } + }, + } +} + +// --------------------------------------------------------------------------- +// CodexAdapter +// --------------------------------------------------------------------------- + +/** Per-model pricing for cost estimation. */ +const CODEX_PRICING: Record = { + 'gpt-4o': { inputPer1M: 2.5, outputPer1M: 10.0 }, + 'gpt-4o-mini': { inputPer1M: 0.15, outputPer1M: 0.6 }, + o3: { inputPer1M: 2.0, outputPer1M: 8.0 }, + 'o4-mini': { inputPer1M: 1.1, outputPer1M: 4.4 }, + 'codex-mini': { inputPer1M: 1.5, outputPer1M: 6.0 }, + 'gpt-5.3': { inputPer1M: 2.0, outputPer1M: 8.0 }, + 'gpt-5.4': { inputPer1M: 2.0, outputPer1M: 8.0 }, + 'gpt-5.3-codex': { inputPer1M: 2.0, outputPer1M: 8.0 }, +} + +export function createCodexAdapter(): LogAdapter { + let codexModel = '' + + return { + agent: 'codex', + + getLogDirs(_projectPath: string): string[] { + return [`~/.codex/sessions/${todayDateDir()}`] + }, + + getFilePattern(): string { + return 'rollout-*.jsonl' + }, + + matchSession(firstLines: string[], cwd: string, _spawnAt: number): boolean { + return anyLineContains(firstLines, cwd) + }, + + parseUsage(line: string, accumulator: TokenUsage): TokenUsage | null { + let parsed: unknown + try { + parsed = JSON.parse(line) + } catch { + return null + } + + if (typeof parsed !== 'object' || parsed === null) return null + + const obj = parsed as Record + const payload = obj['payload'] + if (typeof payload !== 'object' || payload === null) return null + + const p = payload as Record + + // Extract model from turn_context events for pricing lookup + if (p['turn_id'] && typeof p['model'] === 'string') { + // Store model in accumulator metadata — we use cacheWriteTokens as a hack-free approach: + // just remember the model name for later. Actually, return null and let token_count handle it. + // The model is available in turn_context, not in token_count. + // We'll extract it here and store via a side channel. + codexModel = p['model'] as string + } + + if (p['type'] !== 'token_count') return null + + // token_count events have payload.info with total_token_usage + // First event may have info: null (just rate limits) — skip those + const info = p['info'] + if (typeof info !== 'object' || info === null) return null + + const infoObj = info as Record + const usage = infoObj['total_token_usage'] + if (typeof usage !== 'object' || usage === null) return null + + const u = usage as Record + const rawInputTokens = u['input_tokens'] ?? 0 + const outputTokens = u['output_tokens'] ?? 0 + const cachedInputTokens = u['cached_input_tokens'] ?? 0 + + // Codex token_count events are CUMULATIVE — replace accumulator values. + // Codex input_tokens INCLUDES cached_input_tokens as a subset. + // Normalize to non-cached input only (matches Claude adapter semantics). + const nonCachedInput = rawInputTokens - cachedInputTokens + const model = codexModel + const pricing = CODEX_PRICING[model] + const totalCostUsd = + pricing !== undefined + ? (rawInputTokens / 1_000_000) * pricing.inputPer1M + + (outputTokens / 1_000_000) * pricing.outputPer1M + : 0 + + return { + inputTokens: nonCachedInput, + outputTokens, + cacheReadTokens: cachedInputTokens, + cacheWriteTokens: accumulator.cacheWriteTokens, + totalCostUsd, + } + }, + } +} diff --git a/src/preload/index.ts b/src/preload/index.ts index 8dcccee..f530bd7 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -217,6 +217,41 @@ contextBridge.exposeInMainWorld('agentDeck', { releasePrimary: (projectId: string, sessionId: string): Promise => ipcRenderer.invoke('worktree:releasePrimary', projectId, sessionId) as Promise, }, + cost: { + bind: ( + sessionId: string, + opts: { agent: string; projectPath: string; cwd: string; spawnAt: number }, + ): Promise => ipcRenderer.invoke('cost:bind', sessionId, opts), + unbind: (sessionId: string): Promise => ipcRenderer.invoke('cost:unbind', sessionId), + onUpdate: ( + cb: (data: { + sessionId: string + usage: { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number + } + }) => void, + ) => { + const listener = ( + _event: Electron.IpcRendererEvent, + data: { + sessionId: string + usage: { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number + } + }, + ): void => cb(data) + ipcRenderer.on('cost:update', listener) + return () => ipcRenderer.removeListener('cost:update', listener) + }, + }, onFileDrop: (cb: (wslPaths: string[]) => void) => { const listener = (_event: Electron.IpcRendererEvent, wslPaths: string[]): void => cb(wslPaths) ipcRenderer.on('file-dropped', listener) diff --git a/src/renderer/App.tsx b/src/renderer/App.tsx index 3a21bf0..8b81b76 100644 --- a/src/renderer/App.tsx +++ b/src/renderer/App.tsx @@ -136,6 +136,7 @@ export function App(): React.JSX.Element { window.agentDeck.pty.kill(sessionId).catch((err: unknown) => { window.agentDeck.log.send('debug', 'pty', 'Kill failed', { err: String(err) }) }) + window.agentDeck.cost.unbind(sessionId).catch(() => {}) removeSession(sessionId) }, [removeSession], @@ -211,6 +212,7 @@ export function App(): React.JSX.Element { window.agentDeck.pty.kill(sessionId).catch((err: unknown) => { window.agentDeck.log.send('debug', 'pty', 'Kill failed', { err: String(err) }) }) + window.agentDeck.cost.unbind(sessionId).catch(() => {}) window.agentDeck.worktree.discard(sessionId).catch((err: unknown) => { useAppStore .getState() @@ -229,6 +231,7 @@ export function App(): React.JSX.Element { window.agentDeck.pty.kill(sessionId).catch((err: unknown) => { window.agentDeck.log.send('debug', 'pty', 'Kill failed', { err: String(err) }) }) + window.agentDeck.cost.unbind(sessionId).catch(() => {}) window.agentDeck.worktree.keep(sessionId).catch((err: unknown) => { window.agentDeck.log.send('warn', 'worktree', 'Keep failed', { err: String(err) }) }) @@ -291,6 +294,14 @@ export function App(): React.JSX.Element { return unsub }, []) + // Listen for cost/token usage updates from main process + useEffect(() => { + const unsub = window.agentDeck.cost.onUpdate((data) => { + useAppStore.getState().setSessionUsage(data.sessionId, data.usage) + }) + return unsub + }, []) + // Listen for encryption unavailability warning useEffect(() => { const unsub = window.agentDeck.security.onEncryptionUnavailable(() => { diff --git a/src/renderer/components/SplitView/PaneTopbar.css b/src/renderer/components/SplitView/PaneTopbar.css index 39f3988..24bd60c 100644 --- a/src/renderer/components/SplitView/PaneTopbar.css +++ b/src/renderer/components/SplitView/PaneTopbar.css @@ -139,3 +139,15 @@ font-family: var(--font-mono); flex-shrink: 0; } + +.pane-cost-badge { + display: inline-flex; + align-items: center; + gap: 4px; + font-size: 11px; + color: var(--accent); + opacity: 0.85; + margin-left: 8px; + font-family: var(--font-mono); + flex-shrink: 0; +} diff --git a/src/renderer/components/SplitView/PaneTopbar.tsx b/src/renderer/components/SplitView/PaneTopbar.tsx index b4f299f..810569a 100644 --- a/src/renderer/components/SplitView/PaneTopbar.tsx +++ b/src/renderer/components/SplitView/PaneTopbar.tsx @@ -1,9 +1,29 @@ import { memo, useCallback } from 'react' -import { GitBranch } from 'lucide-react' +import { GitBranch, Zap } from 'lucide-react' import { useAppStore } from '../../store/appStore' import { HexDot } from '../shared/HexDot' import './PaneTopbar.css' +function fmtTokens(n: number): string { + if (n < 1000) return String(n) + return (n / 1000).toFixed(1) + 'k' +} + +function fmtCost(usd: number): string { + if (usd <= 0) return '' + return '$' + usd.toFixed(2) +} + +/** Total tokens processed (all types). Consistent with cost computation. */ +function totalTokens(u: { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number +}): number { + return u.inputTokens + u.cacheReadTokens + u.cacheWriteTokens + u.outputTokens +} + interface PaneTopbarProps { sessionId: string focused: boolean @@ -19,13 +39,18 @@ export const PaneTopbar = memo(function PaneTopbar({ projectId ? s.projects.find((p) => p.id === projectId) : undefined, ) const worktreeInfo = useAppStore((s) => s.worktreePaths[sessionId]) + const usage = useAppStore((s) => s.sessionUsage[sessionId]) const restartSession = useAppStore((s) => s.restartSession) + const agentOverride = useAppStore((s) => s.sessions[sessionId]?.agentOverride) const isTerminal = !projectId const accentColor = project?.identity?.accentColor ?? undefined const agentName = isTerminal ? 'shell' - : (project?.agents?.find((a) => a.isDefault)?.agent ?? project?.agent ?? 'claude-code') + : (agentOverride ?? + project?.agents?.find((a) => a.isDefault)?.agent ?? + project?.agent ?? + 'claude-code') // Extract a clean display name: use project name, but if it looks like a path, take the last segment const rawName = isTerminal ? 'Terminal' : (project?.name ?? 'Unknown') @@ -71,6 +96,17 @@ export const PaneTopbar = memo(function PaneTopbar({ {worktreeInfo.branch.split('/').pop()} )} + {usage && (usage.totalCostUsd > 0 || totalTokens(usage) > 0) && ( + + + {fmtCost(usage.totalCostUsd) && {fmtCost(usage.totalCostUsd)}} + {fmtCost(usage.totalCostUsd) && totalTokens(usage) > 0 && · } + {totalTokens(usage) > 0 && {fmtTokens(totalTokens(usage))} tokens} + + )} {showPath && ( <> > diff --git a/src/renderer/components/Terminal/TerminalPane.tsx b/src/renderer/components/Terminal/TerminalPane.tsx index 2aeb7b5..5179174 100644 --- a/src/renderer/components/Terminal/TerminalPane.tsx +++ b/src/renderer/components/Terminal/TerminalPane.tsx @@ -424,7 +424,19 @@ export function TerminalPane({ agentRef.current, agentFlagsRef.current, ) - if (!cancelled) setSessionStatus(sessionId, 'running') + if (cancelled) return + // Bind cost tracking (best-effort, fire-and-forget) + window.agentDeck.cost + .bind(sessionId, { + agent: agentRef.current ?? '', + projectPath: projectPathRef.current ?? '', + cwd: spawnPath ?? projectPathRef.current ?? '', + spawnAt: spawnTimestamp, + }) + .catch(() => { + /* cost tracking is best-effort */ + }) + setSessionStatus(sessionId, 'running') } catch (err: unknown) { if (cancelled) return window.agentDeck.log.send('error', 'terminal', `PTY spawn failed for ${sessionId}`, { @@ -620,6 +632,7 @@ export function TerminalPane({ } else { // Session removed → dispose everything clearWorktreePath(sessionId) + window.agentDeck.cost.unbind(sessionId).catch(() => {}) try { webglAddon?.dispose() } catch { diff --git a/src/renderer/global.d.ts b/src/renderer/global.d.ts index 216597f..0327e5a 100644 --- a/src/renderer/global.d.ts +++ b/src/renderer/global.d.ts @@ -124,6 +124,25 @@ declare global { keep(sessionId: string): Promise releasePrimary(projectId: string, sessionId: string): Promise } + cost: { + bind( + sessionId: string, + opts: { agent: string; projectPath: string; cwd: string; spawnAt: number }, + ): Promise + unbind(sessionId: string): Promise + onUpdate( + cb: (data: { + sessionId: string + usage: { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number + } + }) => void, + ): () => void + } pickFolder: () => Promise log: { send: (level: string, mod: string, message: string, data?: unknown) => Promise diff --git a/src/renderer/store/slices/sessions.ts b/src/renderer/store/slices/sessions.ts index d6d6dcf..fa9938c 100644 --- a/src/renderer/store/slices/sessions.ts +++ b/src/renderer/store/slices/sessions.ts @@ -20,11 +20,34 @@ export interface SessionsSlice { activityFeeds: Record addActivityEvent: (sessionId: string, event: ActivityEvent) => void clearActivityFeed: (sessionId: string) => void + + // Usage tracking (per-session) + sessionUsage: Record< + string, + { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number + } + > + setSessionUsage: ( + sessionId: string, + usage: { + inputTokens: number + outputTokens: number + cacheReadTokens: number + cacheWriteTokens: number + totalCostUsd: number + }, + ) => void } export const createSessionsSlice: StateCreator = (set, get) => ({ sessions: {}, activeSessionId: null, + sessionUsage: {}, addSession: (sessionId, projectId, overrides) => set((state) => { @@ -85,6 +108,7 @@ export const createSessionsSlice: StateCreator set((state) => { const { [sessionId]: _, ...rest } = state.sessions const { [sessionId]: _feed, ...remainingFeeds } = state.activityFeeds + const { [sessionId]: _usage, ...remainingUsage } = state.sessionUsage const remainingIds = Object.keys(rest) // Clear removed session from pane slots, then compact left so pane 0 always // has a session if any exist (prevents empty pane with sessions in hidden slots) @@ -100,6 +124,7 @@ export const createSessionsSlice: StateCreator return { sessions: rest, activityFeeds: remainingFeeds, + sessionUsage: remainingUsage, activeSessionId: state.activeSessionId === sessionId ? newActive : state.activeSessionId, currentView: remainingIds.length === 0 @@ -165,6 +190,10 @@ export const createSessionsSlice: StateCreator return Object.values(sessions).find((s) => s.projectId === projectId) }, + // Usage tracking + setSessionUsage: (sessionId, usage) => + set((s) => ({ sessionUsage: { ...s.sessionUsage, [sessionId]: usage } })), + // Activity Feed activityFeeds: {},