@tduyng/pi-cost 0.0.1 → 0.0.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tduyng/pi-cost",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "Session spend, live provider rate-limit windows, footer segments, and per-message timestamps.",
   "keywords": [
     "pi",

package/src/cost/index.ts

@@ -86,9 +86,12 @@ function errorText(cause: unknown): string {
     return cause instanceof Error ? cause.message : String(cause)
 }
 
-export async function claudeToken(pi: ExtensionAPI): Promise<string | undefined> {
+export async function claudeToken(
+    pi: ExtensionAPI,
+    platform: NodeJS.Platform = process.platform
+): Promise<string | undefined> {
     let raw: string | undefined
-    if (process.platform === 'darwin') {
+    if (platform === 'darwin') {
         const result = await pi.exec('security', ['find-generic-password', '-s', 'Claude Code-credentials', '-w'], {
             cwd: os.homedir(),
             timeout: KEYCHAIN_TIMEOUT_MS
@@ -166,11 +169,18 @@ function modelRates(ctx: ExtensionCommandContext): ModelRates | undefined {
         : undefined
 }
 
-function sessionSpend(ctx: ExtensionCommandContext): SessionSpend {
+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 rates = modelRates(ctx)
+    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) {
@@ -181,7 +191,7 @@ function sessionSpend(ctx: ExtensionCommandContext): SessionSpend {
                     cacheRead: collapsed.cacheRead + usage.cacheRead,
                     cacheWrite: collapsed.cacheWrite + usage.cacheWrite
                 }
-                costUsd = addTurnCost(costUsd, usage, rates)
+                costUsd = addTurnCost(costUsd, usage, currentRates)
             }
             continue
         }
@@ -191,7 +201,7 @@ function sessionSpend(ctx: ExtensionCommandContext): SessionSpend {
             continue
         }
         const usage = entry.message.usage
-        costUsd = addTurnCost(costUsd, usage, rates)
+        costUsd = addTurnCost(costUsd, usage, ratesFor(entry.message.provider, entry.message.model) ?? currentRates)
         turns.push({ usage, timestamp: entry.timestamp })
     }
     const totals = sessionTotals(turns)
@@ -199,17 +209,14 @@ function sessionSpend(ctx: ExtensionCommandContext): SessionSpend {
     totals.output += collapsed.output
     totals.cacheRead += collapsed.cacheRead
     totals.cacheWrite += collapsed.cacheWrite
-    const overhead = sumOverhead(readOverhead(getAgentDir(), ctx.sessionManager.getSessionId()))
-    return { costUsd, totals, overhead }
+    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(', ')
+    const extensions = [...session.overhead.byExtension].map(([name, spend]) => `${name} ${spend.calls}`).join(', ')
     return [
         `  ${label('Cost')}${formatCostUsd(session.costUsd)}`,
         ...(session.overhead.calls === 0
@@ -370,6 +377,7 @@ export default function costExtension(pi: ExtensionAPI): void {
                 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
@@ -387,7 +395,7 @@ export default function costExtension(pi: ExtensionAPI): void {
                 }
 
                 render = () => {
-                    const session = sessionSpend(ctx)
+                    const session = sessionSpend(ctx, overheadRecords)
                     lines = [
                         `${theme.fg('accent', theme.bold('Cost and limits'))}  ${theme.fg('dim', header)}`,
                         '',

package/src/statusline/index.ts

@@ -57,6 +57,7 @@ export default function statuslineExtension(pi: ExtensionAPI): void {
     })
 
     pi.on('agent_settled', () => {
+        if (!enabled) return
         void refreshGit()
     })
 
@@ -73,7 +74,7 @@ export default function statuslineExtension(pi: ExtensionAPI): void {
 
     async function refreshGit(): Promise<void> {
         const ctx = mountedCtx
-        if (!ctx) return
+        if (!ctx || !enabled) 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
@@ -87,7 +88,7 @@ export default function statuslineExtension(pi: ExtensionAPI): void {
 
     async function refreshPr(): Promise<void> {
         const ctx = mountedCtx
-        if (!ctx || !showPr) return
+        if (!ctx || !showPr || !enabled) 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)
@@ -122,7 +123,7 @@ export default function statuslineExtension(pi: ExtensionAPI): void {
     }
 
     function mount(ctx: ExtensionContext): void {
-        if (!ctx.hasUI) return
+        if (!ctx.hasUI || !enabled) return
         mountedCtx = ctx
 
         void refreshGit()

package/src/timestamps/logic.ts

@@ -37,8 +37,6 @@ export function isToolCallMessage(message: { stopReason?: string; content?: { ty
     )
 }
 
-const DUPLICATE_WINDOW_MS = 2_000
-
 export function hasDuplicateTimestampEntry(
     entries: readonly unknown[],
     role: 'user' | 'assistant',
@@ -49,7 +47,7 @@ export function hasDuplicateTimestampEntry(
         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
+        if (entry.data.timestamp === timestamp) return true
     }
     return false
 }