@tduyng/pi-cost 0.0.0-stage → 0.0.2

package/src/cost/index.ts

@@ -0,0 +1,460 @@
+import fs from 'node:fs/promises'
+import os from 'node:os'
+import path from 'node:path'
+import {
+    DynamicBorder,
+    type ExtensionAPI,
+    getAgentDir,
+    type ExtensionCommandContext,
+    type Theme
+} from '@earendil-works/pi-coding-agent'
+import { Container, matchesKey, Text, truncateToWidth } from '@earendil-works/pi-tui'
+
+import { addTurnCost, formatCostUsd, type ModelRates } from '../statusline/logic.js'
+import { loadCostConfig } from './config.js'
+import {
+    type AnthropicUsagePayload,
+    anthropicRows,
+    bar,
+    cachedQuota,
+    cachedRecords,
+    type CommandCodeCreditsPayload,
+    commandCodeRows,
+    formatCacheHit,
+    formatDuration,
+    formatReset,
+    formatTokenCell,
+    formatTokens,
+    type LimitRow,
+    type OpenCodeUsagePayload,
+    claudeAccessToken,
+    openCodeRows,
+    type ProviderTokenRow,
+    providerTokenRows,
+    readCostCache,
+    scanUsage,
+    type SessionTotals,
+    type SessionTurn,
+    sessionTotals,
+    tokenScanSince,
+    totalTokens,
+    writeCostCache
+} from './logic.js'
+import { readOverhead, sumOverhead, type OverheadTotals } from './overhead.js'
+
+const PANEL_SHARE = 0.7
+const REQUEST_TIMEOUT_MS = 10_000
+const KEYCHAIN_TIMEOUT_MS = 5_000
+const LABEL_WIDTH = 20
+const TOKEN_COLUMN_WIDTH = 9
+const TOKEN_HEADERS = ['in', 'out', 'cache rd', 'cache wr', 'total', 'cost']
+const PROVIDER_WIDTH = 15
+
+interface SessionSpend {
+    costUsd: number
+    totals: SessionTotals
+    overhead: OverheadTotals
+}
+
+interface Section {
+    provider: string
+    title: string
+    rows?: LimitRow[]
+    error?: string
+    fetchedAt?: number
+}
+
+interface QuotaSection {
+    provider: string
+    title: string
+    probe: (ctx: ExtensionCommandContext, pi: ExtensionAPI) => Promise<Section | undefined>
+}
+
+const QUOTA_SECTIONS: readonly QuotaSection[] = [
+    { provider: 'commandcode', title: 'Command Code', probe: commandCodeSection },
+    { provider: 'opencode-go', title: 'OpenCode Go', probe: openCodeGoSection },
+    { provider: 'claude', title: 'Claude', probe: (_ctx, pi) => claudeSection(pi) }
+]
+
+async function getJson<T>(url: string, headers: Record<string, string>): Promise<T> {
+    const response = await fetch(url, { headers, signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS) })
+    if (!response.ok) throw new Error(`HTTP ${response.status}`)
+    return (await response.json()) as T
+}
+
+function errorText(cause: unknown): string {
+    return cause instanceof Error ? cause.message : String(cause)
+}
+
+export async function claudeToken(
+    pi: ExtensionAPI,
+    platform: NodeJS.Platform = process.platform
+): Promise<string | undefined> {
+    let raw: string | undefined
+    if (platform === 'darwin') {
+        const result = await pi.exec('security', ['find-generic-password', '-s', 'Claude Code-credentials', '-w'], {
+            cwd: os.homedir(),
+            timeout: KEYCHAIN_TIMEOUT_MS
+        })
+        if (result.code && result.code !== 44) throw new Error(`keychain lookup failed (exit ${result.code})`)
+        raw = result.stdout
+    }
+    if (!raw) {
+        raw = await fs.readFile(path.join(os.homedir(), '.claude', '.credentials.json'), 'utf8').catch(() => undefined)
+    }
+    return claudeAccessToken(raw)
+}
+
+export async function fetchAnthropicUsage(token: string): Promise<AnthropicUsagePayload> {
+    return getJson<AnthropicUsagePayload>('https://api.anthropic.com/api/oauth/usage', {
+        authorization: `Bearer ${token}`,
+        'anthropic-beta': 'oauth-2025-04-20',
+        accept: 'application/json'
+    })
+}
+
+async function commandCodeSection(ctx: ExtensionCommandContext): Promise<Section | undefined> {
+    const key = await ctx.modelRegistry.getApiKeyForProvider('commandcode').catch(() => undefined)
+    if (!key) return undefined
+    const section: Section = { provider: 'commandcode', title: 'Command Code' }
+    try {
+        const payload = await getJson<CommandCodeCreditsPayload>('https://api.commandcode.ai/alpha/billing/credits', {
+            authorization: `Bearer ${key}`,
+            accept: 'application/json'
+        })
+        section.rows = commandCodeRows(payload)
+        if (section.rows.length === 0) section.error = 'no quota reported for this key'
+    } catch (error) {
+        section.error = errorText(error)
+    }
+    return section
+}
+
+async function openCodeGoSection(ctx: ExtensionCommandContext): Promise<Section | undefined> {
+    const key =
+        (await ctx.modelRegistry.getApiKeyForProvider('opencode-go').catch(() => undefined)) ||
+        process.env['OPENCODE_API_KEY']
+    if (!key) return undefined
+    const section: Section = { provider: 'opencode-go', title: 'OpenCode Go' }
+    try {
+        const payload = await getJson<OpenCodeUsagePayload>('https://opencode.ai/zen/go/v1/usage', {
+            authorization: `Bearer ${key}`,
+            accept: 'application/json'
+        })
+        section.rows = openCodeRows(payload)
+        if (section.rows.length === 0) section.error = 'no usage windows reported'
+    } catch (error) {
+        section.error = errorText(error)
+    }
+    return section
+}
+
+async function claudeSection(pi: ExtensionAPI): Promise<Section | undefined> {
+    const token = await claudeToken(pi)
+    if (!token) return undefined
+    const section: Section = { provider: 'claude', title: 'Claude' }
+    try {
+        section.rows = anthropicRows(await fetchAnthropicUsage(token))
+        if (section.rows.length === 0) section.error = 'no usage windows reported'
+    } catch (error) {
+        section.error = errorText(error)
+    }
+    return section
+}
+
+function modelRates(ctx: ExtensionCommandContext): ModelRates | undefined {
+    const cost = ctx.model?.cost
+    return cost
+        ? { input: cost.input, output: cost.output, cacheRead: cost.cacheRead, cacheWrite: cost.cacheWrite }
+        : undefined
+}
+
+function sessionSpend(ctx: ExtensionCommandContext, overheadRecords: ReturnType<typeof readOverhead>): SessionSpend {
+    const turns: SessionTurn[] = []
+    let costUsd = 0
+    let collapsed = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }
+    const currentRates = modelRates(ctx)
+    const ratesFor = (provider: string, model: string): ModelRates | undefined => {
+        const found = ctx.modelRegistry.find(provider, model)
+        const cost = found?.cost
+        return cost
+            ? { input: cost.input, output: cost.output, cacheRead: cost.cacheRead, cacheWrite: cost.cacheWrite }
+            : undefined
+    }
+    for (const entry of ctx.sessionManager.getBranch()) {
+        if (entry.type === 'branch_summary') {
+            if (entry.usage) {
+                const usage = entry.usage
+                collapsed = {
+                    input: collapsed.input + usage.input,
+                    output: collapsed.output + usage.output,
+                    cacheRead: collapsed.cacheRead + usage.cacheRead,
+                    cacheWrite: collapsed.cacheWrite + usage.cacheWrite
+                }
+                costUsd = addTurnCost(costUsd, usage, currentRates)
+            }
+            continue
+        }
+        if (entry.type !== 'message') continue
+        if (entry.message.role !== 'assistant') {
+            turns.push({ usage: undefined, timestamp: entry.timestamp })
+            continue
+        }
+        const usage = entry.message.usage
+        costUsd = addTurnCost(costUsd, usage, ratesFor(entry.message.provider, entry.message.model) ?? currentRates)
+        turns.push({ usage, timestamp: entry.timestamp })
+    }
+    const totals = sessionTotals(turns)
+    totals.input += collapsed.input
+    totals.output += collapsed.output
+    totals.cacheRead += collapsed.cacheRead
+    totals.cacheWrite += collapsed.cacheWrite
+    return { costUsd, totals, overhead: sumOverhead(overheadRecords) }
+}
+
+function sessionLines(session: SessionSpend, theme: Theme): string[] {
+    const label = (text: string) => theme.fg('dim', text.padEnd(LABEL_WIDTH))
+    const totals = session.totals
+    const hit = formatCacheHit(totals)
+    const extensions = [...session.overhead.byExtension].map(([name, spend]) => `${name} ${spend.calls}`).join(', ')
+    return [
+        `  ${label('Cost')}${formatCostUsd(session.costUsd)}`,
+        ...(session.overhead.calls === 0
+            ? []
+            : [
+                  `  ${label('Extensions')}${formatCostUsd(session.overhead.costUsd)} · ${formatTokens(session.overhead.input)} in · ${formatTokens(session.overhead.output)} out · ${session.overhead.calls} calls (${extensions})`
+              ]),
+        `  ${label('Tokens')}${formatTokens(totals.input)} in · ${formatTokens(totals.output)} out · ${formatTokens(totals.cacheRead)} cache read · ${formatTokens(totals.cacheWrite)} cache write`,
+        ...(hit === undefined
+            ? []
+            : [`  ${label('Cache hit')}${hit} ${theme.fg('dim', 'of prompt tokens served from cache')}`]),
+        `  ${label('Turns')}${totals.turns} in ${formatDuration(totals.wallMs)}`
+    ]
+}
+
+function tokenRowLine(row: ProviderTokenRow, repeated: boolean, cost: string, theme: Theme): string {
+    const label = `${(repeated ? '' : row.provider).padEnd(PROVIDER_WIDTH)}${row.window}`
+    const breakdown = [row.input, row.output, row.cacheRead, row.cacheWrite, totalTokens(row)]
+        .map((value) => formatTokenCell(value).padStart(TOKEN_COLUMN_WIDTH))
+        .join('')
+    return `  ${theme.fg('dim', label.padEnd(LABEL_WIDTH))}${theme.fg('dim', breakdown + cost.padStart(TOKEN_COLUMN_WIDTH))}`
+}
+
+function tokenLines(
+    rows: ProviderTokenRow[],
+    theme: Theme,
+    note: string | undefined,
+    costShown: (provider: string) => boolean
+): string[] {
+    const headers = TOKEN_HEADERS.map((header) => header.padStart(TOKEN_COLUMN_WIDTH)).join('')
+    const lines = [
+        `${theme.fg('accent', 'Tokens by provider')}  ${theme.fg('dim', 'all projects')}`,
+        `  ${' '.repeat(LABEL_WIDTH)}${theme.fg('dim', headers)}`,
+        ...rows.map((row, index) => {
+            const cost = !costShown(row.provider) ? '' : row.costUsd > 0 ? formatCostUsd(row.costUsd) : '·'
+            return tokenRowLine(row, row.provider === rows[index - 1]?.provider, cost, theme)
+        })
+    ]
+    if (note !== undefined) lines.push(`  ${theme.fg('dim', note)}`)
+    return lines
+}
+
+function severity(percent: number): 'error' | 'warning' | 'success' {
+    if (percent >= 90) return 'error'
+    if (percent >= 70) return 'warning'
+    return 'success'
+}
+
+function rowLine(row: LimitRow, theme: Theme): string {
+    const detail = [row.detail, formatReset(row.resetAt)].filter(Boolean).join('  ')
+    const label = `  ${theme.fg('dim', row.label.padEnd(LABEL_WIDTH))}`
+    if (row.percent === undefined) return `${label}${detail}`
+    const percent = `${Math.round(row.percent)}%`.padStart(4)
+    return [
+        label,
+        theme.fg(severity(row.percent), bar(row.percent)),
+        ` ${percent}`,
+        detail ? `  ${theme.fg('dim', detail)}` : ''
+    ].join('')
+}
+
+function sectionLines(section: Section, theme: Theme): string[] {
+    const age =
+        section.fetchedAt === undefined
+            ? ''
+            : `  ${theme.fg('dim', `${formatDuration(Date.now() - section.fetchedAt)} old`)}`
+    const lines = [`${theme.fg('accent', `${section.title} (${section.provider})`)}${age}`]
+    if (section.error !== undefined) lines.push(`  ${theme.fg('warning', section.error)}`)
+    if (section.rows) lines.push(...section.rows.map((row) => rowLine(row, theme)))
+    else if (section.error === undefined) lines.push(`  ${theme.fg('dim', 'loading…')}`)
+    return lines
+}
+
+export default function costExtension(pi: ExtensionAPI): void {
+    pi.registerCommand('cost', {
+        description: 'Show session spend and live provider rate-limit windows',
+        handler: async (_args, ctx) => {
+            if (ctx.mode !== 'tui') {
+                ctx.ui.notify('/cost requires TUI mode.', 'warning')
+                return
+            }
+
+            const cacheFile = path.join(getAgentDir(), 'cost-cache.json')
+            const cache = await readCostCache(cacheFile)
+            const { config, problem } = loadCostConfig()
+            if (problem !== undefined) ctx.ui.notify(problem, 'warning')
+            const costShown = (provider: string) => config.cost === undefined || config.cost.includes(provider)
+            const quotaSections = QUOTA_SECTIONS.filter(
+                ({ provider }) => config.limits === undefined || config.limits.includes(provider)
+            )
+
+            const sections = new Map<string, Section>(
+                quotaSections.map(({ provider, title }) => {
+                    const cached = cachedQuota(cache, provider)
+                    return [provider, cached ? { provider, title, ...cached } : { provider, title }]
+                })
+            )
+            let render: (() => void) | undefined
+
+            const settle = (provider: string, section: Section) => {
+                const previous = sections.get(provider)
+                if (section.error !== undefined && previous?.rows) {
+                    sections.set(provider, { ...section, rows: previous.rows, fetchedAt: previous.fetchedAt })
+                } else {
+                    sections.set(provider, section)
+                    if (section.rows) cache.quota[provider] = { rows: section.rows, fetchedAt: Date.now() }
+                }
+                render?.()
+            }
+
+            const resolve = (probe: Promise<Section | undefined>, provider: string) =>
+                probe.then(
+                    (section) => {
+                        if (section) {
+                            settle(provider, section)
+                            return
+                        }
+                        sections.delete(provider)
+                        delete cache.quota[provider]
+                        render?.()
+                    },
+                    (cause: unknown) =>
+                        settle(provider, {
+                            provider,
+                            title: sections.get(provider)?.title ?? provider,
+                            error: errorText(cause)
+                        })
+                )
+
+            const probes = quotaSections.map(({ provider, probe }) => resolve(probe(ctx, pi), provider))
+
+            const sessionDir = ctx.sessionManager.getSessionDir()
+            const agentDir = getAgentDir()
+            const defaultRoot = path.join(agentDir, 'sessions')
+            // Background jobs keep their own sessions under the agent directory, and their
+            // spend is this suite's least visible cost, so they are scanned with the rest.
+            const scanRoots = sessionDir.startsWith(defaultRoot + path.sep)
+                ? [defaultRoot, path.join(agentDir, 'jobs')]
+                : [sessionDir]
+            const since = tokenScanSince()
+            const seeded = cachedRecords(cache, since)
+            let windowRows = seeded.length > 0 ? providerTokenRows(seeded, Date.now(), config.providers) : undefined
+            const scan = scanUsage(scanRoots, since, cache).then((records) => {
+                windowRows = providerTokenRows(records, Date.now(), config.providers)
+                render?.()
+            })
+
+            void Promise.allSettled([...probes, scan]).then(() => writeCostCache(cacheFile, cache))
+
+            await ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
+                const container = new Container()
+                const border = new DynamicBorder((s: string) => theme.fg('accent', s))
+                const body = new Text('', 1, 0)
+                const footer = new Text('', 1, 0)
+                container.addChild(border)
+                container.addChild(body)
+                container.addChild(footer)
+                container.addChild(border)
+
+                const usage = ctx.getContextUsage()
+                const overheadRecords = readOverhead(getAgentDir(), ctx.sessionManager.getSessionId())
+                const header = [
+                    ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : undefined,
+                    usage?.percent != null ? `context ${Math.round(usage.percent)}%` : undefined
+                ]
+                    .filter(Boolean)
+                    .join('  ·  ')
+
+                let lines: string[] = []
+                let scroll = 0
+
+                const viewport = (): number => {
+                    const rows = tui.terminal.rows || 24
+                    const panelRows = Math.max(5, Math.min(rows - 2, Math.floor(rows * PANEL_SHARE)))
+                    return Math.max(1, panelRows - 3)
+                }
+
+                render = () => {
+                    const session = sessionSpend(ctx, overheadRecords)
+                    lines = [
+                        `${theme.fg('accent', theme.bold('Cost and limits'))}  ${theme.fg('dim', header)}`,
+                        '',
+                        theme.fg('accent', 'Session'),
+                        ...sessionLines(session, theme)
+                    ]
+                    for (const section of sections.values()) lines.push('', ...sectionLines(section, theme))
+                    if (sections.size === 0) {
+                        lines.push('', theme.fg('warning', 'No provider credentials found for a quota lookup.'))
+                    }
+                    lines.push(
+                        '',
+                        ...tokenLines(
+                            windowRows ?? [],
+                            theme,
+                            windowRows ? undefined : 'reading past sessions…',
+                            costShown
+                        )
+                    )
+                    tui.requestRender()
+                }
+                render()
+
+                return {
+                    render: (width: number) => {
+                        const height = viewport()
+                        const maxScroll = Math.max(0, lines.length - height)
+                        scroll = Math.min(scroll, maxScroll)
+                        const inner = Math.max(10, width - 2)
+                        const visible = lines.slice(scroll, scroll + height)
+                        while (visible.length < height) visible.push('')
+                        body.setText(visible.map((line) => truncateToWidth(line, inner, '...')).join('\n'))
+                        const position =
+                            maxScroll > 0
+                                ? `${scroll + 1}-${Math.min(lines.length, scroll + height)}/${lines.length}  ·  ↑/↓ scroll  ·  `
+                                : ''
+                        footer.setText(theme.fg('dim', `${position}Enter or Esc to close`))
+                        return container.render(width)
+                    },
+                    invalidate: () => container.invalidate(),
+                    handleInput: (data: string) => {
+                        if (matchesKey(data, 'enter') || matchesKey(data, 'escape')) {
+                            render = undefined
+                            done(undefined)
+                            return
+                        }
+                        const height = viewport()
+                        const maxScroll = Math.max(0, lines.length - height)
+                        let next = scroll
+                        if (matchesKey(data, 'down')) next += 1
+                        else if (matchesKey(data, 'up')) next -= 1
+                        else if (matchesKey(data, 'pageDown')) next += height
+                        else if (matchesKey(data, 'pageUp')) next -= height
+                        else return
+                        scroll = Math.min(Math.max(0, next), maxScroll)
+                        tui.requestRender()
+                    }
+                }
+            })
+        }
+    })
+}