@tduyng/pi-cost 0.0.0-stage → 0.0.1

package/src/statusline/index.ts

@@ -0,0 +1,161 @@
+import fs from 'node:fs'
+import path from 'node:path'
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent'
+
+import { averageTps, type SpeedSample } from './logic.js'
+
+export default function statuslineExtension(pi: ExtensionAPI): void {
+    pi.registerFlag('statusline-no-pr', { description: 'Skip the GitHub PR lookup in the statusline', type: 'boolean' })
+
+    const showPr = pi.getFlag('statusline-no-pr') !== true
+
+    let enabled = true
+    const speedSamples: SpeedSample[] = []
+    let responseStartedAt: number | undefined
+
+    let changedFiles: number | undefined
+    let prNumber: number | undefined
+    let runningAgents = 0
+
+    let gitHeadWatcher: fs.FSWatcher | undefined
+    let gitWatchDebounce: ReturnType<typeof setTimeout> | undefined
+    let mountedCtx: ExtensionContext | undefined
+
+    pi.on('session_start', async (_event, ctx) => {
+        if (!ctx.hasUI) return
+        runningAgents = 0
+        speedSamples.length = 0
+        mount(ctx)
+    })
+
+    pi.on('session_shutdown', async () => {
+        unmount()
+    })
+
+    pi.on('message_start', (event) => {
+        if (event.message.role === 'assistant') responseStartedAt = Date.now()
+    })
+
+    pi.on('message_end', async (event) => {
+        const message = event.message
+        if (message.role !== 'assistant') return
+        const usage = message.usage
+        if (!usage) return
+        const ms = responseStartedAt !== undefined ? Date.now() - responseStartedAt : 0
+        if (ms > 0) speedSamples.push({ outputTokens: usage.output, ms })
+        responseStartedAt = undefined
+        sync()
+    })
+
+    pi.events.on('subagent:started', () => {
+        runningAgents += 1
+        sync()
+    })
+    pi.events.on('subagent:completed', () => {
+        runningAgents = Math.max(0, runningAgents - 1)
+        sync()
+    })
+
+    pi.on('agent_settled', () => {
+        void refreshGit()
+    })
+
+    pi.registerCommand('statusline', {
+        description: 'Toggle the extra statusline segments',
+        handler: async (_args, ctx) => {
+            enabled = !enabled
+            if (!ctx.hasUI) return
+            if (enabled) mount(ctx)
+            else unmount()
+            ctx.ui.notify(`statusline ${enabled ? 'enabled' : 'disabled'}`, 'info')
+        }
+    })
+
+    async function refreshGit(): Promise<void> {
+        const ctx = mountedCtx
+        if (!ctx) return
+        try {
+            const status = await pi.exec('git', ['status', '--porcelain'], { cwd: ctx.cwd })
+            changedFiles = status.stdout.split('\n').filter((line) => line.trim().length > 0).length
+        } catch {
+            changedFiles = undefined
+        }
+        try {
+            sync()
+        } catch {}
+    }
+
+    async function refreshPr(): Promise<void> {
+        const ctx = mountedCtx
+        if (!ctx || !showPr) return
+        try {
+            const pr = await pi.exec('gh', ['pr', 'view', '--json', 'number', '-q', '.number'], { cwd: ctx.cwd })
+            const parsed = Number.parseInt(pr.stdout.trim(), 10)
+            prNumber = Number.isNaN(parsed) ? undefined : parsed
+        } catch {
+            prNumber = undefined
+        }
+        try {
+            sync()
+        } catch {}
+    }
+
+    function sync(): void {
+        const ctx = mountedCtx
+        if (!ctx?.hasUI || !enabled) return
+        const theme = ctx.ui.theme
+
+        const tps = averageTps(speedSamples)
+        ctx.ui.setStatus('tps', tps !== undefined && tps > 0 ? theme.fg('dim', `${tps} tok/s`) : undefined)
+
+        ctx.ui.setStatus(
+            'agents',
+            runningAgents > 0
+                ? theme.fg('accent', `${runningAgents} agent${runningAgents === 1 ? '' : 's'}`)
+                : undefined
+        )
+
+        const gitParts: string[] = []
+        if (changedFiles !== undefined && changedFiles > 0) gitParts.push(`${changedFiles}±`)
+        if (prNumber !== undefined) gitParts.push(`#${prNumber}`)
+        ctx.ui.setStatus('git', gitParts.length > 0 ? theme.fg('success', gitParts.join(' · ')) : undefined)
+    }
+
+    function mount(ctx: ExtensionContext): void {
+        if (!ctx.hasUI) return
+        mountedCtx = ctx
+
+        void refreshGit()
+        void refreshPr()
+        watchGitHead()
+        sync()
+    }
+
+    function watchGitHead(): void {
+        const ctx = mountedCtx
+        gitHeadWatcher?.close()
+        gitHeadWatcher = undefined
+        if (!ctx?.cwd) return
+        try {
+            gitHeadWatcher = fs.watch(path.join(ctx.cwd, '.git', 'HEAD'), { persistent: false }, () => {
+                if (gitWatchDebounce) clearTimeout(gitWatchDebounce)
+                gitWatchDebounce = setTimeout(() => {
+                    void refreshGit()
+                    void refreshPr()
+                }, 250)
+            })
+        } catch {}
+    }
+
+    function unmount(): void {
+        gitHeadWatcher?.close()
+        gitHeadWatcher = undefined
+        if (gitWatchDebounce) clearTimeout(gitWatchDebounce)
+        gitWatchDebounce = undefined
+        const ctx = mountedCtx
+        if (ctx?.hasUI) {
+            for (const key of ['tps', 'agents', 'git']) ctx.ui.setStatus(key, undefined)
+        }
+        mountedCtx = undefined
+    }
+}

