ultramodel

1import { atom, read, update } from 'claude-code'
2import type { Hook, Register, TurnCompleteInput } from 'claude-code'
3
4import type { Status } from '../types'
5import { costUsd, highlight, labelOf, pickModel } from './logic'
6
7type Dollar = Parameters<Hook<'turn.complete'>>[0]
8
9const status = atom({ plugin: 'ultramodel', key: 'status' } as const, {
10  next: null,
11  current: null,
12  lastUsd: null,
13  sessionUsd: 0,
14  switched: 0,
15} as Status)
16
17type Turn = { turnId: string; model: string | null }
18
19let pending: string | null = null
20let active: Turn | null = null
21
22export const register: Register = on => {
23  on('prompt.submit', ($, e, next) => {
24    pending = pickModel(e.text)
25
26    return next(e)
27  })
28
29  on('turn.start', async ($, e, next) => {
30    if (e.agentId === undefined) {
31      const model = pending
32      pending = null
33      active = { turnId: e.turnId, model }
34
35      if (model !== null) {
36        $.ui.toast(`Ultramodel: ${labelOf(model)} for this turn`)
37        await update($, status, s => ({ ...s, current: labelOf(model) }))
38      }
39    }
40
41    return next(e)
42  })
43
44  on('turn.step', async function* ($, e, next) {
45    const isSwitched = active !== null && active.model !== null && e.agentId === undefined && e.turnId === active.turnId
46
47    return yield* next(isSwitched ? { ...e, model: active.model } : e)
48  })
49
50  on('turn.complete', async ($, e, next) => {
51    if (active !== null && e.agentId === undefined && e.turnId === active.turnId) {
52      const turn = active
53      active = null
54
55      await record($, turn, e)
56    }
57
58    return next(e)
59  })
60
61  on('prompt.edit', async ($, e, next) => {
62    const box = await next(e)
63
64    return { ...box, decorations: [...(box.decorations ?? []), ...highlight(box.text)] }
65  })
66
67  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
68    const s = await read($, status)
69
70    if (s.current === null && s.switched === 0) {
71      return next(e)
72    }
73
74    const { Box, Text } = $.ui.resolve(e)
75
76    return (
77      <Box>
78        {s.current === null ? null : <Text color="suggestion">Ultramodel: {s.current} this turn </Text>}
79        <Text dimColor>
80          {s.switched} switched turns, ${s.sessionUsd.toFixed(4)} this session
81          {s.lastUsd === null ? '' : `, last $${s.lastUsd.toFixed(4)}`}
82        </Text>
83      </Box>
84    )
85  })
86}
87
88async function record($: Dollar, turn: Turn, e: TurnCompleteInput): Promise<void> {
89  const usage = e.usage
90  const usd = usage === undefined ? null : costUsd(usage.model, usage)
91  const row = {
92    ts: new Date(await $.clock.now()).toISOString(),
93    turnId: turn.turnId,
94    reason: e.reason,
95    requested: turn.model,
96    answered: usage?.model ?? null,
97    input: usage?.input_tokens ?? 0,
98    output: usage?.output_tokens ?? 0,
99    cacheRead: usage?.cache_read_input_tokens ?? 0,
100    cacheWrite: usage?.cache_creation_input_tokens ?? 0,
101    usd,
102  }
103
104  try {
105    await writeLog($, row)
106  } catch (err) {
107    $.ui.log(`ultramodel: could not write the cost log: ${String(err)}`)
108  }
109
110  await update($, status, s => ({
111    next: s.next,
112    current: null,
113    lastUsd: usd,
114    sessionUsd: s.sessionUsd + (usd ?? 0),
115    switched: s.switched + (turn.model === null ? 0 : 1),
116  }))
117}
118
119// One JSON line per finished main-loop turn, in ~/.claude/ultramodel/.
120// Holds no prompt text: only ids, reasons and token counts.
121async function writeLog($: Dollar, row: object): Promise<void> {
122  const home = await $.env.get('HOME')
123  const day = new Date(await $.clock.now()).toISOString().slice(0, 10)
124  const path = `${home}/.claude/ultramodel/usage-${day}.jsonl`
125  const prior = (await $.fs.exists(path)) ? await $.fs.read(path) : ''
126
127  await $.fs.write(path, `${prior}${JSON.stringify(row)}\n`)
128}
129

