memrove-chat

1import type { EngineInterface, Register } from 'claude-code'
2import {
3  BACKOFF_MS,
4  configFrom,
5  failureText,
6  freshMessages,
7  inBackoff,
8  parseApplied,
9  parseBatch,
10  parseStatus,
11  shortenInstruction,
12  statusLine,
13  viewBlock,
14  withoutView,
15  type Config,
16} from './chat.ts'
17
18/** The mod is transport only: every memory decision lives in the memrove CLI. */
19type State = {
20  running: boolean
21  pending: boolean
22  failures: number
23  lastFailureAt: number
24  nodes?: number
25  calls?: number
26  cap?: number
27  /** The view SessionStart read, for the prompt.context right after it (one read per start). */
28  view?: { at: number; read: Promise<string | undefined> }
29}
30
31const ZOOM_TOOL = {
32  name: 'zoom',
33  description:
34    'Open line id+n of the chat memory view into the two lines of n/2 under it; n = 1 gives message id whole (a long message comes in pages). n is a power of 2 and id a multiple of n. Use it to find the latest mention of anything that may have come up before, before you act, guess or ask.',
35  inputSchema: {
36    type: 'object',
37    properties: {
38      id: { type: ['integer', 'string'], minimum: 0, description: 'The first message of the line (the number before the +); or `rec:<id>` of a saved record from a search hit, which opens the record with its why (n = 1).' },
39      n: { type: 'integer', minimum: 1, description: 'How many messages the line covers (the number after the +); a power of 2.' },
40      page: { type: 'integer', minimum: 1, description: 'Page of a long message (n = 1 only); default 1.' },
41    },
42    required: ['id', 'n'],
43  },
44} as const
45
46const DATE_TOOL = {
47  name: 'date',
48  description: 'The date and time of message id of the chat memory.',
49  inputSchema: { type: 'object', properties: { id: { type: 'integer', minimum: 0 } }, required: ['id'] },
50} as const
51
52const SEARCH_TOOL = {
53  name: 'search',
54  description: 'Search every message of the chat memory, and the saved decisions, lessons and notes, by their words. Returns the best hits as id|date|kind|snippet (a saved record is rec:<id>); a hit is only a snippet that drops detail, so open the message with zoom(id, 1) before you answer from it. Optional since/until limit it to a date range. Use it when the view shows a topic only as a word, or to find where something was said.',
55  inputSchema: {
56    type: 'object',
57    properties: {
58      query: { type: 'string', description: 'A few words the message would contain. May be empty when since or until is given: then the messages of that range are listed, newest first.' },
59      since: { type: 'string', description: 'Only messages from this date on: YYYY-MM-DD or an ISO datetime. Turn words like "last week" into dates yourself.' },
60      until: { type: 'string', description: 'Only messages up to this date, inclusive: YYYY-MM-DD or an ISO datetime.' },
61    },
62    required: ['query'],
63  },
64} as const
65
66const PROJECT_TOOL = {
67  name: 'project',
68  description:
69    'Where a project stands, in a few lines: its latest decisions (the ones a later decision replaced are marked SUPERSEDED), open loops, and the end state of its last session. Every line starts with an id that zoom opens (rec:<id>, th:<id>, ls:<id>). Call it when you start work on a project and need its current state, before you ask the owner to re-explain. Without a name it uses the project of the working directory.',
70  inputSchema: { type: 'object', properties: { name: { type: 'string', description: 'Project name; leave out to use the working directory.' } } },
71} as const
72
73const DOC_TOOL = {
74  name: 'doc',
75  description:
76    'Open a docs write-up that a search doc handle points to (the `doc|date|path|heading|snippet` lines under the hits), whole or one section. Only files under docs/state, docs/design, docs/research and docs/decisions of the owner\'s repos open.',
77  inputSchema: {
78    type: 'object',
79    properties: {
80      path: { type: 'string', description: 'The absolute path from the doc handle.' },
81      section: { type: 'string', description: 'Heading of one section; leave out for the whole file.' },
82    },
83    required: ['path'],
84  },
85} as const
86
87const MODEL_TOKENS = 16_000
88const MODEL_TIMEOUT_MS = 300_000
89const CLI_TIMEOUT_MS = 120_000
90const VIEW_HANDOFF_MS = 30_000
91
92function message(err: unknown): string {
93  return err instanceof Error ? err.message : String(err)
94}
95
96/** Runs the memrove CLI; a non-zero exit throws with the end of its stderr. */
97async function cli($: EngineInterface, cfg: Config, args: string[], stdin?: string): Promise<string> {
98  const r = await $.process.run([...cfg.memrove, ...args], { timeoutMs: CLI_TIMEOUT_MS, ...(stdin === undefined ? {} : { stdin }) })
99  if (r.exitCode !== 0) throw new Error(`memrove ${args.join(' ')} exited ${r.exitCode}: ${(r.stderr || r.stdout).trim().slice(-300)}`)
100  return r.stdout
101}
102
103function show($: EngineInterface, text: string): void {
104  $.ui.status(text)
105}
106
107async function refreshStatus($: EngineInterface, state: State, cfg: Config): Promise<void> {
108  const s = parseStatus(await cli($, cfg, ['chatlog', 'status']))
109  state.nodes = s.nodes
110  state.calls = s.calls_today
111  state.cap = s.cap
112}
113
114/** One batch: prepare, one Haiku call on the session's own auth, apply. Returns true when the queue is empty. */
115async function oneBatch($: EngineInterface, state: State, cfg: Config, n: number): Promise<boolean> {
116  const batch = parseBatch(await cli($, cfg, ['chatlog', 'prepare', '--json', '--no-ingest']))
117  if (batch === undefined) return true
118  show($, statusLine({ building: true, batch: n, nodes: state.nodes, calls: state.calls, cap: state.cap }))
119  const r = await $.model.complete({ model: cfg.model, system: batch.system, prompt: batch.prompt, maxTokens: MODEL_TOKENS, timeoutMs: MODEL_TIMEOUT_MS })
120  if (!r.isAnswered) {
121    // Hands the claim back so it does not wait out its expiry; the CLI decides what an error means.
122    const why = failureText(r as { reason: string; status?: number | null; error?: string })
123    await cli($, cfg, ['chatlog', 'apply', batch.batch_id, '--error', why]).catch(() => undefined)
124    throw new Error(`the model gave no text (${why})`)
125  }
126  const done = parseApplied(await cli($, cfg, ['chatlog', 'apply', batch.batch_id], r.text))
127  if (done.error !== undefined) throw new Error(`apply: ${done.error}`)
128  return done.done
129}
130
131/** Ingest, then up to cfg.batches batches. */
132async function run($: EngineInterface, state: State, cfg: Config): Promise<void> {
133  show($, statusLine({ building: true, nodes: state.nodes, calls: state.calls, cap: state.cap }))
134  await cli($, cfg, ['chatlog', 'ingest'])
135  for (let n = 1; n <= cfg.batches; n++) {
136    if (await oneBatch($, state, cfg, n)) break
137    await refreshStatus($, state, cfg)
138  }
139  await refreshStatus($, state, cfg)
140}
141
142/** One run at a time; a turn that ends during a run asks for one more after it (as memory-save does). */
143async function drain($: EngineInterface, state: State, cfg: Config): Promise<void> {
144  if (state.running) {
145    state.pending = true
146    return
147  }
148  state.running = true
149  try {
150    do {
151      state.pending = false
152      if (inBackoff(state.failures, state.lastFailureAt, await $.clock.now())) {
153        show($, statusLine({ error: `paused ${Math.round(BACKOFF_MS / 60_000)} min after ${state.failures} failed runs` }))
154        break
155      }
156      try {
157        await run($, state, cfg)
158        state.failures = 0
159        show($, statusLine({ nodes: state.nodes, calls: state.calls, cap: state.cap }))
160      } catch (err) {
161        state.failures += 1
162        state.lastFailureAt = await $.clock.now()
163        show($, statusLine({ error: message(err) }))
164      }
165    } while (state.pending)
166  } finally {
167    state.running = false
168  }
169}
170
171/** The whole view as the CLI prints it, or undefined when there is no chat memory yet. */
172async function readView($: EngineInterface, cfg: Config): Promise<string | undefined> {
173  try {
174    const text = (await cli($, cfg, ['wake', '--all'])).trim()
175    return text === '' ? undefined : text
176  } catch (err) {
177    $.ui.log(`memrove-chat: the view was not read: ${message(err)}`, { to: 'debug' })
178    return undefined
179  }
180}
181
182async function toolText($: EngineInterface, cfg: Config, args: string[]): Promise<{ result: string }> {
183  const r = await $.process.run([...cfg.memrove, ...args], { timeoutMs: CLI_TIMEOUT_MS })
184  return { result: r.exitCode === 0 ? r.stdout.trim() : `error: ${(r.stderr || r.stdout).trim()}` }
185}
186
187function intOf(v: unknown): number | undefined {
188  return typeof v === 'number' && Number.isInteger(v) ? v : undefined
189}
190
191export const register: Register = (on, options) => {
192  const cfg = configFrom(options as Record<string, unknown>)
193  const state: State = { running: false, pending: false, failures: 0, lastFailureAt: 0 }
194
195  on('session.start', async ($, e, next) => {
196    const r = await next(e)
197    await $.tool.register(ZOOM_TOOL)
198    await $.tool.register(DATE_TOOL)
199    await $.tool.register(SEARCH_TOOL)
200    await $.tool.register(PROJECT_TOOL)
201    await $.tool.register(DOC_TOOL)
202    return r
203  })
204
205  // A plugin's tool waits behind ToolSearch by default; these are listed, so the model calls them at once.
206  on('tool.describe', { tool: /^mcp__memrove-chat__(zoom|date|search|project|doc)$/ }, async (_, e, next) => ({ ...(await next(e)), isDeferred: false }))
207
208  on('tool.call', { tool: /^mcp__memrove-chat__zoom$/ }, async ($, e) => {
209    const input = e as Record<string, unknown>
210    if (typeof input.id === 'string' && /^(rec|th|ls):[\w:-]+$/.test(input.id.trim())) return toolText($, cfg, ['show', input.id.trim()])
211    const id = intOf(input.id)
212    const n = intOf(input.n)
213    if (id === undefined || n === undefined) return { deny: 'zoom needs integer id and n' }
214    const page = intOf(input.page)
215    return toolText($, cfg, ['zoom', String(id), String(n), ...(page === undefined ? [] : ['--page', String(page)])])
216  })
217
218  on('tool.call', { tool: /^mcp__memrove-chat__date$/ }, async ($, e) => {
219    const id = intOf((e as Record<string, unknown>).id)
220    return id === undefined ? { deny: 'date needs an integer id' } : toolText($, cfg, ['date', String(id)])
221  })
222
223  on('tool.call', { tool: /^mcp__memrove-chat__search$/ }, async ($, e) => {
224    const { query, since, until } = e as Record<string, unknown>
225    const range = [['--since', since], ['--until', until]].flatMap(([flag, v]) => (typeof v === 'string' && v.trim() !== '' ? [flag as string, v.trim()] : []))
226    const text = typeof query === 'string' ? query.trim() : ''
227    return text !== '' || range.length > 0 ? toolText($, cfg, ['chatlog', 'search', text, ...range]) : { deny: 'search needs a query or a since/until date' }
228  })
229
230  on('tool.call', { tool: /^mcp__memrove-chat__project$/ }, async ($, e) => {
231    const name = (e as Record<string, unknown>).name
232    return toolText($, cfg, ['project', ...(typeof name === 'string' && name.trim() !== '' ? [name.trim()] : [])])
233  })
234
235  on('tool.call', { tool: /^mcp__memrove-chat__doc$/ }, async ($, e) => {
236    const { path, section } = e as Record<string, unknown>
237    if (typeof path !== 'string' || path.trim() === '') return { deny: 'doc needs a path' }
238    return toolText($, cfg, ['doc', path.trim(), ...(typeof section === 'string' && section.trim() !== '' ? ['--section', section.trim()] : [])])
239  })
240
241  // The view joins the first user message as a context block of its own: at startup, resume, /clear and after a compaction.
242  on('prompt.context', async ($, e, next) => {
243    const r = await next(e)
244    // Takes the read SessionStart just made (a few seconds old at most); /clear and compaction have none and read again.
245    const held = state.view
246    state.view = undefined
247    const view = await (held !== undefined && (await $.clock.now()) - held.at < VIEW_HANDOFF_MS ? held.read : readView($, cfg))
248    return view === undefined ? r : { ...r, blocks: [...r.blocks, { name: 'memroveChat', text: viewBlock(view) }] }
249  })
250
251  // Subagents inherit the session's context sections (measured: +33k tokens each); the view stays with the main loop.
252  on('prompt.attachment', { type: 'context_sections' }, async (_, e, next) => {
253    const r = await next(e)
254    if (e.agentId === undefined || r.text === null || r.text === undefined) return r
255    return { ...r, text: withoutView(r.text) }
256  })
257
258  // The SessionStart instruction of the memrove hook shrinks to one line once the view is in the context.
259  on('classic.SessionStart', async ($, e, next) => {
260    const r = await next(e)
261    if (e.agent_id !== undefined || r.additionalContext === undefined) return r
262    // SessionStart runs before prompt.context, so it reads the view itself and hands the same read on to that block.
263    const read = readView($, cfg)
264    state.view = { at: await $.clock.now(), read }
265    if ((await read) === undefined) return r
266    return { ...r, additionalContext: shortenInstruction(r.additionalContext) }
267  })
268
269  // The fresh context: the view and the last exchange, nothing else. Off unless the `fresh` option is on.
270  on('session.compact', async ($, e, next) => {
271    if (!cfg.fresh || e.agentId !== undefined) return next(e)
272    const view = await readView($, cfg)
273    const messages = view === undefined ? undefined : freshMessages(view, e.messages)
274    return messages === undefined ? next(e) : { messages }
275  })
276
277  on('turn.complete', async ($, e, next) => {
278    const r = await next(e)
279    if (e.agentId !== undefined || e.reason !== 'answer') return r
280    // On a timer, not in the dispatch: a call still in flight when the dispatch closes is aborted.
281    $.clock.after(0, () => void drain($, state, cfg))
282    if (cfg.fresh) {
283      // A compaction started from a timer skips this plugin's own hook, so it starts here, not awaited.
284      void $.session.compact().catch((err: unknown) => $.ui.log(`memrove-chat: fresh compaction failed: ${message(err)}`))
285    }
286    return r
287  })
288}
289

