INITE Brain — persistent memory

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, McpToolResult, Register } from 'claude-code'
3
4import type { BrainHit, BrainRecall, BrainWorkspace } from '../types'
5
6/**
7 * Brain inside the session, not only at its edges.
8 *
9 * The command hooks beside this module recall once, when a session
10 * starts, and capture when it compacts or ends. Between those two moments
11 * the model only reaches memory if it decides to call a brain tool. This
12 * module closes that gap where the engine allows it:
13 *
14 * - every prompt is looked up in brain before the model reads it, and
15 *   what comes back rides beside the prompt as context the person never
16 *   has to paste (bounded: a slow brain costs at most RECALL_TIMEOUT_MS,
17 *   and a failure costs nothing but the recall);
18 * - a band above the prompt says what was recalled, so memory working is
19 *   visible rather than assumed;
20 * - `/brain` opens a pane to ask memory directly, without a model turn;
21 * - the status line names the workspace this session writes to.
22 *
23 * Every call goes through the plugin's own MCP connection (`$.mcp`), so
24 * whatever credential that connection holds — the API key from the
25 * plugin's settings, or an OAuth sign-in — is the one used. This module
26 * never sees a key.
27 */
28
29const PLUGIN = 'inite-brain'
30const PANE = 'brain'
31const RECALL_TIMEOUT_MS = 2500
32const MIN_PROMPT_CHARS = 15
33const MAX_CONTEXT_CHARS = 2000
34const HITS = 6
35
36const recallAtom = atom({ plugin: 'inite-brain', key: 'recall' } as const, null)
37const foundAtom = atom({ plugin: 'inite-brain', key: 'found' } as const, null)
38const workspaceAtom = atom({ plugin: 'inite-brain', key: 'workspace' } as const, null)
39const hiddenAtom = atom({ plugin: 'inite-brain', key: 'isBandHidden' } as const, false)
40
41type Options = { user_id?: string; recall_each_prompt?: boolean }
42
43/** The JSON a brain tool answered with, or null when it answered none. */
44function payload(result: McpToolResult) {
45  if (result.isError) return null
46  const block = result.content.find(one => one.type === 'text')
47  const text = block && 'text' in block ? block.text : undefined
48  if (typeof text !== 'string' || !text) return null
49  try {
50    return JSON.parse(text) as Record<string, unknown>
51  } catch {
52    return null
53  }
54}
55
56export function hitsOf(out: Record<string, unknown> | null): BrainHit[] {
57  const results = Array.isArray(out?.results) ? out.results : []
58  const hits: BrainHit[] = []
59  for (const raw of results as Array<Record<string, unknown>>) {
60    const facts = (Array.isArray(raw.facts) ? raw.facts : [])
61      .slice(0, 4)
62      .map((f: Record<string, unknown>) => `${String(f.predicate ?? '')} ${String(f.object ?? '')}`.trim())
63      .filter(Boolean)
64    if (typeof raw.canonicalName === 'string' && facts.length > 0) {
65      hits.push({ name: raw.canonicalName, facts })
66    }
67  }
68  return hits.slice(0, HITS)
69}
70
71/** The recall as the model reads it: one line per entity, bounded. */
72export function contextOf(hits: BrainHit[]): string {
73  const lines = [
74    'Brain (long-term memory) recalled this for the prompt above. It may be stale or beside the point — use what bears on the request, and ask brain directly (its MCP tools) for more or to check:',
75  ]
76  for (const hit of hits) {
77    const line = `- ${hit.name}: ${hit.facts.join('; ')}`
78    if (lines.join('\n').length + line.length > MAX_CONTEXT_CHARS) break
79    lines.push(line)
80  }
81  return lines.join('\n')
82}
83
84/**
85 * One search_knowledge call, bounded by `RECALL_TIMEOUT_MS`: a recall that
86 * is late is dropped rather than waited for, because the person is
87 * waiting on the prompt behind it.
88 */
89async function search(
90  $: EngineInterface,
91  at: { server: string | null; userId?: string | undefined },
92  query: string,
93  signal?: AbortSignal,
94): Promise<BrainRecall> {
95  const started = await $.clock.now()
96  if (at.server === null) return { query, hits: [], ms: 0, failure: 'offline' }
97  const args: Record<string, unknown> = { query, limit: HITS }
98  if (at.userId) args.userId = at.userId
99  const call = $.mcp.call(at.server, 'search_knowledge', args).then(
100    (result: McpToolResult) => ({
101      hits: hitsOf(payload(result)),
102      failure: result.isError ? 'error' : undefined,
103    }),
104    () => ({ hits: [] as BrainHit[], failure: 'error' }),
105  )
106  const timeout = $.clock
107    .sleep(RECALL_TIMEOUT_MS, signal ? { signal } : undefined)
108    .then(() => ({ hits: [] as BrainHit[], failure: 'timeout' }))
109    .catch(() => ({ hits: [] as BrainHit[], failure: 'timeout' }))
110  const { hits, failure } = await Promise.race([call, timeout])
111  const ms = (await $.clock.now()) - started
112  return failure ? { query, hits, ms, failure } : { query, hits, ms }
113}
114
115export const register: Register = (on, options) => {
116  const opts = options as Options
117  // The name /mcp lists the plugin's server under; set at session start.
118  let server: string | null = null
119  const at = () => ({ server, userId: opts.user_id })
120
121  on('session.start', async ($, e, next) => {
122    await $.command.register({
123      name: 'brain',
124      description: 'Ask brain (long-term memory) directly: /brain [question]',
125      argumentHint: '[question]',
126    })
127    const connected = await $.mcp.connect('brain')
128    if (!connected.isConnected) {
129      $.ui.status('brain · not connected (/mcp)')
130      return next(e)
131    }
132    server = connected.server
133    const status = payload(await $.mcp.call(server, 'workspace_status', {}).catch(() => ({ content: [], isError: true })))
134    const memory = (status?.memory ?? {}) as Record<string, unknown>
135    const workspace: BrainWorkspace = {
136      name: String(status?.displayName ?? status?.companyId ?? 'workspace'),
137      facts: Number(memory.facts ?? 0),
138      entities: Number(memory.entities ?? 0),
139    }
140    await update($, workspaceAtom, () => workspace)
141    $.ui.status(status ? `brain · ${workspace.name} · ${workspace.facts} facts` : 'brain · connected')
142    return next(e)
143  })
144
145  on('prompt.submit', async ($, e, next) => {
146    if (opts.recall_each_prompt === false || e.command || e.isMeta) return next(e)
147    if (e.text.trim().length < MIN_PROMPT_CHARS) return next(e)
148    const recall = await search($, at(), e.text, next.signal)
149    await update($, recallAtom, () => recall)
150    if (recall.hits.length === 0) return next(e)
151    return next({ ...e, context: [...(e.context ?? []), contextOf(recall.hits)] })
152  })
153    // Memory is an extra, never a gate: a recall that throws or overruns
154    // lets the prompt through as typed (replayed if it already went on).
155    .catch(($, e, next) => next(e))
156
157  on('command.run', { command: 'brain' }, async ($, e) => {
158    const question = e.args.trim()
159    if (question) await update($, foundAtom, () => null)
160    await $.ui.open({ id: PANE, title: 'Brain' })
161    if (question) {
162      const found = await search($, at(), question)
163      await update($, foundAtom, () => found)
164      return { text: `Brain found ${found.hits.length} for "${question}".` }
165    }
166    return { text: 'Brain pane opened.' }
167  })
168
169  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
170    const recall = await read($, recallAtom)
171    if (e.props.hasSurvey || recall === null || recall.hits.length === 0 || (await read($, hiddenAtom))) {
172      return next(e)
173    }
174    const { Box, Button, Text } = $.ui.resolve(e)
175    const names = recall.hits.map(hit => hit.name).join(', ')
176    return (
177      <Box>
178        <Text dimColor wrap="truncate-end">
179          brain recalled {recall.hits.length}: {names}{' '}
180        </Text>
181        <Button key="open" plain label="[open]" onPress={() => void $.ui.open({ id: PANE, title: 'Brain' })} />
182        <Text> </Text>
183        <Button key="hide" plain label="[hide]" onPress={() => update($, hiddenAtom, () => true)} />
184      </Box>
185    )
186  })
187
188  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
189    const elements = $.ui.resolve(e)
190    const { Box, Text } = elements
191    // Mobile draws no text field: the pane there shows the last recall only.
192    const Input = 'Input' in elements ? elements.Input : undefined
193    const workspace = await read($, workspaceAtom)
194    const found = (await read($, foundAtom)) ?? (await read($, recallAtom))
195    const heading = workspace
196      ? `${workspace.name} · ${workspace.entities} entities · ${workspace.facts} facts`
197      : server === null
198        ? 'Not connected — run /mcp to sign in to brain.'
199        : 'Connected.'
200    const ask = (value: string) => {
201      const query = value.trim()
202      if (!query) return
203      void update($, foundAtom, () => ({ query, hits: [], ms: 0 })).then(async () => {
204        const result = await search($, at(), query)
205        await update($, foundAtom, () => result)
206      })
207    }
208    return (
209      <Box flexDirection="column">
210        <Text dimColor>{heading}</Text>
211        {Input && <Input key="ask" label="Ask brain" placeholder="what do we know about…" onSubmit={ask} />}
212        {found === null && <Text dimColor>Nothing asked yet.</Text>}
213        {found !== null && (
214          <Text bold>
215            {found.query.length > 60 ? `${found.query.slice(0, 60)}…` : found.query}
216            {found.failure ? ` — ${found.failure}` : ` — ${found.hits.length} in ${Math.round(found.ms)} ms`}
217          </Text>
218        )}
219        {found?.hits.map(hit => (
220          <Box flexDirection="column" key={hit.name}>
221            <Text>{hit.name}</Text>
222            {hit.facts.map(fact => (
223              <Text dimColor>  {fact}</Text>
224            ))}
225          </Box>
226        ))}
227      </Box>
228    )
229  })
230}
231

1/** One entity brain returned, with the facts the recall shows of it. */
2export type BrainHit = { name: string; facts: string[] }
3
4/** What brain recalled for the last prompt, or why it recalled nothing. */
5export type BrainRecall = {
6  query: string
7  hits: BrainHit[]
8  ms: number
9  /** Set when the recall did not come back: `timeout`, `error`, `offline`. */
10  failure?: string
11}
12
13/** The workspace this connection writes to, from `workspace_status`. */
14export type BrainWorkspace = { name: string; facts: number; entities: number }
15
16declare module 'claude-code' {
17  interface PluginState {
18    'inite-brain': {
19      recall: BrainRecall | null
20      found: BrainRecall | null
21      workspace: BrainWorkspace | null
22      isBandHidden: boolean
23    }
24  }
25}
26