branch-usage

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { BranchView, PrLink, SessionFile } from '../types'
5import {
6  DEFAULT_PRICES,
7  addUsage,
8  aggregate,
9  buildPrices,
10  emptyView,
11  formatUsd,
12  isPrSyncCommand,
13  newSessionFile,
14  prCommentBody,
15  PR_COMMENT_MARKER,
16  relativeTime,
17  sharePct,
18  shortModel,
19  summaryText,
20  tokenLine,
21  usageDir,
22  type PriceTable,
23  type TokenUsage,
24} from './usage'
25
26const PANE = 'branch-usage'
27const TOOL = 'mcp__branch-usage__get_branch_usage'
28const REFRESH_MS = 10_000
29const MIN_REFRESH_GAP_MS = 2_000
30const PR_COLOR = 'blue'
31
32const view = atom({ plugin: 'branch-usage', key: 'view' } as const, emptyView('', null))
33const prLink = atom({ plugin: 'branch-usage', key: 'pr' } as const, null as PrLink | null)
34
35// Module variables start over on a hot reload; everything here is rebuilt
36// from the session files on disk.
37let prices: PriceTable = DEFAULT_PRICES
38let writes: Promise<unknown> = Promise.resolve()
39let lastRefresh = 0
40let prBranch: string | null = null
41const agentTypes = new Map<string, string>()
42const files = new Map<string, SessionFile>()
43
44async function git($: EngineInterface, ...args: string[]): Promise<string | null> {
45  const ran = await $.process.run(['git', ...args], { cwd: await $.session.root() })
46  return ran.exitCode === 0 ? ran.stdout.trim() : null
47}
48
49async function currentBranch($: EngineInterface): Promise<string | null> {
50  const name = await git($, 'rev-parse', '--abbrev-ref', 'HEAD')
51  if (name === null) return null
52  if (name !== 'HEAD') return name
53  const sha = await git($, 'rev-parse', '--short', 'HEAD')
54  return `detached@${sha ?? 'unknown'}`
55}
56
57async function readBranchFiles($: EngineInterface, root: string, branch: string): Promise<SessionFile[]> {
58  const dir = usageDir(root, branch)
59  if (!(await $.fs.exists(dir))) return []
60  const found: SessionFile[] = []
61  for (const entry of await $.fs.list(dir)) {
62    if (entry.kind !== 'file' || !entry.name.endsWith('.json')) continue
63    try {
64      found.push(JSON.parse(await $.fs.read(`${dir}/${entry.name}`)) as SessionFile)
65    } catch (err) {
66      $.ui.log(`branch-usage: skipped unreadable ${dir}/${entry.name}: ${String(err)}`)
67    }
68  }
69  return found
70}
71
72async function compute($: EngineInterface, branchName?: string): Promise<BranchView> {
73  const sessionId = await $.session.id()
74  const repo = await $.session.repo()
75  if (repo === null) return emptyView(sessionId, null)
76  const branch = branchName ?? (await currentBranch($))
77  if (branch === null) return emptyView(sessionId, null)
78  return aggregate(await readBranchFiles($, repo.root, branch), branch, sessionId)
79}
80
81async function refresh($: EngineInterface): Promise<void> {
82  lastRefresh = Date.now()
83  try {
84    const next = await compute($)
85    await update($, view, () => next)
86    if (next.branch !== prBranch) await refreshPr($)
87  } catch (err) {
88    $.ui.log(`branch-usage: refresh failed: ${String(err)}`)
89  }
90}
91
92async function agentLabel($: EngineInterface, agentId: string | undefined): Promise<string> {
93  if (agentId === undefined) return 'main'
94  if (!agentTypes.has(agentId)) {
95    for (const agent of await $.agent.list()) agentTypes.set(agent.id, agent.type)
96  }
97  return agentTypes.get(agentId) ?? 'subagent'
98}
99
100type Attached = { key: string; path: string; file: SessionFile; isNew: boolean }
101
102/** Finds this session's file on the current branch, creating it in memory when absent. */
103async function loadSessionFile($: EngineInterface): Promise<Attached | null> {
104  const repo = await $.session.repo()
105  const branch = await currentBranch($)
106  if (repo === null || branch === null) return null
107  const sessionId = await $.session.id()
108  const key = `${branch}/${sessionId}`
109  const path = `${usageDir(repo.root, branch)}/${sessionId}.json`
110  const known = files.get(key)
111  if (known !== undefined) return { key, path, file: known, isNew: false }
112  if (await $.fs.exists(path)) {
113    return { key, path, file: JSON.parse(await $.fs.read(path)) as SessionFile, isNew: false }
114  }
115  return { key, path, file: newSessionFile(sessionId, branch, Date.now()), isNew: true }
116}
117
118// Attaches the session to its branch: its file exists from the first turn on,
119// and again under the new branch after a checkout or a new id after /clear.
120async function attachNow($: EngineInterface): Promise<void> {
121  const attached = await loadSessionFile($)
122  if (attached === null) return
123  files.set(attached.key, attached.file)
124  if (attached.isNew) {
125    await $.fs.write(attached.path, JSON.stringify(attached.file, null, 2))
126    await refresh($)
127  }
128}
129
130async function recordNow(
131  $: EngineInterface,
132  agentId: string | undefined,
133  model: string,
134  usage: TokenUsage,
135): Promise<void> {
136  const attached = await loadSessionFile($)
137  if (attached === null) return
138  const now = Date.now()
139  const file = addUsage(attached.file, { agent: await agentLabel($, agentId), model, usage }, prices, now)
140  files.set(attached.key, file)
141  await $.fs.write(attached.path, JSON.stringify(file, null, 2))
142  if (now - lastRefresh >= MIN_REFRESH_GAP_MS) await refresh($)
143}
144
145async function gh($: EngineInterface, args: string[], stdin?: string): Promise<string> {
146  const ran = await $.process.run(['gh', ...args], { cwd: await $.session.root(), stdin })
147  if (ran.exitCode !== 0) {
148    throw new Error(`gh ${args.join(' ')} exited ${ran.exitCode}: ${ran.stderr.trim()}`)
149  }
150  return ran.stdout.trim()
151}
152
153/** The open PR of the current branch, or null when there is none or `gh` fails (logged). */
154async function lookupPr($: EngineInterface, branch: string): Promise<PrLink | null> {
155  try {
156    const json = await gh($, ['pr', 'view', '--json', 'number,url'])
157    const { number, url } = JSON.parse(json) as { number: number; url: string }
158    return { branch, number, url }
159  } catch (err) {
160    if (!String(err).includes('no pull requests found')) {
161      $.ui.log(`branch-usage: could not look up the PR of ${branch}: ${String(err)}`)
162    }
163    return null
164  }
165}
166
167// Sets prBranch first, so a failing lookup is not retried on every refresh tick.
168async function refreshPr($: EngineInterface): Promise<PrLink | null> {
169  const branch = await currentBranch($)
170  prBranch = branch
171  const found = branch === null ? null : await lookupPr($, branch)
172  await update($, prLink, () => found)
173  return found
174}
175
176// Creates the PR comment, or edits it when the marker comment already exists.
177async function syncPrCommentNow($: EngineInterface): Promise<void> {
178  const repo = await $.session.repo()
179  const branch = await currentBranch($)
180  if (repo === null || branch === null) return
181  const link = await refreshPr($)
182  if (link === null) return
183  const pr = link.number
184  const body = prCommentBody(aggregate(await readBranchFiles($, repo.root, branch), branch, ''))
185  const found = await gh($, [
186    'api',
187    `repos/{owner}/{repo}/issues/${pr}/comments`,
188    '--paginate',
189    '--jq',
190    `[.[] | select(.body | startswith("${PR_COMMENT_MARKER}")) | .id][0] // empty`,
191  ])
192  const id = found.split('\n')[0]
193  if (id === '') {
194    await gh($, ['api', '-X', 'POST', `repos/{owner}/{repo}/issues/${pr}/comments`, '-F', 'body=@-'], body)
195  } else {
196    await gh($, ['api', '-X', 'PATCH', `repos/{owner}/{repo}/issues/comments/${id}`, '-F', 'body=@-'], body)
197  }
198  $.ui.toast(`branch-usage: PR #${pr} comment ${id === '' ? 'created' : 'updated'}`)
199}
200
201// One write at a time, so parallel subagent steps never lose an update.
202function queue($: EngineInterface, job: () => Promise<void>): Promise<unknown> {
203  writes = writes.then(job).catch(err => $.ui.log(`branch-usage: could not write usage: ${String(err)}`))
204  return writes
205}
206
207function attach($: EngineInterface): Promise<unknown> {
208  return queue($, () => attachNow($))
209}
210
211function record(
212  $: EngineInterface,
213  agentId: string | undefined,
214  model: string,
215  usage: TokenUsage,
216): Promise<unknown> {
217  return queue($, () => recordNow($, agentId, model, usage))
218}
219
220export const register: Register = (on, options) => {
221  prices = buildPrices(String(options.prices ?? ''))
222
223  on('session.start', async ($, e, next) => {
224    await $.command.register({
225      name: 'branch-usage',
226      description: 'Show token usage and cost of all sessions on this branch',
227    })
228    await $.tool.register({
229      name: 'get_branch_usage',
230      description:
231        'Tokens (in/out/cache read/cache write) and $ cost per model, agent and session for all Claude Code sessions on a git branch. Defaults to the current branch.',
232      inputSchema: {
233        type: 'object',
234        properties: { branch: { type: 'string', description: 'Branch name; default is the current branch.' } },
235      },
236    })
237    await attach($)
238    await refresh($)
239    $.clock.every(REFRESH_MS, () => refresh($))
240    void $.ui.open({ id: PANE, title: 'Branch usage' })
241
242    return next(e)
243  })
244
245  // /clear starts a new session id without session.start; attach it now so
246  // the pane lists it before the first prompt.
247  on('classic.SessionStart', async ($, e, next) => {
248    if (e.source === 'clear') await attach($)
249
250    return next(e)
251  })
252
253  // A checkout moves the branch: each turn re-attaches to the current branch.
254  on('turn.start', async ($, e, next) => {
255    await attach($)
256
257    return next(e)
258  })
259
260  on('turn.step', async function* ($, e, next) {
261    const result = yield* next(e)
262    if (result.usage !== null) {
263      const { model, ...usage } = result.usage
264      await record($, e.agentId, model, usage)
265    }
266
267    return result
268  })
269
270  // After a PR is created or the branch is pushed, post or refresh the usage comment.
271  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
272    const ran = await next(e)
273    if (ran.deny === undefined && ran.isError !== true && isPrSyncCommand(e.command)) {
274      void queue($, () => syncPrCommentNow($))
275    }
276
277    return ran
278  })
279
280  on('session.measure', async ($, e, next) => {
281    const cost = e.cost
282    if (cost !== undefined) {
283      const sessionId = await $.session.id()
284      for (const [key, file] of files) {
285        if (key.endsWith(`/${sessionId}`)) files.set(key, { ...file, engineUsd: cost.usd })
286      }
287      await update($, view, v => ({ ...v, engineUsd: cost.usd }))
288    }
289
290    return next(e)
291  })
292
293  on('command.run', { command: 'branch-usage' }, async $ => {
294    await $.ui.open({ id: PANE, title: 'Branch usage', focus: true })
295    await refreshPr($)
296    await refresh($)
297
298    return { text: summaryText(await compute($)) }
299  })
300
301  on('tool.call', { tool: TOOL }, async ($, e) => {
302    const asked = (e as { branch?: unknown }).branch
303    const branch = typeof asked === 'string' && asked !== '' ? asked : undefined
304
305    return { result: summaryText(await compute($, branch)) }
306  })
307
308  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
309    const { branch, totalUsd } = await read($, view)
310    if (e.props.hasSurvey || branch === null) return next(e)
311
312    const { Box, Link, Text } = $.ui.resolve(e)
313    const link = await read($, prLink)
314
315    return (
316      <Box>
317        <Text dimColor>
318          {formatUsd(totalUsd)} · {branch}
319          {link?.branch === branch ? ' ' : ''}
320        </Text>
321        {link?.branch === branch && (
322          <Link href={link.url}>
323            <Text color={PR_COLOR} underline>(#{link.number})</Text>
324          </Link>
325        )}
326      </Box>
327    )
328  })
329
330  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
331    const { Box, Link, Text } = $.ui.resolve(e)
332    const { branch, sessionId, sessions, byModel, byAgent, totalUsd, engineUsd, error } = await read($, view)
333    const link = await read($, prLink)
334    const rows = Math.max(1, (e.viewport?.rows ?? 30) - 4)
335
336    if (error !== undefined) return <Text color="red">{error}</Text>
337    if (branch === null) return <Text dimColor>Not inside a git repository.</Text>
338
339    const now = Date.now()
340    // Label shrinks and truncates; the right part keeps its width, so the
341    // cost column stays on the right edge.
342    const line = (label: string, usd: number, share: string | null, props: { bold?: boolean; color?: string } = {}) => (
343      <Box flexDirection="row" justifyContent="space-between">
344        <Box flexShrink={1}>
345          <Text wrap="truncate-end" {...props}>{label}</Text>
346        </Box>
347        <Box flexShrink={0} marginLeft={1}>
348          {share !== null && <Text dimColor>{share} </Text>}
349          <Text {...props}>{formatUsd(usd)}</Text>
350        </Box>
351      </Box>
352    )
353    // The PR link sits right of the branch name and outside the truncating
354    // label, so a long branch name never cuts it off.
355    const header = (
356      <Box flexDirection="row" justifyContent="space-between">
357        <Box flexShrink={1}>
358          <Box flexShrink={1}>
359            <Text wrap="truncate-end" bold>{branch}</Text>
360          </Box>
361          {link?.branch === branch && (
362            <Box flexShrink={0}>
363              <Text> </Text>
364              <Link href={link.url}>
365                <Text color={PR_COLOR} underline>(#{link.number})</Text>
366              </Link>
367            </Box>
368          )}
369        </Box>
370        <Box flexShrink={0} marginLeft={1}>
371          <Text bold>{formatUsd(totalUsd)}</Text>
372        </Box>
373      </Box>
374    )
375
376    return (
377      <Box flexDirection="column">
378        {header}
379        <Text dimColor>
380          {'  '}{sessions.length} session{sessions.length === 1 ? '' : 's'}
381        </Text>
382        {engineUsd !== undefined && <Text dimColor>  this session, engine total {formatUsd(engineUsd)}</Text>}
383        <Box marginTop={1}>
384          <Text bold>By model</Text>
385        </Box>
386        {byModel.map(r => (
387          <Box flexDirection="column">
388            {line(`  - ${shortModel(r.model)}${r.isEstimated ? ' ~' : ''}`, r.usd, sharePct(r.usd, totalUsd))}
389            <Text dimColor wrap="truncate-end">    {tokenLine(r)}</Text>
390          </Box>
391        ))}
392        <Box marginTop={1}>
393          <Text bold>By agent</Text>
394        </Box>
395        {byAgent.map(r => (
396          <Box flexDirection="column">
397            {line(`  - ${r.agent} · ${r.requests} req`, r.usd, sharePct(r.usd, totalUsd))}
398            <Text dimColor wrap="truncate-end">    {tokenLine(r)}</Text>
399          </Box>
400        ))}
401        <Box marginTop={1}>
402          <Text bold>By session</Text>
403        </Box>
404        {sessions.slice(0, rows).map(s => {
405          const active = s.sessionId === sessionId
406          return (
407            line(
408              `  - ${active ? '● ' : '  '}${s.sessionId.slice(0, 8)}  ${relativeTime(s.updatedAt, now)}`,
409              s.usd,
410              null,
411              active ? { bold: true, color: 'cyan' } : {},
412            )
413          )
414        })}
415      </Box>
416    )
417  })
418}
419

