cc-cost

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register, SessionMessage, SessionRateLimit } from 'claude-code'
3
4import type { Gasto, Limite, Modo, Tokens } from '../types'
5
6const gasto = atom({ plugin: 'cc-cost', key: 'tokens' } as const, null)
7const compactando = atom({ plugin: 'cc-cost', key: 'compactando' } as const, null)
8
9const BARRAS = 12
10const DIAS_GUARDADOS = 7
11const NIVEIS = '▁▂▃▄▅▆▇█'
12const ROTULO: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'limite' }
13const ZERO: Tokens = { entrada: 0, saida: 0, cacheLido: 0, cacheEscrito: 0 }
14const VAZIO: Gasto = { sessao: ZERO, hoje: 0, proprio: { chave: '', n: 0 }, turnos: [], emTurno: false, limites: [], contexto: null, janela: null }
15
16// Folga até a janela de compactação: o vermelho deixa ~30k para compactar antes do autocompact.
17const FOLGA_ATENCAO = 60_000
18const FOLGA_TROCAR = 30_000
19// Chaves que o mod não usa mais: as da versão antiga e o recap do antigo botão "nova sessão".
20const CHAVES_ANTIGAS = ['dia:', 'tokens:', 'ultimoRecap']
21
22// A partir dessa parte do contexto em leituras de arquivo, o micro compensa regravar o cache inteiro.
23const PARTE_MICRO = 0.4
24// Aproximação para comparar partes do contexto, não para medir: a medida real vem do engine.
25const CHARS_POR_TOKEN = 4
26// O marcador que o micro-compaction deixa no lugar do arquivo: o que já foi tirado não conta de novo.
27const ELIDIDO = / elided by compaction: /
28
29const FOCO_COMPACTACAO =
30  'Priorize o que é preciso para continuar o trabalho: objetivo, arquivos alterados (com caminho), ' +
31  'decisões e o porquê, estado atual (testes, o que está quebrado), próximos passos e pendências em aberto.'
32
33/** Data local AAAA-MM-DD; em ordem alfabética é ordem cronológica. */
34function dia(ms: number): string {
35  const d = new Date(ms)
36  const p = (n: number) => String(n).padStart(2, '0')
37
38  return `${d.getFullYear()}-${p(d.getMonth() + 1)}-${p(d.getDate())}`
39}
40
41/** Leitura de cache fica de fora: é o mesmo contexto reenviado a cada requisição, e somá-la conta a conversa várias vezes. */
42function novos(t: Tokens): number {
43  return t.entrada + t.cacheEscrito + t.saida
44}
45
46/** 950 · 12,4k · 312k · 1,2M — uma casa decimal só abaixo de 10. */
47function compacto(n: number): string {
48  const [valor, sufixo] = n >= 1e6 ? [n / 1e6, 'M'] : n >= 1e3 ? [n / 1e3, 'k'] : [n, '']
49  const casas = sufixo && valor < 10 ? 1 : 0
50
51  return valor.toFixed(casas).replace('.', ',') + sufixo
52}
53
54function sparkline(valores: number[]): string {
55  const maior = Math.max(1, ...valores)
56
57  return valores.map(v => NIVEIS[Math.min(NIVEIS.length - 1, Math.floor((v / maior) * NIVEIS.length))]).join('')
58}
59
60function cor(pct: number): string {
61  if (pct >= 85) return 'red'
62  if (pct >= 60) return 'yellow'
63
64  return 'green'
65}
66
67function limiares(janela: number | null): { atencao: number; trocar: number } {
68  const total = janela ?? 200_000
69
70  return { atencao: total - FOLGA_ATENCAO, trocar: total - FOLGA_TROCAR }
71}
72
73function corContexto(tokens: number, janela: number | null): string {
74  const { atencao, trocar } = limiares(janela)
75  if (tokens >= trocar) return 'red'
76  if (tokens >= atencao) return 'yellow'
77
78  return 'green'
79}
80
81/** A janela de compactação (autoCompactWindow), não a do modelo; o breakdown `summary` é estimado localmente, sem requisição. */
82async function janelaCompactacao($: EngineInterface): Promise<number> {
83  const { context } = await $.session.usage({ breakdown: 'summary' })
84
85  return context.breakdown?.rawMaxTokens ?? context.window
86}
87
88function limites(janelas: readonly SessionRateLimit[]): Limite[] {
89  return janelas.map(({ kind, percentUsed, resetsAt }) => ({ kind, percentUsed, ...(resetsAt ? { resetsAt } : {}) }))
90}
91
92/**
93 * Cada sessão grava só a própria chave `novos:<dia>:<sessionId>`: duas sessões abertas
94 * não apagam a soma uma da outra, e o resume da mesma sessão continua de onde parou.
95 */
96async function somaDoDia($: EngineInterface, hoje: string, proprio: Gasto['proprio']): Promise<number> {
97  const prefixo = `novos:${hoje}:`
98  let soma = proprio.chave.startsWith(prefixo) ? proprio.n : 0
99  for (const chave of await $.store.keys()) {
100    if (chave.startsWith(prefixo) && chave !== proprio.chave) soma += Number((await $.store.get(chave)) ?? 0)
101  }
102
103  return soma
104}
105
106async function podar($: EngineInterface, agora: number) {
107  const corte = `novos:${dia(agora - DIAS_GUARDADOS * 86_400_000)}`
108  for (const chave of await $.store.keys()) {
109    const velha = chave.startsWith('novos:') ? chave < corte : CHAVES_ANTIGAS.some(p => chave.startsWith(p))
110    if (velha) await $.store.delete(chave)
111  }
112}
113
114// As escritas passam uma de cada vez: ler o store, somar e gravar não pode intercalar
115// entre passos paralelos (subagentes), senão uma soma apaga a outra.
116let fila: Promise<unknown> = Promise.resolve()
117function enfileirar(tarefa: () => Promise<void>): Promise<void> {
118  const vez = fila.then(tarefa)
119  fila = vez.catch(() => undefined)
120
121  return vez
122}
123
124async function somar($: EngineInterface, usage: ModelUsage) {
125  const t: Tokens = {
126    entrada: usage.input_tokens,
127    saida: usage.output_tokens,
128    cacheLido: usage.cache_read_input_tokens,
129    cacheEscrito: usage.cache_creation_input_tokens,
130  }
131  const n = novos(t)
132  if (n === 0 && t.cacheLido === 0) return
133
134  const hoje = dia(await $.clock.now())
135  const chave = `novos:${hoje}:${await $.session.id()}`
136  const atual = (await read($, gasto)) ?? VAZIO
137  const antes = atual.proprio.chave === chave ? atual.proprio.n : Number((await $.store.get(chave)) ?? 0)
138  const proprio = { chave, n: antes + n }
139  await $.store.set(chave, proprio.n)
140  const totalHoje = await somaDoDia($, hoje, proprio)
141
142  await update($, gasto, a => {
143    const g = a ?? VAZIO
144    // Subagentes do turno caem na barra dele; compactação e trabalho fora de turno só no total.
145    const turnos = g.emTurno && g.turnos.length > 0
146      ? [...g.turnos.slice(0, -1), (g.turnos.at(-1) ?? 0) + n]
147      : g.turnos
148
149    return {
150      ...g,
151      sessao: {
152        entrada: g.sessao.entrada + t.entrada,
153        saida: g.sessao.saida + t.saida,
154        cacheLido: g.sessao.cacheLido + t.cacheLido,
155        cacheEscrito: g.sessao.cacheEscrito + t.cacheEscrito,
156      },
157      hoje: totalHoje,
158      proprio,
159      turnos,
160    }
161  })
162}
163
164/** Tokens estimados das leituras de arquivo (Read) ainda inteiras: o que o /compact micro tira. */
165function leituras(msgs: readonly SessionMessage[]): number {
166  const lidos = new Set(msgs.flatMap(m => m.toolUses.filter(u => u.tool === 'Read').map(u => u.tool_use_id)))
167  let chars = 0
168  for (const m of msgs) {
169    for (const r of m.toolResults ?? []) {
170      if (lidos.has(r.tool_use_id) && !r.isError && !ELIDIDO.test(r.text)) chars += r.text.length
171    }
172  }
173
174  return Math.round(chars / CHARS_POR_TOKEN)
175}
176
177/** Tokens estimados do texto das mensagens; serve para a diferença entre antes e depois de compactar. */
178function tamanho(msgs: readonly SessionMessage[]): number {
179  let chars = 0
180  for (const m of msgs) chars += m.text.length + (m.toolResults ?? []).reduce((s, r) => s + r.text.length, 0)
181
182  return Math.round(chars / CHARS_POR_TOKEN)
183}
184
185/** Micro quando as leituras pesam no contexto; senão o resumo, que encolhe a conversa em si. */
186function modo(lidos: number, contexto: number | null): Modo {
187  return contexto !== null && lidos >= contexto * PARTE_MICRO ? 'micro' : 'resumo'
188}
189
190/**
191 * Compacta do jeito que mais rende agora. O micro (plugin micro-compaction) vai como /compact digitado,
192 * o único gatilho que o plugin atende. O resumo vai direto, e o ctx recebe aqui o tamanho que ele deixou,
193 * porque o hook session.compact deste mod não roda na chamada que o próprio mod faz.
194 */
195async function compactar($: EngineInterface) {
196  if ((await read($, compactando)) !== null) return
197  const escolhido = modo(leituras(await $.session.messages()), (await read($, gasto))?.contexto ?? null)
198  await update($, compactando, () => escolhido)
199  try {
200    if (escolhido === 'micro') {
201      await $.command.run({ command: 'compact', args: 'micro' })
202      return
203    }
204    const r = await $.session.compact({ instructions: FOCO_COMPACTACAO })
205    // Vetada por um hook: o próprio engine mostra a linha com o motivo.
206    if (r.skip !== undefined) return
207    const depois = r.tokensAfter ?? null
208    await enfileirar(() => update($, gasto, a => (a ? { ...a, contexto: depois } : a)).then(() => undefined))
209  } catch (erro) {
210    $.ui.toast(`Não consegui compactar (${erro instanceof Error ? erro.message : String(erro)}); a sessão continua como está.`)
211  } finally {
212    await update($, compactando, () => null)
213  }
214}
215
216export const register: Register = on => {
217  on('session.start', async ($, e, next) => {
218    await enfileirar(async () => {
219      const agora = await $.clock.now()
220      await podar($, agora)
221      const atual = (await read($, gasto)) ?? VAZIO
222      const hoje = await somaDoDia($, dia(agora), atual.proprio)
223      const { rateLimits, context } = await $.session.usage()
224      const janela = await janelaCompactacao($)
225      await update($, gasto, a => ({
226        ...(a ?? VAZIO),
227        hoje,
228        limites: rateLimits.length > 0 ? limites(rateLimits) : (a?.limites ?? []),
229        contexto: context.tokens ?? a?.contexto ?? null,
230        janela,
231      }))
232    })
233
234    return next(e)
235  })
236
237  // Só a thread principal abre turno: cada turno é uma barra.
238  on('turn.start', async ($, e, next) => {
239    await enfileirar(() =>
240      update($, gasto, a => {
241        const g = a ?? VAZIO
242
243        return { ...g, turnos: [...g.turnos, 0].slice(-BARRAS), emTurno: true }
244      }).then(() => undefined),
245    )
246
247    return next(e)
248  })
249
250  on('turn.complete', async ($, e, next) => {
251    if (e.agentId === undefined) {
252      await enfileirar(() => update($, gasto, a => (a ? { ...a, emTurno: false } : a)).then(() => undefined))
253    }
254
255    return next(e)
256  })
257
258  // Cada requisição ao modelo: thread principal, subagente ou processo do engine (compactação).
259  on('turn.step', async function* ($, e, next) {
260    const resposta = yield* next(e)
261    const usage = resposta.usage
262    if (usage) await enfileirar(() => somar($, usage))
263
264    return resposta
265  })
266
267  on('session.measure', async ($, e, next) => {
268    await enfileirar(async () => {
269      const atual = (await read($, gasto)) ?? VAZIO
270      const hoje = await somaDoDia($, dia(await $.clock.now()), atual.proprio)
271      const contexto = e.context.tokens ?? atual.contexto
272      await update($, gasto, a => ({
273        ...(a ?? VAZIO),
274        hoje,
275        limites: e.rateLimits.length > 0 ? limites(e.rateLimits) : (a?.limites ?? []),
276        contexto: e.context.tokens ?? a?.contexto ?? null,
277      }))
278
279      // Só na subida pelo limiar: depois do /clear ou de uma compactação o contexto cai e o aviso volta a valer.
280      const { trocar } = limiares(atual.janela)
281      if (contexto !== null && contexto >= trocar && (atual.contexto === null || atual.contexto < trocar)) {
282        const de = atual.janela ? ` de ${compacto(atual.janela)}` : ''
283        const lidos = leituras(await $.session.messages())
284        const faz = modo(lidos, contexto) === 'micro' ? `tira ~${compacto(lidos)} de leituras de arquivo (micro)` : 'resume a conversa'
285        $.ui.toast(`Contexto em ${compacto(contexto)}${de}: "compactar" na faixa (c) ${faz}, antes do autocompact.`)
286      }
287    })
288
289    return next(e)
290  })
291
292  // Depois da compactação o engine só mede o contexto na próxima resposta: até lá, o tamanho que ela deixou.
293  // O micro não informa esse tamanho: vale a estimativa do que saiu das mensagens.
294  on('session.compact', async ($, e, next) => {
295    const r = await next(e)
296    if (e.trigger !== 'precompute' && e.agentId === undefined && r.messages) {
297      const informado = r.tokensAfter
298      const saiu = tamanho(e.messages) - tamanho(r.messages)
299      await enfileirar(() =>
300        update($, gasto, a => {
301          if (!a) return a
302          const estimado = a.contexto === null ? null : Math.max(0, a.contexto - saiu)
303
304          return { ...a, contexto: informado ?? estimado }
305        }).then(() => undefined),
306      )
307      // Só o resumo do núcleo informa o tamanho: "micro" com tamanho é o núcleo resumindo, sem o plugin.
308      if (e.instructions?.trim().toLowerCase() === 'micro' && informado !== undefined) {
309        $.ui.toast('O /compact micro virou resumo: o plugin micro-compaction não está ativo.')
310      }
311    }
312
313    return r
314  })
315
316  // O /clear segue no mesmo processo com outro id e sem session.start: zera só a sessão.
317  on('session.end', async ($, e, next) => {
318    if (e.reason === 'clear') {
319      await enfileirar(() =>
320        update($, gasto, a => (a ? { ...a, sessao: ZERO, turnos: [], emTurno: false, contexto: null } : a)).then(() => undefined),
321      )
322    }
323
324    return next(e)
325  })
326
327  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
328    const g = await read($, gasto)
329    if (e.props.hasSurvey || g === null) return next(e)
330
331    const s = g.sessao
332    const entradaTotal = s.entrada + s.cacheLido + s.cacheEscrito
333    if (entradaTotal + s.saida === 0 && g.hoje === 0 && g.limites.length === 0 && g.contexto === null) return next(e)
334
335    const { Box, Button, Text } = $.ui.resolve(e)
336    const agora = await $.clock.now()
337    const emCurso = await read($, compactando)
338    const ocupado = e.props.isWorking || emCurso !== null
339    const cache = entradaTotal > 0 ? Math.round((s.cacheLido / entradaTotal) * 100) : 0
340    const ponto = <Text dimColor>  ·  </Text>
341
342    return (
343      <Box paddingX={1} flexDirection="row">
344        <Text wrap="truncate-end">
345          {g.turnos.length > 0 && <Text color="cyan">{sparkline(g.turnos)}  </Text>}
346          <Text bold>{compacto(novos(s))}</Text>
347          <Text dimColor> tokens</Text>
348          {ponto}
349          <Text dimColor>↑ </Text>
350          <Text>{compacto(s.entrada + s.cacheEscrito)}</Text>
351          <Text dimColor>  ↓ </Text>
352          <Text>{compacto(s.saida)}</Text>
353          {ponto}
354          <Text dimColor>cache </Text>
355          <Text>{cache}%</Text>
356          {ponto}
357          <Text dimColor>hoje </Text>
358          <Text bold>{compacto(g.hoje)}</Text>
359          {g.limites.map(l => {
360            // A janela já virou e nenhuma resposta nova trouxe o valor: não mostrar o antigo.
361            const pct = l.resetsAt && Date.parse(l.resetsAt) <= agora ? 0 : l.percentUsed
362
363            return (
364              <Text key={l.kind}>
365                {ponto}
366                <Text dimColor>{ROTULO[l.kind] ?? l.kind} </Text>
367                <Text color={cor(pct)}>{Math.round(pct)}%</Text>
368              </Text>
369            )
370          })}
371        </Text>
372        {/* Fora do texto que trunca: em terminal estreito o contexto é o último a sumir. */}
373        {g.contexto !== null && (
374          <Text>
375            {ponto}
376            <Text dimColor>ctx </Text>
377            <Text color={corContexto(g.contexto, g.janela)}>{compacto(g.contexto)}</Text>
378            {g.janela !== null && <Text dimColor>/{compacto(g.janela)}</Text>}
379          </Text>
380        )}
381        <Text>  </Text>
382        {ocupado ? (
383          e.props.isWorking ? null : <Text dimColor>compactando ({emCurso})…</Text>
384        ) : (
385          <Button
386            key="compactar"
387            label="compactar"
388            hotkey="c"
389            variant={g.contexto !== null && g.contexto >= limiares(g.janela).trocar ? 'primary' : undefined}
390            dimColor={g.contexto === null || g.contexto < limiares(g.janela).atencao}
391            onPress={() => compactar($)}
392          />
393        )}
394      </Box>
395    )
396  })
397}
398