1/**
2 * The pure part of memrove-chat: the options, what the CLI prints, the status line and the message
3 * list of a fresh context. No `$` here, so the tests need no engine.
4 */
5import type { SessionMessage } from 'claude-code'
6
7export type Config = {
8  /** argv of the memrove entry point, split on spaces. */
9  memrove: string[]
10  /** Batches (one model call each) after a turn. */
11  batches: number
12  model: string
13  /** Compact to the view plus the last exchange after every turn. */
14  fresh: boolean
15}
16
17export const DEFAULTS = { memrove: 'memrove', batches: 3, model: 'haiku', fresh: false }
18
19export function splitCommand(text: string): string[] {
20  return text.split(/\s+/).filter(p => p !== '')
21}
22
23export function configFrom(options: Record<string, unknown>): Config {
24  const memrove = typeof options.memrove === 'string' && options.memrove.trim() !== '' ? options.memrove : DEFAULTS.memrove
25  const b = options.batches
26  const batches = typeof b === 'number' && Number.isInteger(b) && b >= 1 && b <= 20 ? b : DEFAULTS.batches
27  const model = typeof options.model === 'string' && options.model.trim() !== '' ? options.model.trim() : DEFAULTS.model
28  return { memrove: splitCommand(memrove), batches, model, fresh: options.fresh === true || options.fresh === 'true' }
29}
30
31export type Batch = { batch_id: string; system: string; prompt: string; node_ids: string[] }
32export type Applied = { applied: number; retry_ids: string[]; done: boolean; error?: string }
33export type Status = { nodes: number; calls_today: number; cap: number; messages: number }
34
35function parseObject(text: string): Record<string, unknown> | undefined {
36  try {
37    const v: unknown = JSON.parse(text)
38    return typeof v === 'object' && v !== null && !Array.isArray(v) ? (v as Record<string, unknown>) : undefined
39  } catch {
40    return undefined
41  }
42}
43
44/** `chatlog prepare --json`: a batch, or undefined for `{}` (nothing ready, locked, or the cap is spent). */
45export function parseBatch(text: string): Batch | undefined {
46  const o = parseObject(text)
47  if (o === undefined) throw new Error(`prepare printed no JSON: ${text.slice(0, 120)}`)
48  if (typeof o.batch_id !== 'string') return undefined
49  if (typeof o.system !== 'string' || typeof o.prompt !== 'string') throw new Error('prepare printed a batch without system or prompt')
50  return { batch_id: o.batch_id, system: o.system, prompt: o.prompt, node_ids: Array.isArray(o.node_ids) ? o.node_ids.map(String) : [] }
51}
52
53export function parseApplied(text: string): Applied {
54  const o = parseObject(text)
55  if (o === undefined || typeof o.applied !== 'number') throw new Error(`apply printed no result: ${text.slice(0, 120)}`)
56  return {
57    applied: o.applied,
58    retry_ids: Array.isArray(o.retry_ids) ? o.retry_ids.map(String) : [],
59    done: o.done === true,
60    ...(typeof o.error === 'string' ? { error: o.error } : {}),
61  }
62}
63
64export function parseStatus(text: string): Status {
65  const o = parseObject(text)
66  const n = (k: string): number => (typeof o?.[k] === 'number' ? o[k] : 0)
67  if (o === undefined) throw new Error(`status printed no JSON: ${text.slice(0, 120)}`)
68  return { nodes: n('nodes'), calls_today: n('calls_today'), cap: n('cap'), messages: n('messages') }
69}
70
71/** Why a model call gave no text, as one short phrase for `apply --error` and the status line. */
72export function failureText(r: { reason: string; status?: number | null; error?: string }): string {
73  if (r.reason !== 'api-error') return r.reason
74  return `api-error${r.status == null ? '' : ` ${r.status}`}${r.error === undefined ? '' : ` ${r.error}`}`
75}
76
77/** The status line: building, nodes done and calls made today, or the error. */
78export function statusLine(s: { building?: boolean; nodes?: number; calls?: number; cap?: number; error?: string; batch?: number }): string {
79  if (s.error !== undefined) return `chat memory: ${s.error}`
80  const parts = [s.building ? `building${s.batch === undefined ? '' : ` batch ${s.batch}`}` : 'idle']
81  if (s.nodes !== undefined) parts.push(`${s.nodes} nodes`)
82  if (s.calls !== undefined) parts.push(`${s.calls}${s.cap === undefined ? '' : `/${s.cap}`} calls today`)
83  return `chat memory: ${parts.join(' · ')}`
84}
85
86/** After a run of failures the turns stop asking for a batch for a while, so a poisoned node does not cost 3 calls a turn. */
87export const BACKOFF_FAILURES = 3
88export const BACKOFF_MS = 10 * 60_000
89
90export function inBackoff(failures: number, lastFailureAt: number, now: number): boolean {
91  return failures >= BACKOFF_FAILURES && now - lastFailureAt < BACKOFF_MS
92}
93
94/** What the classic SessionStart instruction becomes when the view itself is in the context. */
95export const ONE_LINE = 'Before you answer or start a task that may build on earlier work, search first; a hit or a doc line is a lead, not evidence; open it with zoom or memrove doc and answer from what you opened. On one topic the newer statement wins, and the user\'s word beats an agent\'s. Tools: search, zoom (id, n), date.'
96const INSTRUCTION_MARK = 'Memrove chat memory is on'
97
98export function shortenInstruction(texts: readonly string[] | undefined): string[] | undefined {
99  if (texts === undefined) return undefined
100  const out: string[] = []
101  for (const t of texts) {
102    if (!t.includes(INSTRUCTION_MARK)) out.push(t)
103    else if (!out.includes(ONE_LINE)) out.push(ONE_LINE)
104  }
105  return out
106}
107
108export const VIEW_NOTE = 'This is an index of the whole chat so far, one line per run of messages (id+n|summary), not evidence: the newest lines are the most exact, older ones keep little. Open a line with the zoom tool before you rely on it.'
109
110/** The block text of the view. */
111export function viewBlock(view: string): string {
112  return `${VIEW_NOTE}\n${view}`
113}
114
115/**
116 * The text of a `context_sections` attachment without the view: subagents get the context blocks of the
117 * session, and the view (about 32k tokens) is for the main loop only. Sections are `# name` headers; the
118 * view's lines start with a number, so the next core header ends it. Empty text comes back as null.
119 */
120export function withoutView(text: string): string | null {
121  const head = '# memroveChat\n'
122  const at = text.startsWith(head) ? 0 : text.indexOf(`\n${head}`) + 1
123  if (at === 0 && !text.startsWith(head)) return text
124  const end = text.slice(at).search(/\n# (claudeMd|userEmail|attachedProject|currentDate)\n/)
125  const rest = end < 0 ? text.slice(0, at).trimEnd() : text.slice(0, at) + text.slice(at + end + 1)
126  return rest.trim() === '' ? null : rest
127}
128
129/**
130 * The messages of a fresh context: the last user message and everything after it verbatim (its tool
131 * calls and answer); the view comes back as the prompt.context block. Undefined when there is no view or
132 * user message with text, so the compaction falls back to the engine's.
133 */
134export function freshMessages(view: string, messages: readonly SessionMessage[]): SessionMessage[] | undefined {
135  let last = -1
136  messages.forEach((m, i) => {
137    if (m.role === 'user' && m.text.trim() !== '' && (m.toolResults ?? []).length === 0) last = i
138  })
139  const user = messages[last]
140  if (user === undefined || view.trim() === '') return undefined
141  // The view itself is not put in: prompt.context adds it again after every compaction (measured: both copies cost ~2x the view per turn).
142  return messages.slice(last)
143}
144