open-science-context

1// open-science context management inside Claude Code (a Claude Code mod).
2//
3// It does from inside Claude Code what the tmux path does by typing into the pane: clear
4// the session and run the resume prompt (session jumps), wake a waiting session (queued
5// SLURM wakers, the cache-cold notice), rename the session after its task, and read the
6// context size, also each subagent's. It also shows the context size and how long the
7// prompt cache stays warm (a line above the prompt, on the terminal and in the Desktop app;
8// Remote Control shows no mod's drawing), asks the user before a prompt they send to a cold
9// cache, and shows the jump report and a cache-cold note as rows of the conversation, which
10// Remote Control shows too (/opsci-note). The policy and the records stay in the plugin's
11// shell scripts, shared with the tmux path and Codex: `cm_stop.sh --mod` decides at each
12// stop, `cm_mod.sh` keeps the records, `wait_slurm.sh --check` polls a waker,
13// `session_name.sh want` names the session.
14//
15// Loading the mod sets OPSCI_MOD=1 for Claude Code and everything it starts, which turns
16// the plain Stop and naming hooks off and makes jump.sh and wait_slurm.sh key their
17// records by session id. A Claude Code that does not load the mod never sets it, and the
18// tmux path runs as before.
19
20const SUB_LIMIT_DEFAULT = 200000
21const POLL_DEFAULT_S = 60
22// The bar and the question call the cache cold from here (OPSCI_CACHE_TTL_MIN). The cache
23// lives 60 min from the start of the last request that read it; the bar's clock starts at
24// that request's end, so a minute less covers most responses. (The cache-cold notice to the
25// agent comes a minute earlier still, so its turn starts warm: OPSCI_CACHE_COLD_MIN, 58, in
26// cm_lib.sh.)
27const TTL_MIN_DEFAULT = 59
28const BAR_TICK_MS = 30000
29const STORE_KEY = 'lastRequest' // { <session id>: ms of the main agent's last model request }
30
31let sid = null              // the session id this process runs now
32let busy = false            // a main-agent turn is running
33let jumping = false         // a clear and resume is in progress
34let waiting = false         // a wait jump left this session waiting for a waker
35let clearedFrom = null      // the id a /clear the user typed left behind
36let atStop = null           // what cm_stop.sh --mod decided, acted on at turn.complete
37let coldTimer = null
38let pollTimer = null
39let polling = false
40let lastName = null
41let barSid = null           // the session lastReq belongs to
42let lastReq = null          // ms of the main agent's last model request in barSid; null: none yet
43let barText = ''            // what the bar shows now
44let barPhase = 'none'       // none | warm | cold: the bar's color
45let lastTokens = null       // the context after the main agent's last response, in barSid
46let barTimer = null
47const warned = new Set()    // subagents already told to checkpoint
48
49const key = s => String(s).replace(/[^A-Za-z0-9._-]/g, '_')
50
51async function stateDir($) {
52  const d = await $.env.get('OPSCI_STATE_DIR')
53  if (d) return d
54  const x = await $.env.get('XDG_STATE_HOME')
55  return (x || (await $.env.get('HOME')) + '/.local/state') + '/open-science'
56}
57
58async function sh($, script, args, stdin, timeoutMs) {
59  try {
60    return await $.process.run(['bash', $.plugin.root + '/scripts/' + script, ...args],
61      { stdin: stdin || '', timeoutMs: timeoutMs || 30000 })
62  } catch (err) {
63    return { exitCode: -1, stdout: '', stderr: String(err) }
64  }
65}
66
67const log = ($, text) => sh($, 'cm_mod.sh', ['log', text])
68
69// The resume command, built here from the context file: no prompt or command is ever taken
70// from a record in the state directory. The file must be an absolute path with no
71// whitespace or control character, to a file that exists; otherwise continue-context runs
72// with no file and finds the session's registration.
73const RESUME_CMD = 'open-science-context:continue-context'
74const ctxOk = c => typeof c === 'string' && /^\/[^\s\x00-\x1f\x7f-\x9f]+$/.test(c)
75
76async function resume($, ctx) {
77  const ok = ctxOk(ctx) && (await $.fs.exists(ctx))
78  if (ctx && !ok) await log($, 'resume: context path refused; continuing from the registration')
79  return $.command.run({ command: RESUME_CMD, args: ok ? ctx : '' })
80}
81
82// The records cm_mod.sh and pane_context.sh keep, by session id.
83const record = async ($, ...parts) => [await stateDir($), ...parts].join('/')
84
85async function registered($, s) {
86  return $.fs.exists(await record($, 'session_context', 'claude__' + key(s) + '.json'))
87}
88
89async function rename($, always) {
90  const r = await sh($, 'session_name.sh', ['want', ...(always ? ['--always'] : []), sid, await $.session.cwd()])
91  const name = r.stdout.trim()
92  if (!name || (name === lastName && !always)) return
93  lastName = name
94  await $.command.run({ command: 'rename', args: name })
95}
96
97// ---- notes in the conversation ---------------------------------------------------------
98// A note the user sees: the output row of the mod's /opsci-note command, the one kind of row
99// a mod adds that Remote Control shows too (a notice row shows nowhere). It starts no turn;
100// the model reads it with the next prompt, as it reads any command's output. Typed by the
101// user, /opsci-note shows the last note again.
102//
103// showNote is called from inside a hook only: a command the mod runs from a timer skips the
104// mod's own command.run hook (Claude Code 2.1.287), and Claude Code answers it instead.
105const NOTE_CMD = 'opsci-note'
106let noteText = null         // what the next /opsci-note shows
107let lastNote = ''           // the last note shown, for /opsci-note typed by the user
108let pendingJump = null      // the jump's note, shown when the clear starts the new session
109let pendingCold = null      // the cache-cold note, shown at the next draw of the bar
110
111function showNote($, text) {
112  noteText = text
113  lastNote = text
114  $.command.run({ command: NOTE_CMD, args: '' }).catch(err => log($, 'note FAILED: ' + err))
115}
116
117// The jump's report (jump.sh --report), shown at the top of the cleared session. For a wait
118// jump it first says that the session is cleared and waiting.
119function jumpNote(j) {
120  const report = String(j.report || '').trim()
121  if (j.kind === 'active') return report
122  const n = Number(j.wakers) || 0
123  const what = n ? n + ' running background task' + (n === 1 ? '' : 's') + ', cron' + (n === 1 ? '' : 's') + ' or SLURM waker' + (n === 1 ? '' : 's') : 'its background work'
124  return 'Wait jump: this session was cleared and is waiting for ' + what +
125    '. It resumes by itself when that work reports back; until then it does nothing.' + (report ? '\n\n' + report : '')
126}
127
128// The clear did not start the new session through SessionStart, so no note was shown: the
129// agent still reads the report, as a user-role row the user does not see.
130async function jumpNoteFallback($, text) {
131  await log($, 'jump note not shown: no SessionStart after the clear; stored for the agent')
132  try {
133    await $.session.append({ message: { type: 'user', content: [{ type: 'text', text: '[open-science] The previous session wrote this when it jumped:\n\n' + text }] } })
134  } catch (err) { await log($, 'jump note FAILED: ' + err) }
135}
136
137// Clear the session with Claude Code's own /clear, hand the registration and wakers to
138// the new session, show the jump's report at its top (from the SessionStart the clear
139// fires), then run the resume prompt (active) or leave it waiting (wait).
140function jump($, j) {
141  jumping = true
142  coldTimer?.cancel(); coldTimer = null
143  $.clock.after(0, async () => {
144    try {
145      const old = await $.session.id()
146      pendingJump = jumpNote(j) || null
147      await $.command.run({ command: 'clear', args: '' })
148      const now = await $.session.id()
149      if (now === old) {
150        pendingJump = null
151        await log($, 'jump FAILED: /clear did not start a new session (still ' + old.slice(0, 8) + ')')
152        void $.prompt.submit({ text: '[open-science] session jump FAILED before anything was changed; this session keeps running: the clear did not happen.' })
153        return
154      }
155      sid = now
156      await sh($, 'cm_mod.sh', ['handover', old, now])
157      await log($, j.kind + ' jump: cleared ' + old.slice(0, 8) + ' -> ' + now.slice(0, 8))
158      if (pendingJump) { const text = pendingJump; pendingJump = null; await jumpNoteFallback($, text) }
159      if (j.kind === 'active') {
160        await resume($, j.context)
161      } else {
162        waiting = true
163        await sh($, 'cm_mod.sh', ['waiting', now, j.context || ''])
164      }
165      await rename($)
166    } catch (err) {
167      await log($, 'jump FAILED: ' + err)
168    } finally {
169      jumping = false
170    }
171  })
172}
173
174// One cache-cold timer per process. It fires only into the session it was armed for,
175// with no turn and no jump since.
176function armCold($, secs, notice) {
177  coldTimer?.cancel(); coldTimer = null
178  if (!secs || !notice) return
179  const armedFor = sid
180  coldTimer = $.clock.after(secs * 1000, async () => {
181    coldTimer = null
182    if (busy || jumping || (await $.session.id()) !== armedFor) return
183    await log($, 'cache-cold notice sent to ' + armedFor.slice(0, 8))
184    void $.prompt.submit({ text: notice })
185  })
186}
187
188// Queued SLURM wakers of this session: when their jobs have left the queue, send the
189// report (the same text a Codex waker sends), which starts a turn once the session is idle.
190// The reports of all wakers done in one poll go in one message, each distinct report once,
191// so wakers queued twice for the same jobs wake the session once.
192async function pollWakers($) {
193  if (polling) return
194  polling = true
195  try {
196    const cur = await $.session.id()
197    if (!(await $.fs.exists(await record($, 'wakers', 'claude-sid__' + key(cur))))) return
198    const files = (await sh($, 'cm_mod.sh', ['wakers', cur])).stdout.split('\n').filter(Boolean)
199    const reports = new Set()
200    for (const f of files) {
201      const r = await sh($, 'wait_slurm.sh', ['--check', f], '', 120000)
202      if (r.exitCode === 0 || r.exitCode === 1) {
203        await log($, 'waker ' + f.split('/').pop() + ': jobs left the queue; report sent')
204        reports.add(r.stdout.trim().replace(/\s*\n\s*/g, ' '))
205      } else if (r.exitCode !== 10) {
206        await log($, 'waker ' + f + ': check failed (' + r.exitCode + '): ' + r.stderr.trim())
207      }
208    }
209    if (reports.size) void $.prompt.submit({ text: '[open-science] ' + [...reports].join(' ') })
210  } finally {
211    polling = false
212  }
213}
214
215// ---- the context bar and the cold-cache question --------------------------------------
216// The prompt cache lives TTL minutes from the last request that read it, so the clock runs
217// from the main agent's last model request (each turn.step), not from the last prompt.
218
219async function ttlMin($) {
220  return Number(await $.env.get('OPSCI_CACHE_TTL_MIN')) || TTL_MIN_DEFAULT
221}
222
223// Follows the session through clears and resumes: a new session id reads its own record
224// (a cleared session has none: no cache yet; a resumed one has its last request's time).
225async function syncSid($) {
226  const now = await $.session.id()
227  if (now === barSid) return
228  barSid = now
229  lastTokens = null
230  const all = (await $.store.get(STORE_KEY)) || {}
231  lastReq = typeof all[now] === 'number' ? all[now] : null
232}
233
234// At the end of each main-agent request: its time, and the context it leaves (what it read
235// plus what it wrote, which the next request reads).
236async function noteRequest($, usage) {
237  await syncSid($)
238  lastReq = await $.clock.now()
239  if (usage) lastTokens = (usage.input_tokens || 0) + (usage.cache_read_input_tokens || 0) +
240    (usage.cache_creation_input_tokens || 0) + (usage.output_tokens || 0)
241  const all = { ...((await $.store.get(STORE_KEY)) || {}) }
242  delete all[barSid]
243  all[barSid] = lastReq
244  const keys = Object.keys(all)
245  for (const k of keys.slice(0, Math.max(0, keys.length - 50))) delete all[k]
246  await $.store.set(STORE_KEY, all)
247}
248
249// { phase: none | warm | cold, idleMin, leftMin, tokens }
250async function cacheState($) {
251  await syncSid($)
252  const ttl = await ttlMin($)
253  let tokens = lastTokens
254  if (tokens === null) {
255    try { tokens = (await $.session.usage()).context.tokens } catch { tokens = undefined }
256  }
257  if (lastReq === null) return { phase: 'none', tokens }
258  if (busy) return { phase: 'warm', idleMin: 0, leftMin: ttl, tokens }
259  const idleMin = Math.max(0, ((await $.clock.now()) - lastReq) / 60000)
260  return { phase: idleMin >= ttl ? 'cold' : 'warm', idleMin, leftMin: Math.max(0, Math.ceil(ttl - idleMin)), tokens }
261}
262
263// 987, then 1.0k, 123.4k
264const fmtTokens = n => n >= 1000 ? (n / 1000).toFixed(1) + 'k' : String(n)
265
266function describe(st) {
267  const ctx = typeof st.tokens === 'number' ? 'context ' + fmtTokens(st.tokens) + ' tokens' : 'context size unknown'
268  if (st.phase === 'cold') return '⚠ ' + ctx + ' · cache cold (' + Math.floor(st.idleMin) + ' min idle)'
269  return ctx + ' · ' + (st.phase === 'warm' ? 'cache warm, ' + st.leftMin + ' min left' : 'no cache yet')
270}
271
272const BAR_COLOR = { warm: 'green', cold: 'yellow' }
273
274function bar($, e) {
275  const { Text } = $.ui.resolve(e)
276  const color = BAR_COLOR[barPhase]
277  return h(Text, color ? { color } : { dimColor: true }, barText)
278}
279
280async function refreshBar($) {
281  const st = await cacheState($)
282  const text = describe(st)
283  if (text === barText && st.phase === barPhase) return
284  // The redraw is a hook, where the note can be shown (showNote).
285  if (barPhase === 'warm' && st.phase === 'cold') pendingCold = coldNote(st)
286  barText = text
287  barPhase = st.phase
288  $.ui.invalidate('ui.render')
289}
290
291// Shown once, when the cache of an idle session goes cold.
292function coldNote(st) {
293  const ctx = typeof st.tokens === 'number' ? ' (' + fmtTokens(st.tokens) + ' tokens)' : ''
294  return 'Cache cold: ' + Math.floor(st.idleMin) + ' min since the last request. The next prompt reads the whole context' +
295    ctx + ' again at the full price; /clear starts a fresh session.'
296}
297
298// A prompt the user sent (typed, or from Remote Control) to an idle session whose cache
299// is cold is held, the model not woken, until they confirm it. A slash command passes:
300// /clear is the usual answer to a cold cache.
301//
302// A typed prompt gets the mod's own dialog in the terminal, with exactly two answers. A hook
303// may not wait on its own code for more than 10 s, so the prompt is dropped at once and the
304// dialog sends it again on Submit. Claude Code shows it as a prompt the plugin sent, and no
305// hook can make it the user's (2.1.287: the mod's own hooks are skipped for a prompt sent from
306// a button's handler; from a ui.press hook, a result without the plugin origin is still shown
307// as the plugin's, and next() refuses another origin). The mod's own prompt.submit hook does
308// not see it, so it is not asked about again. Claude Code's question dialog is used where the
309// mod's cannot be: a prompt from Remote Control (the app draws only that dialog, and adds its
310// own Other answers), or one with attachments, which sending it again would lose.
311const ASK_PANE = 'opsci-cold-ask'
312let asking = null           // { question, text }: the typed prompt the dialog holds
313
314function coldQuestion(st) {
315  const ctx = typeof st.tokens === 'number' ? ' The whole context (' + fmtTokens(st.tokens) + ' tokens) will be read again at the full price.' : ''
316  return 'The cache is cold (' + Math.floor(st.idleMin) + ' min since the last request).' + ctx + ' Are you sure you want to submit this prompt?'
317}
318
319async function answerAsk($, submit) {
320  const held = asking
321  asking = null
322  await $.ui.close({ id: ASK_PANE })
323  if (!held) return
324  if (submit) {
325    await $.prompt.submit({ text: held.text })
326  } else {
327    await $.prompt.fill({ text: held.text })
328    await log($, 'cold-cache prompt held back')
329  }
330}
331
332async function coldGate($, e, next) {
333  if (e.turnId || (e.origin.kind !== 'composer' && e.origin.kind !== 'bridge') || /^\s*\//.test(e.text)) return next(e)
334  const st = await cacheState($)
335  if (st.phase !== 'cold') return next(e)
336  const q = coldQuestion(st)
337  if (e.origin.kind === 'composer' && !(e.attachments && e.attachments.length)) {
338    asking = { question: q, text: e.text }
339    const opened = await $.ui.open({ id: ASK_PANE, title: 'Cache cold', focus: true, closeOnEscape: true, holdToasts: true, rows: 6 })
340    if (opened.isPlaced) return { drop: 'Not sent yet: the cache is cold. Answer below.' }
341    asking = null
342  }
343  let answer = ''
344  try { answer = await $.ui.ask(q, { header: 'Cache cold', options: ['Submit', 'Do not submit'] }) } catch { answer = '' }
345  if (answer === 'Submit') return next(e)
346  if (e.origin.kind === 'composer') await $.prompt.fill({ text: e.text })
347  await log($, 'cold-cache prompt held back (' + Math.floor(st.idleMin) + ' min idle)')
348  return { drop: 'Not submitted: the cache is cold.' +
349    (e.origin.kind === 'composer' ? ' Your prompt is back in the prompt box.' : '') + ' /clear starts a fresh session.' }
350}
351
352export function register(on) {
353  on('session.start', async ($, e, next) => {
354    await $.env.set('OPSCI_MOD', '1')
355    sid = await $.session.id()
356    const poll = Number(await $.env.get('OPSCI_WAIT_POLL')) || POLL_DEFAULT_S
357    pollTimer?.cancel()
358    pollTimer = $.clock.every(poll * 1000, () => { void pollWakers($) })
359    barTimer?.cancel()
360    barText = ''; barPhase = 'none'; barSid = null; lastTokens = null
361    barTimer = $.clock.every(BAR_TICK_MS, () => { void refreshBar($) })
362    try { await $.command.register({ name: NOTE_CMD, description: 'Show the last open-science note again (a jump report, cache cold)' }) } catch (err) { await log($, 'note command not registered: ' + err) }
363    $.clock.after(0, async () => {
364      // A session left waiting by another process (resumed after a SLURM resurrection or
365      // by hand) lost the background tasks that were to wake it: wake it now.
366      const r = await sh($, 'cm_mod.sh', ['resumed', sid])
367      const m = /^\/open-science-context:continue-context(?: (\S*))?$/.exec(r.stdout.trim())
368      if (m) await resume($, m[1] || '')
369      // Only when the name should change: session.start also fires at every reload of a mod,
370      // and a resumed session keeps its name (the session record has it; a session without
371      // one gets it here).
372      await rename($)
373    })
374    $.clock.after(0, () => { void refreshBar($) })
375    return next(e)
376  })
377
378  on('classic.Stop', async ($, e, next) => {
379    if (e.agent_id) return next(e)
380    let tokens
381    try { tokens = (await $.session.usage()).context.tokens } catch { tokens = undefined }
382    const r = await sh($, 'cm_stop.sh', ['--mod'], JSON.stringify({ ...e, opsci_tokens: tokens }))
383    let out = {}
384    try { out = JSON.parse(r.stdout || '{}') } catch { out = {} }
385    // Claude Code refuses a command or prompt started from the Stop hook (it would wait on
386    // the turn the hook holds), so the decision is carried out at turn.complete.
387    atStop = out.opsci || null
388    const res = await next(e)
389    return out.decision === 'block' ? { ...(res || {}), block: out.reason } : res
390  })
391
392  on('turn.start', async ($, e, next) => {
393    if (e.agentId) return next(e)
394    busy = true
395    pendingCold = null
396    coldTimer?.cancel(); coldTimer = null
397    $.clock.after(0, () => { void refreshBar($) })
398    const now = await $.session.id()
399    if (clearedFrom && now !== clearedFrom && !jumping) {
400      // A /clear the user typed: the registration stays with this process, as a pane's does.
401      await sh($, 'cm_mod.sh', ['handover', clearedFrom, now])
402      clearedFrom = null
403    }
404    sid = now
405    if (waiting) { waiting = false; await sh($, 'cm_mod.sh', ['woken', now]) }
406    return next(e)
407  })
408
409  on('turn.complete', async ($, e, next) => {
410    if (e.agentId) return next(e)
411    busy = false
412    $.clock.after(0, () => { void refreshBar($) })
413    const o = atStop || {}
414    atStop = null
415    if (o.jump) jump($, o.jump)
416    else {
417      armCold($, o.cold || 0, o.notice)
418      $.clock.after(0, () => { void rename($) })
419    }
420    return next(e)
421  })
422
423  // The new session a jump's clear started: its report, at the top.
424  on('classic.SessionStart', async ($, e, next) => {
425    if (e.source === 'clear' && pendingJump) { const text = pendingJump; pendingJump = null; showNote($, text) }
426    return next(e)
427  })
428
429  on('command.run', { command: NOTE_CMD }, async () => {
430    const text = noteText || lastNote || 'No open-science note in this session yet.'
431    noteText = null
432    return { text }
433  })
434
435  on('session.end', async ($, e, next) => {
436    if (e.reason === 'clear' && !jumping) clearedFrom = e.sessionId
437    return next(e)
438  })
439
440  on('prompt.submit', coldGate)
441  on('ui.render', { component: 'Pane', requestId: ASK_PANE }, async ($, e) => {
442    const { Box, Text, Button } = $.ui.resolve(e)
443    return h(Box, { flexDirection: 'column' },
444      h(Text, { color: 'yellow' }, '⚠ ' + (asking ? asking.question : 'The cache is cold.')),
445      h(Box, { flexDirection: 'row', gap: 2 },
446        h(Button, { key: 'submit', label: 'Submit', hotkey: '1', variant: 'primary', autoFocus: true, onPress: () => answerAsk($, true) }),
447        h(Button, { key: 'cancel', label: 'Do not submit', hotkey: '2', role: 'dismiss', onPress: () => answerAsk($, false) })))
448  })
449  // Esc, or the pane's close mark: not submitted.
450  on('ui.close', async ($, e, next) => {
451    const r = await next(e)
452    if (e.id === ASK_PANE && asking && e.origin.kind !== 'plugin') {
453      const held = asking
454      asking = null
455      await $.prompt.fill({ text: held.text })
456      await log($, 'cold-cache prompt held back')
457    }
458    return r
459  })
460
461  // The bar: a line above the prompt on the terminal and in the Desktop app. Green while the
462  // cache is warm, yellow with a warning sign once it is cold.
463  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
464    if (pendingCold) { const text = pendingCold; pendingCold = null; if (!busy) showNote($, text) }
465    if (!barText || e.props.hasSurvey) return next(e)
466    return bar($, e)
467  })
468
469  // Each subagent's own context size, from each request it sends. Above the limit, it is
470  // told once to checkpoint, as the dispatch contract asks (open-science-context:
471  // context-management, "Subagents"). Only in a session that drives a task.
472  on('turn.step', async function* ($, e, next) {
473    const result = yield* next(e)
474    if (!e.agentId) { await noteRequest($, result && result.usage); await refreshBar($) }
475    const u = result && result.usage
476    if (e.agentId && u && !warned.has(e.agentId)) {
477      const tokens = (u.input_tokens || 0) + (u.cache_read_input_tokens || 0) + (u.cache_creation_input_tokens || 0)
478      const limit = Number(await $.env.get('OPSCI_SUBAGENT_LIMIT')) || SUB_LIMIT_DEFAULT
479      if (tokens > limit && sid && (await registered($, sid))) {
480        warned.add(e.agentId)
481        const text = 'open-science: your context is ' + tokens + ' tokens, above ' + limit +
482          '. Stop at the next clean boundary, bring your subcontext document up to date, and end your report PAUSED - <doc path> - <exact next step>.'
483        const sent = await $.session.send({ to: { agentId: e.agentId }, text })
484        await log($, 'subagent ' + e.agentId + ' at ' + tokens + ' tokens: checkpoint ' + (sent.isDelivered ? 'sent' : 'NOT delivered: ' + sent.reason))
485      }
486    }
487    return result
488  })
489}
490