assumption-tracker

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Assumption } from '../types'
5
6const PANE = 'assumptions'
7const TOOL = 'mcp__assumption-tracker__record_assumption'
8
9const assumptions = atom({ plugin: 'assumption-tracker', key: 'assumptions' } as const, [] as Assumption[])
10
11const SECTION = {
12  id: 'assumption-tracker:assumptions',
13  scope: 'session',
14  text: `# Assumptions
15Whenever you act on an assumption, something about what the person wants, a project convention, or a fact about the code or environment that you were not told and have not checked, call \`${TOOL}\` with it in one sentence. Then keep working; do not stop to ask. The person reviews these in a side pane, and their verdicts come back to you: follow every correction, and stop relying on any assumption marked wrong.`,
16} as const
17
18const verdictLine = (a: Assumption) =>
19  a.verdict === 'right'
20    ? `- #${a.id} "${a.text}": right, keep it.`
21    : a.correction
22      ? `- #${a.id} "${a.text}": wrong. Correction: ${a.correction}`
23      : `- #${a.id} "${a.text}": wrong (no correction given), so stop relying on it.`
24
25/** The verdicts that have not reached Claude yet, as one message, and their ids. */
26function feedback(list: readonly Assumption[]): { text: string; ids: number[] } | undefined {
27  const unsent = list.filter(a => a.verdict !== undefined && !a.isSent)
28  if (unsent.length === 0) return undefined
29  return {
30    text: `The person reviewed assumptions you recorded:\n${unsent.map(verdictLine).join('\n')}\nAdjust your work to match.`,
31    ids: unsent.map(a => a.id),
32  }
33}
34
35const markSent = ($: EngineInterface, ids: number[]) =>
36  update($, assumptions, list => list.map(a => (ids.includes(a.id) ? { ...a, isSent: true } : a)))
37
38const judge = ($: EngineInterface, id: number, verdict: 'right' | 'wrong', correction?: string) =>
39  update($, assumptions, list =>
40    list.map(a => {
41      if (a.id !== id) return a
42      const { correction: _, ...rest } = a
43      return { ...rest, verdict, isSent: false, ...(verdict === 'wrong' && correction ? { correction } : {}) }
44    }),
45  )
46
47export const register: Register = on => {
48  on('session.start', async ($, e, next) => {
49    await $.tool.register({
50      name: 'record_assumption',
51      description:
52        'Record an assumption you are acting on (one you were not told and have not verified) so the person can confirm or correct it. Call it each time, then carry on without waiting.',
53      inputSchema: {
54        type: 'object',
55        properties: {
56          assumption: { type: 'string', description: 'The assumption in one sentence, e.g. "Dates are stored in UTC".' },
57          basis: { type: 'string', description: 'Optional: why you assumed it, in a few words.' },
58        },
59        required: ['assumption'],
60      },
61      isDeferred: false,
62    })
63    return next(e)
64  })
65
66  on('prompt.compose', async ($, e, next) => {
67    const composed = await next(e)
68    return { ...composed, sections: [...composed.sections, SECTION] }
69  })
70
71  on('tool.call', { tool: TOOL }, async ($, e) => {
72    const text = typeof e.assumption === 'string' ? e.assumption.trim() : ''
73    if (!text) return { deny: 'Pass `assumption`: the assumption in one sentence.' }
74    const basis = typeof e.basis === 'string' ? e.basis.trim() : ''
75
76    const list = await update($, assumptions, l => [
77      ...l,
78      { id: (l[l.length - 1]?.id ?? 0) + 1, text, isSent: false, ...(basis ? { basis } : {}) },
79    ])
80    return { result: `Recorded as assumption #${list[list.length - 1]?.id}. Keep going; the person will confirm or correct it.` }
81  })
82
83  // Verdicts not sent yet ride along with the next prompt, unseen by the person.
84  on('prompt.submit', async ($, e, next) => {
85    const pending = feedback(await read($, assumptions))
86    if (pending === undefined) return next(e)
87
88    const submitted = await next({ ...e, context: [...(e.context ?? []), pending.text] })
89    if (submitted.drop === undefined) await markSent($, pending.ids)
90    return submitted
91  })
92
93  // A /clear starts a new conversation: its assumptions start over too.
94  on('session.end', async ($, e, next) => {
95    if (e.reason === 'clear') await update($, assumptions, () => [])
96    return next(e)
97  })
98
99  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
100    const list = await read($, assumptions)
101    if (e.props.hasSurvey || list.length === 0) return next(e)
102
103    const { Box, Button, Text } = $.ui.resolve(e)
104    const toReview = list.filter(a => a.verdict === undefined).length
105
106    return (
107      <Box>
108        <Text dimColor>
109          Assumptions ({list.length}){toReview > 0 ? `, ${toReview} to review` : ''}{' '}
110        </Text>
111        <Button
112          key="show"
113          label="Show"
114          hotkey="a"
115          onPress={() => $.ui.open({ id: PANE, title: 'Assumptions', focus: true, closeOnEscape: true })}
116        />
117      </Box>
118    )
119  })
120
121  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
122    const list = await read($, assumptions)
123    const { Box, Button, Text } = $.ui.resolve(e)
124    const Input = e.surface === 'mobile' ? undefined : $.ui.resolve(e).Input
125    const pending = feedback(list)
126
127    const send = async () => {
128      const now = feedback(await read($, assumptions))
129      if (now === undefined) return
130      await markSent($, now.ids) // first, so the prompt.submit hook doesn't attach them again
131      await $.prompt.submit({ text: now.text })
132    }
133
134    return (
135      <Box flexDirection="column">
136        {list.length === 0 && <Text dimColor>No assumptions recorded yet.</Text>}
137        {list.map(a => (
138          <Box key={`a:${a.id}`} flexDirection="column" marginBottom={1}>
139            <Text>
140              <Text bold>#{a.id}</Text> {a.text}
141            </Text>
142            {a.basis ? <Text dimColor>because {a.basis}</Text> : null}
143            {a.verdict === undefined ? null : (
144              <Text color={a.verdict === 'right' ? 'success' : 'error'}>
145                {a.verdict === 'right' ? '✓ right' : `✗ wrong${a.correction ? `: ${a.correction}` : ''}`}
146                <Text dimColor>{a.isSent ? ' (sent to Claude)' : ' (not sent yet)'}</Text>
147              </Text>
148            )}
149            <Box>
150              <Button
151                key={`right:${a.id}`}
152                label="✓ Right"
153                variant={a.verdict === 'right' ? 'primary' : 'secondary'}
154                onPress={() => judge($, a.id, 'right')}
155              />
156              <Text> </Text>
157              <Button
158                key={`wrong:${a.id}`}
159                label="✗ Wrong"
160                variant={a.verdict === 'wrong' ? 'primary' : 'secondary'}
161                onPress={() => judge($, a.id, 'wrong', a.correction)}
162              />
163            </Box>
164            {Input === undefined ? null : (
165              <Input
166                key={`fix:${a.id}`}
167                label="Correction: "
168                placeholder="what is actually true"
169                submitLabel="mark wrong"
170                value={a.correction ?? ''}
171                onSubmit={value => judge($, a.id, 'wrong', value.trim())}
172              />
173            )}
174          </Box>
175        ))}
176        {pending === undefined ? null : (
177          <Box>
178            <Button key="send" label={`Send feedback (${pending.ids.length})`} variant="primary" onPress={send} />
179            <Text dimColor> or it goes with your next prompt</Text>
180          </Box>
181        )}
182      </Box>
183    )
184  })
185}
186

1/** One assumption Claude recorded, and the person's verdict on it once given. */
2export type Assumption = {
3  id: number
4  text: string
5  /** Why Claude assumed it, when it said. */
6  basis?: string
7  verdict?: 'right' | 'wrong'
8  /** What is actually true, when the person typed it. */
9  correction?: string
10  /** Whether the verdict has reached Claude. */
11  isSent: boolean
12}
13
14declare module 'claude-code' {
15  interface PluginState {
16    'assumption-tracker': { assumptions: Assumption[] }
17  }
18}
19