1import type { BranchView, SessionFile, SessionSummary, UsageRow } from '../types'
2
3export type Price = { in: number; out: number; cacheRead: number; cacheWrite: number }
4export type PriceTable = Record<string, Price>
5export type TokenUsage = {
6  input_tokens: number
7  output_tokens: number
8  cache_read_input_tokens: number
9  cache_creation_input_tokens: number
10}
11
12// $ per million tokens. Cache write = 1.25x input and cache read = 0.1x input
13// unless the model's row says otherwise.
14const tier = (input: number, output: number, cacheRead = input * 0.1): Price => ({
15  in: input,
16  out: output,
17  cacheRead,
18  cacheWrite: input * 1.25,
19})
20
21export const DEFAULT_PRICES: PriceTable = {
22  'claude-fable': tier(10, 50, 0.25),
23  'claude-mythos': tier(10, 50, 0.25),
24  'claude-opus-5-5': tier(4, 20, 0.2),
25  'claude-opus-5': tier(5, 25),
26  'claude-opus-4-1': tier(15, 75),
27  'claude-opus-4-2': tier(15, 75),
28  'claude-opus-4': tier(5, 25),
29  'claude-sonnet-5': tier(2, 10, 0.2),
30  'claude-sonnet-4': tier(3, 15),
31  'claude-sonnet-3': tier(3, 15),
32  'claude-haiku-4': tier(1, 5),
33  'claude-haiku-3-5': tier(0.8, 4),
34  'claude-haiku-3': tier(0.25, 1.25),
35}
36
37const FALLBACK_PRICE: Price = tier(3, 15)
38
39/** Parses the `prices` option over the defaults; throws with the reason on bad JSON. */
40export const buildPrices = (raw: string): PriceTable => {
41  if (raw.trim() === '') return DEFAULT_PRICES
42  let parsed: unknown
43  try {
44    parsed = JSON.parse(raw)
45  } catch (err) {
46    throw new Error(`branch-usage: option "prices" is not valid JSON (${String(err)}); fix it in /config`)
47  }
48  if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
49    throw new Error('branch-usage: option "prices" must be an object keyed by model id prefix')
50  }
51  const table: PriceTable = { ...DEFAULT_PRICES }
52  for (const [prefix, value] of Object.entries(parsed)) {
53    const v = value as Partial<Price>
54    if (typeof v?.in !== 'number' || typeof v?.out !== 'number') {
55      throw new Error(`branch-usage: prices["${prefix}"] needs numeric "in" and "out"`)
56    }
57    table[prefix] = {
58      in: v.in,
59      out: v.out,
60      cacheRead: v.cacheRead ?? v.in * 0.1,
61      cacheWrite: v.cacheWrite ?? v.in * 1.25,
62    }
63  }
64  return table
65}
66
67/** Longest matching prefix wins; no match gives the fallback tier, flagged. */
68export const priceFor = (model: string, table: PriceTable): { price: Price; isKnown: boolean } => {
69  let best: string | undefined
70  for (const prefix of Object.keys(table)) {
71    if (model.startsWith(prefix) && (best === undefined || prefix.length > best.length)) best = prefix
72  }
73  return best === undefined
74    ? { price: FALLBACK_PRICE, isKnown: false }
75    : { price: table[best], isKnown: true }
76}
77
78export const costUsd = (u: TokenUsage, p: Price): number =>
79  (u.input_tokens * p.in +
80    u.output_tokens * p.out +
81    u.cache_read_input_tokens * p.cacheRead +
82    u.cache_creation_input_tokens * p.cacheWrite) /
83  1e6
84
85export const newSessionFile = (sessionId: string, branch: string, now: number): SessionFile => ({
86  version: 1,
87  sessionId,
88  branch,
89  startedAt: now,
90  updatedAt: now,
91  rows: {},
92})
93
94/** Returns a copy of `file` with one model response added. */
95export const addUsage = (
96  file: SessionFile,
97  entry: { agent: string; model: string; usage: TokenUsage },
98  table: PriceTable,
99  now: number,
100): SessionFile => {
101  const key = `${entry.agent}|${entry.model}`
102  const { price, isKnown } = priceFor(entry.model, table)
103  const prev = file.rows[key]
104  const row: UsageRow = {
105    model: entry.model,
106    agent: entry.agent,
107    requests: (prev?.requests ?? 0) + 1,
108    input: (prev?.input ?? 0) + entry.usage.input_tokens,
109    output: (prev?.output ?? 0) + entry.usage.output_tokens,
110    cacheRead: (prev?.cacheRead ?? 0) + entry.usage.cache_read_input_tokens,
111    cacheWrite: (prev?.cacheWrite ?? 0) + entry.usage.cache_creation_input_tokens,
112    usd: (prev?.usd ?? 0) + costUsd(entry.usage, price),
113    isEstimated: (prev?.isEstimated ?? false) || !isKnown,
114  }
115  return { ...file, updatedAt: now, rows: { ...file.rows, [key]: row } }
116}
117
118const merge = (into: Map<string, UsageRow>, key: string, row: UsageRow, label: Partial<UsageRow>) => {
119  const prev = into.get(key)
120  into.set(key, {
121    model: label.model ?? row.model,
122    agent: label.agent ?? row.agent,
123    requests: (prev?.requests ?? 0) + row.requests,
124    input: (prev?.input ?? 0) + row.input,
125    output: (prev?.output ?? 0) + row.output,
126    cacheRead: (prev?.cacheRead ?? 0) + row.cacheRead,
127    cacheWrite: (prev?.cacheWrite ?? 0) + row.cacheWrite,
128    usd: (prev?.usd ?? 0) + row.usd,
129    isEstimated: (prev?.isEstimated ?? false) || row.isEstimated,
130  })
131}
132
133const byUsdDesc = (a: UsageRow, b: UsageRow) => b.usd - a.usd
134
135/** Folds every session file of a branch into the view the pane draws. */
136export const aggregate = (
137  files: readonly SessionFile[],
138  branch: string,
139  sessionId: string,
140): BranchView => {
141  const models = new Map<string, UsageRow>()
142  const agents = new Map<string, UsageRow>()
143  const sessions: SessionSummary[] = []
144  let totalUsd = 0
145  for (const file of files) {
146    let sessionUsd = 0
147    for (const row of Object.values(file.rows)) {
148      merge(models, row.model, row, { agent: '' })
149      merge(agents, row.agent, row, { model: '' })
150      sessionUsd += row.usd
151    }
152    totalUsd += sessionUsd
153    sessions.push({ sessionId: file.sessionId, updatedAt: file.updatedAt, usd: sessionUsd })
154  }
155  sessions.sort((a, b) => b.updatedAt - a.updatedAt)
156  return {
157    branch,
158    sessionId,
159    sessions,
160    byModel: [...models.values()].sort(byUsdDesc),
161    byAgent: [...agents.values()].sort(byUsdDesc),
162    totalUsd,
163    engineUsd: files.find(f => f.sessionId === sessionId)?.engineUsd,
164  }
165}
166
167export const emptyView = (sessionId: string, branch: string | null, error?: string): BranchView => ({
168  branch,
169  sessionId,
170  sessions: [],
171  byModel: [],
172  byAgent: [],
173  totalUsd: 0,
174  error,
175})
176
177/** Directory name for a branch: `feat/x` becomes `feat%2Fx`. */
178export const branchDirName = (branch: string): string => encodeURIComponent(branch)
179
180export const usageDir = (repoRoot: string, branch: string): string =>
181  `${repoRoot}/.claude/branch-usage/${branchDirName(branch)}`
182
183export const formatTokens = (n: number): string =>
184  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${(n / 1e3).toFixed(1)}k` : String(n)
185
186export const formatUsd = (n: number): string => `$${n < 10 ? n.toFixed(3) : n.toFixed(2)}`
187
188/** Pane label for a model id: `claude-sonnet-5-5` becomes `sonnet-5-5`. */
189export const shortModel = (model: string): string => model.replace(/^claude-/, '')
190
191export const relativeTime = (ms: number, now: number): string => {
192  const mins = Math.floor((now - ms) / 60_000)
193  if (mins < 1) return 'just now'
194  if (mins < 60) return `${mins}m ago`
195  const hours = Math.floor(mins / 60)
196  return hours < 24 ? `${hours}h ago` : `${Math.floor(hours / 24)}d ago`
197}
198
199export const sharePct = (usd: number, total: number): string =>
200  total === 0 ? '0%' : `${Math.round((usd / total) * 100)}%`
201
202export const tokenLine = (r: UsageRow): string =>
203  `in ${formatTokens(r.input)} out ${formatTokens(r.output)} cr ${formatTokens(r.cacheRead)} cw ${formatTokens(r.cacheWrite)}`
204
205const TOKEN_HEAD = '| Req | In | Out | Cache read | Cache write | Cost |'
206const TOKEN_ALIGN = '| ---: | ---: | ---: | ---: | ---: | ---: |'
207
208const tokenCells = (r: UsageRow): string =>
209  `${r.requests} | ${formatTokens(r.input)} | ${formatTokens(r.output)} | ${formatTokens(r.cacheRead)} | ${formatTokens(r.cacheWrite)} | ${formatUsd(r.usd)}`
210
211const usageTable = (label: string, rows: readonly UsageRow[], name: (r: UsageRow) => string): string[] => [
212  `| ${label} ${TOKEN_HEAD}`,
213  `| --- ${TOKEN_ALIGN}`,
214  ...rows.map(r => `| ${name(r)} | ${tokenCells(r)} |`),
215]
216
217/** First line of the PR comment; marks the one comment the mod owns. */
218export const PR_COMMENT_MARKER = '<!-- branch-usage -->'
219
220// `gh pr create` or `git [opts] push` at the start of a command or after && ; | newline.
221const PR_SYNC_COMMAND = /(^|[;&|\n]\s*)(gh\s+pr\s+create|git(\s+-\S+(\s+[^\s-]\S*)?)*\s+push)(\s|$)/
222
223/** True when a Bash command creates a PR or pushes a branch. */
224export const isPrSyncCommand = (command: string): boolean => PR_SYNC_COMMAND.test(command)
225
226/** PR comment text: the marker, then the report without the "this session" mark. */
227export const prCommentBody = (view: BranchView): string =>
228  `${PR_COMMENT_MARKER}\n${summaryText({ ...view, sessionId: '' })}`
229
230/** Markdown summary for the command and the model tool. */
231export const summaryText = (view: BranchView): string => {
232  if (view.branch === null) return 'branch-usage: not inside a git repository, nothing is tracked.'
233  const lines = [
234    `## Branch \`${view.branch}\``,
235    '',
236    `**${formatUsd(view.totalUsd)}** over ${view.sessions.length} session${view.sessions.length === 1 ? '' : 's'}`,
237    '',
238    '### By model',
239    '',
240    ...usageTable('Model', view.byModel, r => `\`${r.model}\`${r.isEstimated ? ' (~price)' : ''}`),
241    '',
242    '### By agent',
243    '',
244    ...usageTable('Agent', view.byAgent, r => r.agent),
245    '',
246    '### By session',
247    '',
248    '| Session | Updated | Cost |',
249    '| --- | --- | ---: |',
250    ...view.sessions.map(
251      s =>
252        `| \`${s.sessionId.slice(0, 8)}\`${s.sessionId === view.sessionId ? ' ●' : ''} | ${new Date(s.updatedAt).toISOString()} | ${formatUsd(s.usd)} |`,
253    ),
254  ]
255  return lines.join('\n')
256}
257

