cache-battery

1// Claude Code scans this file before loading it: every hook must be a function
2// literal here and every host call spelled $.noun.event(...), so the host API
3// stays in this file and ./lib holds only pure helpers.
4import { modOptions, toInk } from './lib/cc-mod.mjs'
5import { renderBattery } from './lib/core/render.mjs'
6import { defaultTier, fromSamples, sampleFromAnthropicUsage, tierOf } from './lib/core/state.mjs'
7
8export function register(on) {
9  let state
10  let tier = '1h'
11  let now = 0
12  let lastKey = ''
13  let options = modOptions({})
14  let ticking = false
15  // Main-conversation requests in flight keep their cached prefix alive.
16  let inFlight = 0
17
18  const shown = () => (inFlight > 0 && state ? { ...state, anchorAt: now, chargingUntil: undefined } : state)
19
20  on('session.start', async ($, e, next) => {
21    // Claude Code does not pass the 5m/1h write split to mods, so the tier follows its own
22    // settings. ANTHROPIC_API_KEY and ANTHROPIC_AUTH_TOKEN are only tested for presence.
23    const env = {
24      CACHE_BATTERY_TTL: await $.env.get('CACHE_BATTERY_TTL'),
25      CACHE_BATTERY_NUMBERS: await $.env.get('CACHE_BATTERY_NUMBERS'),
26      CACHE_BATTERY_CELLS: await $.env.get('CACHE_BATTERY_CELLS'),
27      FORCE_PROMPT_CACHING_5M: await $.env.get('FORCE_PROMPT_CACHING_5M'),
28      CLAUDE_CODE_PROMPT_CACHE_TTL: await $.env.get('CLAUDE_CODE_PROMPT_CACHE_TTL'),
29      ENABLE_PROMPT_CACHING_1H: await $.env.get('ENABLE_PROMPT_CACHING_1H'),
30      ANTHROPIC_API_KEY: await $.env.get('ANTHROPIC_API_KEY'),
31      ANTHROPIC_AUTH_TOKEN: await $.env.get('ANTHROPIC_AUTH_TOKEN'),
32      CLAUDE_CODE_USE_BEDROCK: await $.env.get('CLAUDE_CODE_USE_BEDROCK'),
33      CLAUDE_CODE_USE_VERTEX: await $.env.get('CLAUDE_CODE_USE_VERTEX'),
34      CLAUDE_CODE_USE_FOUNDRY: await $.env.get('CLAUDE_CODE_USE_FOUNDRY'),
35    }
36    tier = defaultTier(env)
37    options = modOptions(env)
38    now = await $.clock.now()
39    if (!ticking) {
40      ticking = true
41      $.clock.every(1000, async () => {
42        now = await $.clock.now()
43        const key = JSON.stringify(renderBattery(shown(), now, options))
44        if (key === lastKey) return
45        lastKey = key
46        $.ui.invalidate('ui.render')
47      })
48    }
49    return next(e)
50  })
51
52  on('classic.SessionStart', { source: ['clear', 'resume', 'fork'] }, async ($, e, next) => {
53    state = undefined
54    return next(e)
55  })
56
57  on('turn.step', async function* ($, e, next) {
58    const startedAt = await $.clock.now()
59    if (!e.agentId) {
60      inFlight++
61      $.ui.invalidate('ui.render')
62    }
63    let result
64    try {
65      result = yield* next(e)
66    } finally {
67      if (!e.agentId) {
68        inFlight--
69        $.ui.invalidate('ui.render')
70      }
71    }
72    if (e.agentId || !result || !result.usage) return result
73    const sample = sampleFromAnthropicUsage(result.usage, startedAt)
74    tier = tierOf(sample) ?? tier
75    // A failed or cancelled step throws or has no usage and keeps the state above; a
76    // finished step with no cache use means a model that does not cache.
77    state = fromSamples([sample], tier)
78    $.ui.invalidate('ui.render')
79    return result
80  })
81
82  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
83    const below = await next(e)
84    if (e.props && e.props.hasSurvey) return below
85    now = await $.clock.now()
86    const segments = renderBattery(shown(), now, options)
87    if (!segments.length) return below
88    const { Box, Text } = $.ui.resolve(e)
89    const battery = Box({ flexDirection: 'row', flexGrow: 0, flexShrink: 0, children: segments.map((s) => toInk(s, Text)) })
90    const row = Box({ flexDirection: 'row', children: [battery, Box({ flexGrow: 1 })] })
91    return Box({ flexDirection: 'column', children: below ? [row, below] : [row] })
92  })
93}
94