1export type Tokens = { entrada: number; saida: number; cacheLido: number; cacheEscrito: number }
2
3/** micro: tira as leituras de arquivo (plugin micro-compaction); resumo: o /compact nativo. */
4export type Modo = 'micro' | 'resumo'
5
6export type Limite ={ kind: string; percentUsed: number; resetsAt?: string }
7
8export type Gasto = {
9  sessao: Tokens
10  /** Tokens novos do dia, somando todas as sessões (chaves `novos:<dia>:*` do $.store). */
11  hoje: number
12  /** Chave do $.store desta sessão no dia e o quanto ela já gravou ali. */
13  proprio: { chave: string; n: number }
14  /** Tokens novos por turno da thread principal, os mais recentes, para a sparkline. */
15  turnos: number[]
16  /** Se um turno da thread principal está rodando: só então o uso entra na última barra. */
17  emTurno: boolean
18  limites: Limite[]
19  /** Tokens do contexto atual (entrada da última resposta); null antes da primeira e depois do /clear. */
20  contexto: number | null
21  /** Janela de compactação (autoCompactWindow), de onde saem os limiares do ctx; null até a primeira leitura. */
22  janela: number | null
23}
24
25declare module 'claude-code' {
26  interface PluginState {
27    /** `compactando`: qual compactação o botão está rodando; null fora dela. */
28    'cc-cost': { tokens: Gasto | null; compactando: Modo | null }
29  }
30}
31