1export type UsageRow = {
2  model: string
3  /** 'main' or the subagent type. */
4  agent: string
5  requests: number
6  input: number
7  output: number
8  cacheRead: number
9  cacheWrite: number
10  usd: number
11  /** True when no price matched the model and the default tier was used. */
12  isEstimated: boolean
13}
14
15export type SessionFile = {
16  version: 1
17  sessionId: string
18  branch: string
19  startedAt: number
20  updatedAt: number
21  /** The engine's own session total at the last measurement. */
22  engineUsd?: number
23  /** Keyed `${agent}|${model}`. */
24  rows: Record<string, UsageRow>
25}
26
27export type SessionSummary = {
28  sessionId: string
29  updatedAt: number
30  usd: number
31}
32
33export type BranchView = {
34  /** Null when the directory is not in a git repository. */
35  branch: string | null
36  sessionId: string
37  sessions: SessionSummary[]
38  byModel: UsageRow[]
39  byAgent: UsageRow[]
40  totalUsd: number
41  /** The engine's total for this session, to compare with ours. */
42  engineUsd?: number
43  error?: string
44}
45
46export type PrLink = {
47  branch: string
48  number: number
49  url: string
50}
51
52declare module 'claude-code' {
53  interface PluginState {
54    'branch-usage': { view: BranchView; pr: PrLink | null }
55  }
56}
57