jsc-marketplace/plugins/cockpit/hooks/hud/index.tsx

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

import type { Mode, Sys, Usage } from '../../types'
import { fitCompanion } from '../companion/band'
import { createCompanion } from '../companion/core'
import { nodeFs } from '../companion/node-fs'
import { bandLines, cpuPercent, parseCpu, parseDf, parseGit, parseMem, parseOffset } from './format'
import type { CpuSample } from './format'

const usageAtom = atom({ plugin: 'cockpit', key: 'usage' } as const, {})
const gitAtom = atom({ plugin: 'cockpit', key: 'git' } as const, null)
const sysAtom = atom({ plugin: 'cockpit', key: 'sys' } as const, {})
const modeAtom = atom({ plugin: 'cockpit', key: 'mode' } as const, 'auto')
const alertedAtom = atom({ plugin: 'cockpit', key: 'alerted' } as const, [])
const effortAtom = atom({ plugin: 'cockpit', 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 天' }

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,
    ctxWindow: context.window,
    ctxPercent: context.percent,
  }
}

let home: string | undefined
let companion: ReturnType<typeof createCompanion> | 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: 'cockpit',
      description: 'Cockpit 控制：助理（status、switch、companion off、setup）與 hud 狀態列',
      argumentHint: 'status | switch <角色> | companion off | setup | hud on | off | compact',
    })
    // 別名：/hud <arg> 等同 /cockpit hud <arg>
    await $.command.register({
      name: 'hud',
      description: '切換輸入框上方的狀態列（/cockpit hud 的別名）',
      argumentHint: 'on | off | compact',
    })
    home = await $.env.get('HOME')
    companion = createCompanion({ fs: nodeFs, clock: { now: Date.now }, root: `${home ?? ''}/.cockpit` })
    const first = await companion.ensureActivated()
    if (first.first) $.ui.toast(`已啟用 ${first.name}。用 /cockpit companion off 可關閉`)
    try {
      const r = await $.process.run(['date', '+%z'], { timeoutMs: 3000 })
      tzOffsetMin = parseOffset(r.stdout)
    } catch {}
    const u = await $.session.usage()
    await update($, usageAtom, () => toUsage(u.context, u.rateLimits))
    void refreshGit($)
    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.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: 'hud' }, ($, e) => runHud($, e.args))

  on('command.run', { command: 'cockpit' }, async ($, e, next) => {
    const [sub, ...rest] = e.args.trim().split(/\s+/)
    if (sub === 'hud') return runHud($, rest.join(' '))

    return { text: companion ? await companion.command(e.args) : '尚未就緒。' }
  })

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

    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 view = await companion?.view()
    const cols = e.props.bodyColumns
    const lines = bandLines(data, cols, mode === 'compact', view ? fitCompanion(view, cols, mode === 'compact') : null)
    const { Box, Text } = $.ui.resolve(e)

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