clauguard

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Phase, Run } from '../types'
5import {
6  CONTINUE,
7  GRACE_CALLS,
8  HANDOFF,
9  STOPPED,
10  WRAPUP,
11  isRemember,
12  label,
13  nextPhase,
14  parseStart,
15  settle,
16} from './guard'
17
18const IDLE: Phase = { kind: 'idle' }
19
20const run = atom({ plugin: 'clauguard', key: 'run' } as const, null as Run | null)
21const phase = atom({ plugin: 'clauguard', key: 'phase' } as const, IDLE as Phase)
22const turnId = atom({ plugin: 'clauguard', key: 'turnId' } as const, null as string | null)
23
24const RETRY_MS = 250
25const RETRIES = 40
26
27const sleep = ($: EngineInterface, ms: number) =>
28  new Promise<void>((resolve) => $.clock.after(ms, resolve))
29
30// Compaction is refused while a turn still winds down after the abort, so it is retried briefly.
31const compactWhenIdle = async ($: EngineInterface) => {
32  for (let attempt = 0; attempt < RETRIES; attempt++) {
33    try {
34      await $.session.compact({ instructions: 'Keep the current task, its plan and the handoff.' })
35      return true
36    } catch {
37      await sleep($, RETRY_MS)
38    }
39  }
40  return false
41}
42
43const endTurn = async ($: EngineInterface) => {
44  const id = await read($, turnId)
45  if (id !== null) await $.turn.abort({ turnId: id }).catch(() => undefined)
46}
47
48const handoffDone = async ($: EngineInterface) => {
49  await endTurn($)
50  if (await compactWhenIdle($)) await $.prompt.submit({ text: CONTINUE })
51  else $.ui.toast('clauguard: compact failed, please run /compact')
52  await update($, phase, () => IDLE)
53}
54
55const stop = async ($: EngineInterface) => {
56  await update($, phase, () => ({ kind: 'stopped' }) as Phase)
57  $.ui.toast('clauguard: Limit erreicht', { timeoutMs: 10_000 })
58  $.clock.after(0, () => void endTurn($))
59}
60
61export const register: Register = (on) => {
62  on('session.start', async ($, e, next) => {
63    await $.command.register({
64      name: 'guard',
65      description: 'Usage budget for this run: start [points], stop, status',
66      argumentHint: 'start [points] | stop | status',
67    })
68    return next(e)
69  })
70
71  on('command.run', { command: 'guard' }, async ($, e) => {
72    const args = e.args.trim()
73    const budget = parseStart(args)
74    if (budget !== undefined) {
75      const { rateLimits } = await $.session.usage()
76      const weekly = rateLimits.find((r) => r.kind === 'seven_day')?.percentUsed
77      const started = settle({ budget }, weekly)
78      await update($, run, () => started)
79      await update($, phase, () => IDLE)
80      $.ui.toast(
81        started.cap === undefined
82          ? `clauguard: +${budget} points, the stop is set with the first response`
83          : `clauguard: ${Math.round(started.baseline ?? 0)}% now, stop at ${Math.round(started.cap)}%`,
84      )
85    } else if (args === 'stop') {
86      await update($, run, () => null)
87      await update($, phase, () => IDLE)
88      $.ui.toast('clauguard: budget off')
89    } else {
90      const current = await read($, run)
91      $.ui.toast(
92        label(current, await read($, phase)) ?? 'clauguard: no budget set (/guard start 30)',
93      )
94    }
95    $.ui.invalidate('ui.render')
96    return {}
97  })
98
99  on('turn.start', async ($, e, next) => {
100    if (e.agentId === undefined) await update($, turnId, () => e.turnId)
101    return next(e)
102  })
103
104  on('session.measure', async ($, e, next) => {
105    const weekly = e.rateLimits.find((r) => r.kind === 'seven_day')?.percentUsed
106    const { percent, tokens, window } = e.context
107    const context =
108      percent ?? (tokens !== undefined && window > 0 ? (tokens / window) * 100 : undefined)
109
110    const current = await read($, run)
111    const settled = current && settle(current, weekly)
112    if (settled !== current) await update($, run, () => settled)
113
114    const before = await read($, phase)
115    const after = nextPhase(before, settled, weekly, context)
116    if (after !== before) {
117      await update($, phase, () => after)
118      if (after.kind === 'wrapup') $.ui.toast('clauguard: Budget erreicht, Claude schließt ab')
119      $.ui.invalidate('ui.render')
120    }
121    return next(e)
122  })
123
124  on('tool.call', async ($, e, next) => {
125    const current = await read($, phase)
126
127    if (current.kind === 'stopped') return { deny: STOPPED }
128
129    if (current.kind === 'wrapup') {
130      if (current.callsLeft === 0) {
131        await stop($)
132        return { deny: STOPPED }
133      }
134      await update($, phase, () => ({ kind: 'wrapup', callsLeft: current.callsLeft - 1 }) as Phase)
135      if (current.callsLeft === GRACE_CALLS) return { deny: WRAPUP }
136      return next(e)
137    }
138
139    // Subagents can't compact the main conversation, so the handoff is the main agent's job.
140    if (current.kind !== 'handoff' || e.agentId !== undefined) return next(e)
141
142    if (!current.asked) {
143      await update($, phase, () => ({ kind: 'handoff', asked: true }) as Phase)
144      return { deny: HANDOFF }
145    }
146
147    const result = await next(e)
148    if (e.tool === 'Skill' && isRemember(e.skill)) $.clock.after(0, () => void handoffDone($))
149    return result
150  })
151
152  // Any compaction, also /compact or the automatic one, ends a pending handoff.
153  on('session.compact', async ($, e, next) => {
154    const result = await next(e)
155    if (e.agentId === undefined && (await read($, phase)).kind === 'handoff') {
156      await update($, phase, () => IDLE)
157    }
158    return result
159  })
160
161  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
162    const text = label(await read($, run), await read($, phase))
163    if (text === undefined) return next(e)
164    const tail = e.props.tail ? `${e.props.tail}  ${text}` : text
165    return next({ ...e, props: { ...e.props, tail } })
166  })
167}
168

