jsc-marketplace/plugins/jsc-mod/hooks/register.tsx

import { atom, read, update } from 'claude-code'
import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'

import type { Mode, Sys, Usage, Window } from '../types'
import { fit, cpuPercent, parseCpu, parseDf, parseGit, parseMem, parseOffset } from './format'
import type { CpuSample } from './format'

const usageAtom = atom({ plugin: 'jsc-mod', key: 'usage' } as const, {})
const gitAtom = atom({ plugin: 'jsc-mod', key: 'git' } as const, null)
const sysAtom = atom({ plugin: 'jsc-mod', key: 'sys' } as const, {})
const modeAtom = atom({ plugin: 'jsc-mod', key: 'mode' } as const, 'auto')
const alertedAtom = atom({ plugin: 'jsc-mod', key: 'alerted' } as const, [])
const effortAtom = atom({ plugin: 'jsc-mod', key: 'effort' } as const, null)

const SYS_MS = 3000
const GIT_MS = 15_000
const DISK_EVERY = 10 // 每 10 次系統取樣（30 秒）讀一次磁碟
const ALERT_AT = 90
const GIT_TOOLS = new Set(['Bash', 'Edit', 'Write', 'NotebookEdit'])
const LABEL: Record<string, string> = { five_hour: '5 小時', seven_day: '7 天' }

// 引擎要等第一次 API 回應才有額度與 ctx；先前的額度存在 $.store，開新對話時先顯示，ctx 先當作 0
function toUsage(context: SessionContextUsage, rateLimits: SessionRateLimit[]): Usage {
  const win = (kind: string) => {
    const r = rateLimits.find(x => x.kind === kind)
    return r && { percent: r.percentUsed, resetsAt: r.resetsAt }
  }
  return {
    fiveHour: win('five_hour'),
    sevenDay: win('seven_day'),
    ctxTokens: context.tokens ?? 0,
    ctxWindow: context.window,
    ctxPercent: context.percent ?? 0,
  }
}

// 已過重置時間的視窗用量已歸零，不再沿用舊值
const live = (w: Window | undefined, now: number) =>
  w && (!w.resetsAt || Date.parse(w.resetsAt) > now) ? w : undefined

async function restoreRates($: EngineInterface, usage: Usage): Promise<Usage> {
  const saved = (await $.store.get('rates')) as Pick<Usage, 'fiveHour' | 'sevenDay'> | undefined
  if (!saved) return usage
  const now = await $.clock.now()

  return {
    ...usage,
    fiveHour: usage.fiveHour ?? live(saved.fiveHour, now),
    sevenDay: usage.sevenDay ?? live(saved.sevenDay, now),
  }
}

let home: string | undefined
let tzOffsetMin = 0
let prevCpu: CpuSample | null = null
let ticks = 0

async function refreshGit($: EngineInterface) {
  try {
    const r = await $.process.run(['git', 'status', '--porcelain=v2', '-b'], { timeoutMs: 5000 })
    const git = r.exitCode === 0 ? parseGit(r.stdout) : null
    await update($, gitAtom, () => git)
  } catch {
    await update($, gitAtom, () => null)
  }
}

async function refreshSys($: EngineInterface) {
  const next: Sys = { ...(await read($, sysAtom)) }
  try {
    const cur = parseCpu(await $.fs.read('/proc/stat'))
    if (cur) {
      if (prevCpu) next.cpu = cpuPercent(prevCpu, cur) ?? next.cpu
      prevCpu = cur
    }
  } catch {}
  try {
    const mem = parseMem(await $.fs.read('/proc/meminfo'))
    if (mem) [next.memUsed, next.memTotal] = [mem.used, mem.total]
  } catch {}
  if (ticks++ % DISK_EVERY === 0) {
    try {
      const r = await $.process.run(['df', '-P', await $.session.cwd()], { timeoutMs: 5000 })
      if (r.exitCode === 0) next.disk = parseDf(r.stdout)
    } catch {}
  }
  await update($, sysAtom, () => next)
}