1// Pure helpers for the ultramodel mod. No engine imports, so plain Node tests can run them.
2import type { PromptDecoration } from 'claude-code'
3
4// Model names a prompt can mention, by the word typed.
5export const MODELS: Record<string, { id: string; label: string }> = {
6  opus: { id: 'claude-opus-5-5', label: 'Opus 5.5' },
7  sonnet: { id: 'claude-sonnet-5-5', label: 'Sonnet 5.5' },
8  haiku: { id: 'claude-haiku-5-5', label: 'Haiku 5.5' },
9}
10
11// USD per million tokens. Input and output are the published rates. The cache
12// write and Haiku cache read rates are assumptions (1.25x and 0.1x of input);
13// confirm them before trusting the logged cost.
14export const PRICES: Record<string, { input: number; output: number; cacheRead: number; cacheWrite: number }> = {
15  'claude-opus-5-5': { input: 4, output: 20, cacheRead: 0.2, cacheWrite: 5 },
16  'claude-sonnet-5-5': { input: 2, output: 10, cacheRead: 0.2, cacheWrite: 2.5 },
17  'claude-haiku-5-5': { input: 0.1, output: 0.5, cacheRead: 0.01, cacheWrite: 0.125 },
18}
19
20export const RAINBOW = ['error', 'warning', 'success', 'suggestion', 'permission', 'claude'] as const
21
22export const NAME = /\b(opus|sonnet|haiku)\b/gi
23
24// The first model named in a prompt, as an API id; null when none is named.
25export function pickModel(text: string): string | null {
26  const match = text.match(/\b(opus|sonnet|haiku)\b/i)
27
28  return match === null ? null : MODELS[match[1].toLowerCase()].id
29}
30
31export function labelOf(id: string): string {
32  return Object.values(MODELS).find(model => model.id === id)?.label ?? id
33}
34
35// Paints each letter of a model name in a cycling color, as ultrathink does.
36export function highlight(text: string): PromptDecoration[] {
37  const runs: PromptDecoration[] = []
38
39  for (const match of text.matchAll(NAME)) {
40    const start = match.index ?? 0
41
42    for (let i = 0; i < match[0].length; i += 1) {
43      runs.push({
44        start: start + i,
45        end: start + i + 1,
46        color: RAINBOW[(start + i) % RAINBOW.length],
47        bold: true,
48      })
49    }
50  }
51
52  return runs
53}
54
55export function costUsd(model: string, usage: { input_tokens: number; output_tokens: number; cache_read_input_tokens: number; cache_creation_input_tokens: number }): number | null {
56  const price = PRICES[model]
57
58  if (price === undefined) {
59    return null
60  }
61
62  const usd =
63    usage.input_tokens * price.input +
64    usage.output_tokens * price.output +
65    usage.cache_read_input_tokens * price.cacheRead +
66    usage.cache_creation_input_tokens * price.cacheWrite
67
68  return usd / 1_000_000
69}
70

1export type Status = {
2  // The model the box has picked for the next turn, by its display label.
3  next: string | null
4  // The model the running turn was switched to; null when no switch.
5  current: string | null
6  // Cost of the last finished main-loop turn, in USD; null before any.
7  lastUsd: number | null
8  // Cost of every finished main-loop turn this session, in USD.
9  sessionUsd: number
10  // Count of finished turns that ran on a switched model.
11  switched: number
12}
13
14declare module 'claude-code' {
15  interface PluginState {
16    ultramodel: { status: Status }
17  }
18}
19