use-haiku

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PromptOrigin, Register } from 'claude-code'
3
4import type { AgentIds } from '../types'
5
6const MARK = '_'
7const HAIKU = 'claude-haiku-5-5'
8const COLOR = 'suggestion'
9
10const MARKED = { plugin: 'use-haiku', key: 'marked' } as const
11const TURN = { plugin: 'use-haiku', key: 'turn' } as const
12const AGENTS = { plugin: 'use-haiku', key: 'agents' } as const
13
14const marked = atom(MARKED, false)
15const haikuTurn = atom(TURN, null as string | null)
16const haikuAgents = atom(AGENTS, [] as AgentIds)
17
18// `_summarize the diff` is marked; `__init__.py fails` is not.
19const isMarked = (text: string) => text.startsWith(MARK) && !text.startsWith(MARK + MARK)
20
21// Prompts in the person's own words: typed, from a phone or through -p, scheduled, or sent by a plugin for them (After's queue).
22const OWN = new Set<PromptOrigin['kind']>(['composer', 'bridge', 'sdk', 'scheduled-trigger', 'plugin'])
23
24const NOTE = `The person sent this prompt to Haiku: you are ${HAIKU} while you answer it, whatever model earlier context names.`
25
26// What `marked` holds, so a keystroke writes it only when it flips.
27let lastMarked: boolean | undefined
28
29async function setMarked ($: EngineInterface, value: boolean) {
30  if (value === lastMarked) return
31  lastMarked = value
32  await $.state.set(MARKED, value)
33}
34
35export const register: Register = on => {
36  // A marked prompt sent while Claude was idle: the next turn of the main loop is its own.
37  let pending = false
38  // The line under a Haiku answer is for a person at the prompt, not for -p output.
39  let isInteractive = false
40
41  on('session.start', ($, e, next) => {
42    isInteractive = e.isInteractive
43    return next(e)
44  })
45
46  // The mark lights up in the box as soon as it is typed.
47  on('prompt.edit', async ($, e, next) => {
48    const box = await next(e)
49    const isNow = isMarked(box.text)
50    await setMarked($, isNow)
51    if (!isNow) return box
52    return { ...box, decorations: [...(box.decorations ?? []), { start: 0, end: MARK.length, color: COLOR, bold: true }] }
53  })
54
55  on('prompt.submit', async ($, e, next) => {
56    await setMarked($, false)
57    const text = e.text.slice(MARK.length).trimStart()
58    if (!isMarked(e.text) || text === '' || !OWN.has(e.origin.kind)) {
59      if (e.turnId === undefined) pending = false
60      return next(e)
61    }
62
63    // Typed while a turn runs, it is folded into that turn or queued behind it, and which is not
64    // known here: rather than guess, it keeps the current model, and says so.
65    if (e.turnId !== undefined) {
66      $.ui.toast('⚡ Sent mid-turn, so not on Haiku: it keeps the current model')
67      return next({ ...e, text })
68    }
69
70    // The model learns which model it is once per conversation, from the session's: asked "which
71    // model are you?" it would name that one. The prompt itself tells it, unseen by the person.
72    pending = true
73    const result = await next({ ...e, text, context: [...(e.context ?? []), NOTE] })
74    if (result.drop !== undefined) pending = false
75    return result
76  })
77
78  // Only the main loop raises turn.start; a subagent's run does not.
79  on('turn.start', async ($, e, next) => {
80    await $.state.set(TURN, pending ? e.turnId : null)
81    pending = false
82    return next(e)
83  })
84
85  // Every request of a Haiku turn goes to Haiku: the main loop's and those of the subagents it started.
86  on('turn.step', async function* ($, e, next) {
87    const isHaiku = e.agentId === undefined
88      ? e.turnId === (await read($, haikuTurn))
89      : (await read($, haikuAgents)).includes(e.agentId)
90    return yield* next(isHaiku ? { ...e, model: HAIKU } : e)
91  })
92
93  on('agent.spawn', async ($, e, next) => {
94    const isHaiku = !e.isTeammate && (e.parentAgentId === undefined
95      ? (await read($, haikuTurn)) !== null
96      : (await read($, haikuAgents)).includes(e.parentAgentId))
97    if (!isHaiku) return next(e)
98
99    // A fork inherits its parent's model and a workflow's agent keeps its own: turn.step catches theirs.
100    const result = await next(e.fork || e.workflow ? e : { ...e, model: HAIKU })
101    const id = result.deny === undefined ? result.agentId : undefined
102    if (id !== undefined) await update($, haikuAgents, ids => [...ids, id])
103    return result
104  })
105
106  on('turn.complete', async ($, e, next) => {
107    if (e.agentId !== undefined) return next(e)
108
109    // Cleared before next(e): a plugin beneath may start the next turn from here (After does).
110    const turn = await read($, haikuTurn)
111    if (turn !== null) await $.state.set(TURN, null)
112    const result = await next(e)
113    if (turn !== e.turnId || e.reason !== 'answer' || e.usage === undefined || !isInteractive) return result
114
115    // What the API says answered, so a Haiku turn is seen to be one.
116    const model = e.usage.model
117    const line = model.includes('haiku') ? `⚡ Answered by ${model}` : `⚠ Asked for Haiku, but ${model} answered`
118    return { ...result, text: result.text === e.answer ? line : `${result.text}\n${line}` }
119  })
120
121  on('ui.render', { component: 'Spinner' }, async ($, e, next) => {
122    if ((await read($, haikuTurn)) === null) return next(e)
123    // The first frame of a turn comes before the word is sampled.
124    const word = e.props.word === '' ? 'Haiku' : `${e.props.word} on Haiku`
125    return next({ ...e, props: { ...e.props, word } })
126  })
127
128  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
129    const isMarkedDraft = await read($, marked)
130    if (!e.props.isDraft || !isMarkedDraft) return next(e)
131
132    const { Box, Text } = $.ui.resolve(e)
133    return (
134      <Box flexDirection="row">
135        {e.props.isWorking
136          ? <Text color="warning">⚡ not on Haiku mid-turn: wait for Claude to finish</Text>
137          : <Text color={COLOR}>⚡ this prompt runs on Haiku</Text>}
138        <Text>  </Text>
139        {await next(e)}
140      </Box>
141    )
142  })
143}
144

1/** The subagents a Haiku turn started, by id: their requests go to Haiku too. */
2export type AgentIds = readonly string[]
3
4declare module 'claude-code' {
5  interface PluginState {
6    // marked: the draft in the prompt box starts with the mark.
7    // turn: the main loop's turn that runs on Haiku, while it runs.
8    'use-haiku': { marked: boolean; turn: string | null; agents: AgentIds }
9  }
10}
11