// 每個額度視窗（以重置時間區分）只提醒一次
async function alert($: EngineInterface, rateLimits: SessionRateLimit[]) {
  const done = await read($, alertedAtom)
  for (const r of rateLimits) {
    const id = `${r.kind}@${r.resetsAt ?? ''}`
    if (r.percentUsed < ALERT_AT || done.includes(id)) continue
    await update($, alertedAtom, list => [...list, id].slice(-20))
    const label = LABEL[r.kind] ?? r.kind
    $.ui.toast(`${label}額度已用 ${Math.round(r.percentUsed)}%`)
  }
}

async function runHud($: EngineInterface, args: string) {
  const arg = args.trim()
  const mode: Mode = arg === 'off' ? 'off' : arg === 'compact' ? 'compact' : 'auto'
  await update($, modeAtom, () => mode)
  const said: Record<Mode, string> = { auto: '狀態列已開啟。', off: '狀態列已隱藏。', compact: '狀態列改用極窄版。' }

  return { text: said[mode] }
}

export const register: Register = on => {
  on('session.start', async ($, e, next) => {
    await $.command.register({
      name: 'statusline',
      description: '切換輸入框上方的狀態列',
      argumentHint: 'on | off | compact',
    })
    home = await $.env.get('HOME')
    try {
      const r = await $.process.run(['date', '+%z'], { timeoutMs: 3000 })
      tzOffsetMin = parseOffset(r.stdout)
    } catch {}
    const u = await $.session.usage()
    const usage = await restoreRates($, toUsage(u.context, u.rateLimits))
    await update($, usageAtom, () => usage)
    void refreshGit($)
    // CPU 要兩次取樣才算得出來，開場先隔 0.5 秒補一次，免得等到 3 秒後才出現
    void refreshSys($).then(() => $.clock.after(500, () => void refreshSys($)))
    $.clock.every(SYS_MS, () => void refreshSys($))
    $.clock.every(GIT_MS, () => void refreshGit($))

    return next(e)
  })

  on('session.measure', async ($, e, next) => {
    await update($, usageAtom, () => toUsage(e.context, e.rateLimits))
    if (e.rateLimits.length) {
      const w = (kind: string) => e.rateLimits.find(x => x.kind === kind)
      const pick = (r: SessionRateLimit | undefined) => r && { percent: r.percentUsed, resetsAt: r.resetsAt }
      await $.store.set('rates', { fiveHour: pick(w('five_hour')), sevenDay: pick(w('seven_day')) })
    }
    if (e.changed.includes('rateLimits')) await alert($, e.rateLimits)

    return next(e)
  })

  // 思考等級只在送出模型請求時看得到；只記主對話的，不記 subagent
  on('turn.step', async function* ($, e, next) {
    if (!e.agentId) {
      const effort = e.effort === undefined ? null : String(e.effort)
      if (effort !== (await read($, effortAtom))) await update($, effortAtom, () => effort)
    }

    return yield* next(e)
  })

  on('tool.call', async ($, e, next) => {
    const ran = await next(e)
    if (GIT_TOOLS.has(e.tool)) void refreshGit($)

    return ran
  })

  on('command.run', { command: 'statusline' }, ($, e) => runHud($, e.args))

  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
    const mode = await read($, modeAtom)
    const rest = await next(e)
    if (e.props.hasSurvey || mode === 'off') return rest

    const data = {
      model: await $.session.model(),
      cwd: await $.session.cwd(),
      home,
      git: await read($, gitAtom),
      usage: await read($, usageAtom),
      sys: await read($, sysAtom),
      tzOffsetMin,
      effort: await read($, effortAtom),
    }
    const line = fit(data, e.props.bodyColumns, mode === 'compact')
    const { Box, Text } = $.ui.resolve(e)

    // 引擎在 band 與輸入框之間已留一行空白，上方補一行讓上下間距相同
    return (
      <Box marginTop={1} flexDirection="column">
        <Text wrap="truncate">
          {line.map(s => (
            <Text color={s.color} dimColor={s.dim} bold={s.bold}>
              {s.text}
            </Text>
          ))}
        </Text>
        {rest}
      </Box>
    )
  })
}