package/src/statusline/logic.ts

@@ -0,0 +1,68 @@
+export interface UsageCost {
+    input: number
+    output: number
+    cacheRead: number
+    cacheWrite: number
+    total: number
+}
+
+export interface TokenUsage {
+    input: number
+    output: number
+    cacheRead: number
+    cacheWrite: number
+    cost?: UsageCost
+}
+
+export interface ModelRates {
+    input: number
+    output: number
+    cacheRead: number
+    cacheWrite: number
+}
+
+function hasPricing(rates: ModelRates | undefined): boolean {
+    if (!rates) return false
+    return rates.input > 0 || rates.output > 0 || rates.cacheRead > 0 || rates.cacheWrite > 0
+}
+
+export function turnCostUsd(usage: TokenUsage, rates?: ModelRates): number {
+    if (usage.cost && usage.cost.total > 0) return usage.cost.total
+    const r = rates
+    if (!r || !hasPricing(r)) return 0
+    return (
+        (r.input / 1_000_000) * usage.input +
+        (r.output / 1_000_000) * usage.output +
+        (r.cacheRead / 1_000_000) * usage.cacheRead +
+        (r.cacheWrite / 1_000_000) * usage.cacheWrite
+    )
+}
+
+export function addTurnCost(totalUsd: number, usage: TokenUsage | undefined, rates?: ModelRates): number {
+    if (!usage) return totalUsd
+    return totalUsd + Math.max(0, turnCostUsd(usage, rates))
+}
+
+export function formatCostUsd(amount: number): string {
+    if (!Number.isFinite(amount) || amount <= 0) return '$0.00'
+    if (amount < 0.01) return `$${amount.toFixed(4)}`
+    if (amount < 1) return `$${amount.toFixed(3)}`
+    return `$${amount.toFixed(2)}`
+}
+
+export interface SpeedSample {
+    outputTokens: number
+    ms: number
+}
+
+export function averageTps(samples: SpeedSample[]): number | undefined {
+    let tokens = 0
+    let seconds = 0
+    for (const s of samples) {
+        if (s.ms <= 0) continue
+        tokens += s.outputTokens
+        seconds += s.ms / 1000
+    }
+    if (seconds <= 0) return undefined
+    return Math.round(tokens / seconds)
+}

package/src/timestamps/index.ts