1import type { Phase, Run } from '../types'
2
3export const DEFAULT_BUDGET = 30
4export const HANDOFF_AT = 85
5export const GRACE_CALLS = 10
6
7export const HANDOFF = `clauguard: the context is at ${HANDOFF_AT}%. Finish the step you are on, update your memory, then run the remember skill (/remember). Once it is saved, clauguard compacts the conversation and you carry on.`
8export const WRAPUP = `clauguard: the usage budget for this run is reached. Bring the current step to a clean end within ${GRACE_CALLS} tool calls, write a short handoff, start nothing new.`
9export const STOPPED =
10  'clauguard: Limit erreicht. The run is stopped, no more tool calls until Flo starts a new one.'
11export const CONTINUE = 'Context was compacted after your handoff. Carry on where you left off.'
12
13export const isRemember = (skill: string) => /(^|:)remember$/.test(skill)
14
15// The budget counts from the first weekly figure the run sees, so a start before any response works.
16export const settle = (run: Run, weekly: number | undefined): Run =>
17  run.baseline !== undefined || weekly === undefined
18    ? run
19    : { ...run, baseline: weekly, cap: weekly + run.budget }
20
21export const parseStart = (args: string): number | undefined => {
22  const match = /^start(?:\s+(\d+(?:\.\d+)?))?$/.exec(args.trim())
23  if (!match) return undefined
24  return match[1] === undefined ? DEFAULT_BUDGET : Number(match[1])
25}
26
27// Budget beats handoff: a run over budget wraps up, it does not start a compaction.
28export const nextPhase = (
29  phase: Phase,
30  run: Run | null,
31  weekly: number | undefined,
32  context: number | undefined,
33): Phase => {
34  if (phase.kind === 'wrapup' || phase.kind === 'stopped') return phase
35  if (run?.cap !== undefined && weekly !== undefined && weekly >= run.cap) {
36    return { kind: 'wrapup', callsLeft: GRACE_CALLS }
37  }
38  if (phase.kind === 'idle' && context !== undefined && context >= HANDOFF_AT) {
39    return { kind: 'handoff', asked: false }
40  }
41  return phase
42}
43
44export const label = (run: Run | null, phase: Phase) => {
45  if (phase.kind === 'stopped') return 'guard ⛔ Limit erreicht'
46  if (phase.kind === 'wrapup') return 'guard ⏳ schließt ab'
47  if (!run) return undefined
48  return run.cap === undefined ? 'guard …' : `guard ≤${Math.round(run.cap)}%`
49}
50

1// Weekly usage in percent points: the run started at `baseline` and stops at `cap`.
2export type Run = { budget: number; baseline?: number; cap?: number }
3
4export type Phase =
5  | { kind: 'idle' }
6  | { kind: 'handoff'; asked: boolean }
7  | { kind: 'wrapup'; callsLeft: number }
8  | { kind: 'stopped' }
9
10declare module 'claude-code' {
11  interface PluginState {
12    clauguard: { run: Run | null; phase: Phase; turnId: string | null }
13  }
14}
15