1import { parseNumbers } from './core/render.mjs';
2const TRACK = '#5c5c5c';
3const TRACK_GLYPH = '█';
4const INK_COLOR = { green: 'green', yellow: 'yellow', red: 'red', cyan: 'cyan' };
5export function modOptions(env) {
6    const cells = Number(env.CACHE_BATTERY_CELLS ?? 8);
7    return { cells: Number.isInteger(cells) && cells > 0 ? cells : 8, numbers: parseNumbers(env.CACHE_BATTERY_NUMBERS) };
8}
9/** No backgroundColor: in the band a background fills the cell's whole flex area, so empty cells are a track glyph instead. */
10export function toInk(segment, Text) {
11    if (segment.kind === 'empty')
12        return Text({ color: TRACK, children: [TRACK_GLYPH.repeat(segment.count)] });
13    const tone = segment.tone;
14    const style = tone === 'dim' ? { dimColor: true } : tone ? { color: INK_COLOR[tone] } : {};
15    if (segment.kind === 'fill')
16        return Text({ ...style, children: [segment.text] });
17    return Text({ ...style, bold: segment.bold, children: [segment.text] });
18}
19

1import { fraction, remainingMs } from './state.mjs';
2const EIGHTHS = ['', '▏', '▎', '▍', '▌', '▋', '▊', '▉'];
3const BADGE = { '5m': '◔', '1h': '●' };
4const ESTIMATED_BADGE = '○';
5const CAP = '▌';
6const EMPTY_GLYPH = '░';
7const MINUTE = 60_000;
8export function renderBattery(state, now, options = {}) {
9    if (!state)
10        return [];
11    const cells = options.cells ?? 8;
12    const left = remainingMs(state, now);
13    const cap = { kind: 'text', text: CAP, tone: 'dim' };
14    if (left <= 0) {
15        return [{ kind: 'text', text: '❄ ', tone: 'cyan' }, { kind: 'empty', count: cells }, cap];
16    }
17    const level = fraction(state, now);
18    const tone = level > 0.5 ? 'green' : level > 0.2 ? 'yellow' : 'red';
19    const eighths = Math.min(cells * 8, Math.ceil(level * cells * 8));
20    const full = Math.floor(eighths / 8);
21    const partial = EIGHTHS[eighths % 8];
22    const used = full + (partial ? 1 : 0);
23    const segments = [
24        { kind: 'text', text: `${state.estimated ? ESTIMATED_BADGE : BADGE[state.tier]} `, tone: 'dim' },
25        { kind: 'fill', text: '█'.repeat(full) + partial, tone },
26    ];
27    if (used < cells)
28        segments.push({ kind: 'empty', count: cells - used });
29    segments.push(cap);
30    const label = countdownLabel(left, state.ttlMs, options.numbers ?? 'end');
31    if (label)
32        segments.push({ kind: 'text', text: ` ${label}`, tone, bold: left <= MINUTE });
33    if (state.chargingUntil !== undefined && now < state.chargingUntil)
34        segments.push({ kind: 'text', text: ' ⚡', tone: 'yellow' });
35    return segments;
36}
37function countdownLabel(left, ttlMs, mode) {
38    if (mode === 'never')
39        return undefined;
40    if (left <= MINUTE)
41        return `${Math.ceil(left / 1000)}s`;
42    const nearEnd = ttlMs > 5 * MINUTE && left <= 5 * MINUTE;
43    if (mode === 'always' || nearEnd)
44        return `${Math.ceil(left / MINUTE)}m`;
45    return undefined;
46}
47export function toPlain(segments) {
48    return segments.map((s) => (s.kind === 'empty' ? EMPTY_GLYPH.repeat(s.count) : s.text)).join('');
49}
50/** `dim` is an explicit grey: SGR 2 is nearly invisible in some hosts' footers. */
51const FG = { green: '32', yellow: '33', red: '31', dim: '38;5;245', cyan: '36' };
52const TRACK_BG = '48;5;239';
53const TRACK_FG = '38;5;239';
54/** Empty cells are solid blocks in the track colour, not spaces: pi's footer collapses runs of spaces. */
55export function toAnsi(segments, color) {
56    if (!color)
57        return toPlain(segments);
58    const out = segments.map((s) => {
59        if (s.kind === 'empty')
60            return `\x1b[${TRACK_FG}m${'█'.repeat(s.count)}\x1b[0m`;
61        if (s.kind === 'fill')
62            return `\x1b[${FG[s.tone]};${TRACK_BG}m${s.text}\x1b[0m`;
63        const codes = [s.tone ? FG[s.tone] : '', s.bold ? '1' : ''].filter(Boolean).join(';');
64        return codes ? `\x1b[${codes}m${s.text}\x1b[0m` : s.text;
65    });
66    return out.join('');
67}
68export function colorEnabled(env) {
69    return !env.NO_COLOR && env.FORCE_COLOR !== '0';
70}
71export function parseNumbers(value) {
72    return value === 'always' || value === 'never' ? value : 'end';
73}
74