@@ -0,0 +1,43 @@
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+import { Text } from '@earendil-works/pi-tui'
+
+import {
+    ENTRY_TYPE,
+    formatTimestampLine,
+    hasDuplicateTimestampEntry,
+    isToolCallMessage,
+    recentEntries,
+    type TimestampEntryData,
+    type TimestampKind
+} from './logic.js'
+
+export default function timestampsExtension(pi: ExtensionAPI): void {
+    let enabled = true
+
+    pi.on('message_end', async (event, ctx) => {
+        const message = event.message
+
+        if (enabled && (message.role === 'user' || message.role === 'assistant')) {
+            const timestamp = message.timestamp ?? Date.now()
+            const kind: TimestampKind =
+                message.role === 'user' ? 'user' : isToolCallMessage(message) ? 'tool' : 'assistant'
+            if (!hasDuplicateTimestampEntry(recentEntries(ctx.sessionManager), message.role, timestamp)) {
+                pi.appendEntry(ENTRY_TYPE, { role: message.role, kind, timestamp } satisfies TimestampEntryData)
+            }
+        }
+    })
+
+    pi.registerEntryRenderer(ENTRY_TYPE, (entry, _options, theme) => {
+        if (!enabled || !entry.data) return undefined
+        const data = entry.data as TimestampEntryData
+        return new Text(theme.fg('dim', formatTimestampLine(data, Date.now())), 0, 0)
+    })
+
+    pi.registerCommand('timestamps', {
+        description: 'Toggle message timestamps',
+        handler: async (_args, ctx) => {
+            enabled = !enabled
+            if (ctx.hasUI) ctx.ui.notify(`timestamps ${enabled ? 'enabled' : 'disabled'}`, 'info')
+        }
+    })
+}

package/src/timestamps/logic.ts

@@ -0,0 +1,79 @@
+import type { SessionEntry } from '@earendil-works/pi-coding-agent'
+
+export type TimestampKind = 'user' | 'assistant' | 'tool'
+export type TimestampEntryData = {
+    role: 'user' | 'assistant'
+    kind?: TimestampKind
+    timestamp: number
+}
+
+export function formatClock(ts: number): string {
+    const d = new Date(ts)
+    const pad = (n: number) => String(n).padStart(2, '0')
+    return `${pad(d.getHours())}:${pad(d.getMinutes())}:${pad(d.getSeconds())}`
+}
+
+export function formatRelative(ts: number, now: number): string {
+    const diff = Math.max(0, Math.floor((now - ts) / 1000))
+    if (diff < 5) return 'now'
+    if (diff < 60) return `${diff}s ago`
+    if (diff < 3600) return `${Math.floor(diff / 60)}m ago`
+    return `${Math.floor(diff / 3600)}h ago`
+}
+
+export function describeKind(kind: TimestampKind | undefined, role: TimestampEntryData['role']): string {
+    const resolved = kind ?? (role === 'user' ? 'user' : 'assistant')
+    return resolved === 'tool' ? 'tool call' : resolved === 'user' ? 'user message' : 'ai response'
+}
+
+export function formatTimestampLine(data: TimestampEntryData, now: number): string {
+    return `${formatClock(data.timestamp)} (${formatRelative(data.timestamp, now)}) · ${describeKind(data.kind, data.role)}`
+}
+
+export function isToolCallMessage(message: { stopReason?: string; content?: { type?: string }[] }): boolean {
+    return (
+        message.stopReason === 'toolUse' ||
+        (Array.isArray(message.content) && message.content.some((block) => block?.type === 'toolCall'))
+    )
+}
+
+const DUPLICATE_WINDOW_MS = 2_000
+
+export function hasDuplicateTimestampEntry(
+    entries: readonly unknown[],
+    role: 'user' | 'assistant',
+    timestamp: number
+): boolean {
+    const lookBack = 10
+    for (let i = entries.length - 1; i >= 0 && i >= entries.length - lookBack; i--) {
+        const entry = entries[i] as { type?: string; customType?: string; data?: TimestampEntryData } | undefined
+        if (!entry || entry.type !== 'custom' || entry.customType !== ENTRY_TYPE) continue
+        if (!entry.data || entry.data.role !== role) continue
+        if (Math.abs((entry.data.timestamp ?? Number.NaN) - timestamp) <= DUPLICATE_WINDOW_MS) return true
+    }
+    return false
+}
+
+export const ENTRY_TYPE = 'timestamps'
+
+/** The entries the duplicate check reads, walked back by parent id instead of copied whole. */
+export interface SessionWalk {
+    getLeafEntry?: () => SessionEntry | undefined
+    getEntry?: (id: string) => SessionEntry | undefined
+    getEntries?: () => SessionEntry[]
+}
+
+export function recentEntries(session: SessionWalk, limit = 10): SessionEntry[] {
+    const leaf = session.getLeafEntry?.()
+    if (leaf !== undefined && session.getEntry !== undefined) {
+        const out: SessionEntry[] = []
+        let entry: SessionEntry | undefined = leaf
+        while (entry !== undefined && out.length < limit) {
+            out.push(entry)
+            if (out.length === limit) break
+            entry = entry.parentId === null ? undefined : session.getEntry(entry.parentId)
+        }
+        return out
+    }
+    return (session.getEntries?.() ?? []).slice(-limit).toReversed()
+}