1export const TTL_MS = { '5m': 5 * 60_000, '1h': 60 * 60_000 };
2const CHARGING_MS = 5000;
3export function sampleFromAnthropicUsage(usage, at) {
4    return {
5        at,
6        cacheRead: usage.cache_read_input_tokens ?? 0,
7        cacheWrite: usage.cache_creation_input_tokens ?? 0,
8        write1h: usage.cache_creation?.ephemeral_1h_input_tokens,
9        write5m: usage.cache_creation?.ephemeral_5m_input_tokens,
10    };
11}
12export function tierOf(sample) {
13    if ((sample.write1h ?? 0) > 0)
14        return '1h';
15    if ((sample.write5m ?? 0) > 0)
16        return '5m';
17    return undefined;
18}
19/**
20 * A newest request with no cache use means a provider that does not cache, so the
21 * state is unknown; adapters drop failed and cancelled requests before this. Cache-reading
22 * requests record empty write buckets, so the tier comes from the newest request that wrote.
23 */
24export function fromSamples(samples, fallback) {
25    const newest = samples.at(-1);
26    if (newest?.estimated)
27        return { tier: '5m', anchorAt: newest.at, ttlMs: TTL_MS['5m'], estimated: true };
28    if (!newest || newest.cacheRead + newest.cacheWrite <= 0)
29        return undefined;
30    let tier = fallback;
31    for (let i = samples.length - 1; i >= 0; i--) {
32        const found = tierOf(samples[i]);
33        if (found) {
34            tier = found;
35            break;
36        }
37    }
38    const state = { tier, anchorAt: newest.at, ttlMs: TTL_MS[tier] };
39    if (newest.refresh)
40        state.chargingUntil = newest.at + CHARGING_MS;
41    return state;
42}
43export function fromPromptCache(field) {
44    if (!field?.caching_observed)
45        return undefined;
46    const tier = field.ttl === '1h' ? '1h' : '5m';
47    const ttlMs = TTL_MS[tier];
48    const anchorAt = typeof field.expires_at === 'number' ? field.expires_at * 1000 - ttlMs : 0;
49    return { tier, anchorAt, ttlMs };
50}
51const isTier = (value) => value === '5m' || value === '1h';
52const truthy = (value) => /^(1|true|yes|on)$/i.test(value?.trim() ?? '');
53/**
54 * Mirrors Claude Code's TTL choice in its order: FORCE_PROMPT_CACHING_5M, then
55 * CLAUDE_CODE_PROMPT_CACHE_TTL, then ENABLE_PROMPT_CACHING_1H (or its Bedrock variant),
56 * then 1h on a subscription and 5m on API keys and cloud providers. The promptCacheTtl
57 * setting and subscription overage are not visible here.
58 */
59export function defaultTier(env) {
60    if (isTier(env.CACHE_BATTERY_TTL))
61        return env.CACHE_BATTERY_TTL;
62    if (truthy(env.FORCE_PROMPT_CACHING_5M))
63        return '5m';
64    if (isTier(env.CLAUDE_CODE_PROMPT_CACHE_TTL))
65        return env.CLAUDE_CODE_PROMPT_CACHE_TTL;
66    if (truthy(env.ENABLE_PROMPT_CACHING_1H))
67        return '1h';
68    if (truthy(env.CLAUDE_CODE_USE_BEDROCK) && truthy(env.ENABLE_PROMPT_CACHING_1H_BEDROCK))
69        return '1h';
70    const metered = env.ANTHROPIC_API_KEY ||
71        env.ANTHROPIC_AUTH_TOKEN ||
72        truthy(env.CLAUDE_CODE_USE_BEDROCK) ||
73        truthy(env.CLAUDE_CODE_USE_VERTEX) ||
74        truthy(env.CLAUDE_CODE_USE_FOUNDRY);
75    return metered ? '5m' : '1h';
76}
77export function remainingMs(state, now) {
78    return state.anchorAt + state.ttlMs - now;
79}
80export function fraction(state, now) {
81    return Math.min(1, Math.max(0, remainingMs(state, now) / state.ttlMs));
82}
83