codebuddy

1import type { EngineInterface, Register, Timer } from 'claude-code'
2
3const USAGE_URL = 'https://www.codebuddy.cn/billing/meter/get-enterprise-user-usage'
4const DAILY_USAGE_URL = 'https://www.codebuddy.cn/billing/meter/get-user-daily-usage'
5const AUTH_PATH = '/AppData/Local/CodeBuddyExtension/Data/Public/auth/Tencent-Cloud.coding-copilot.info'
6
7const object = (value: unknown): Record<string, unknown> =>
8  value !== null && typeof value === 'object' && !Array.isArray(value)
9    ? value as Record<string, unknown>
10    : {}
11
12const points = (value: unknown): number | undefined => {
13  if (typeof value !== 'number' && typeof value !== 'string') return undefined
14  if (typeof value === 'string' && value.trim() === '') return undefined
15  const number = Number(value)
16  return Number.isFinite(number) ? number : undefined
17}
18
19const today = () => {
20  const parts = new Intl.DateTimeFormat('en-US', {
21    timeZone: 'Asia/Taipei', year: 'numeric', month: '2-digit', day: '2-digit',
22  }).formatToParts(new Date())
23  const part = (type: string) => parts.find((value) => value.type === type)?.value
24  return `${part('year')}-${part('month')}-${part('day')}`
25}
26
27const loadDailyUsage = async ($: EngineInterface, headers: Record<string, string>): Promise<string> => {
28  try {
29    const date = today()
30    const response = await $.http.fetch(DAILY_USAGE_URL, {
31      method: 'POST', headers,
32      body: JSON.stringify({
33        startTime: `${date} 00:00:00`, endTime: `${date} 23:59:59`, pageNum: 1, pageSize: 1,
34      }),
35    })
36    if (!response.ok) return '查询失败'
37    const result = object(JSON.parse(response.text))
38    if (result.code !== 0) return '查询失败'
39    const data = object(result.data)
40    if (!Array.isArray(data.data)) return '数据异常'
41    if (data.data.length === 0 && points(data.total) === 0) return '0'
42    const credit = points(object(data.data[0]).credit)
43    return credit === undefined ? '数据异常' : String(Math.trunc(credit))
44  } catch {
45    return '查询失败'
46  }
47}
48
49const loadUsage = async ($: EngineInterface): Promise<string> => {
50  let account: Record<string, unknown>
51  let auth: Record<string, unknown>
52  try {
53    const home = await $.env.get('USERPROFILE')
54    if (!home) return '无法确定用户目录'
55    const info = object(JSON.parse((await $.fs.read(
56      home.replace(/\\/g, '/').replace(/\/$/, '') + AUTH_PATH,
57    )).replace(/^\uFEFF/, '')))
58    account = object(info.account)
59    auth = object(info.auth)
60  } catch {
61    return '无法读取登录信息，请先登录 CodeBuddy'
62  }
63
64  const enterpriseId = account.enterpriseId
65  const accessToken = auth.accessToken
66  if (typeof enterpriseId !== 'string' || enterpriseId.trim() === ''
67    || typeof accessToken !== 'string' || accessToken.trim() === '') {
68    return '登录信息缺少 enterpriseId 或 accessToken'
69  }
70
71  try {
72    const headers = {
73      Authorization: `Bearer ${accessToken}`,
74      'x-enterprise-id': enterpriseId,
75      'Content-Type': 'application/json',
76    }
77    const response = await $.http.fetch(USAGE_URL, {
78      method: 'POST',
79      headers,
80      body: '{}',
81    })
82    if (response.status === 401 || response.status === 403) return '登录失效或无权限，请重新登录 CodeBuddy'
83    if (!response.ok) return `积分查询失败（HTTP ${response.status}）`
84    const result = object(JSON.parse(response.text))
85    if (result.code !== 0) return '积分接口返回业务错误，请稍后重试'
86    const data = object(result.data)
87    const credit = points(data.credit)
88    const limitNum = points(data.limitNum)
89    if (credit === undefined || limitNum === undefined) return '积分接口返回的数据格式异常'
90    const dailyUsage = await loadDailyUsage($, headers)
91    return `today ${dailyUsage} · remain ${Math.trunc(credit)}/${Math.trunc(limitNum)}`
92  } catch {
93    // 不输出原始异常、响应或认证信息，避免泄露令牌。
94    return '积分查询失败，请检查网络或登录状态'
95  }
96}
97
98let label = '正在查询积分…'
99let timer: Timer | undefined
100let pending: Promise<void> | undefined
101let generation = 0
102
103const publishUsage = async ($: EngineInterface) => {
104  try {
105    const home = await $.env.get('USERPROFILE')
106    if (!home) return
107    await $.fs.write(`${home.replace(/\\/g, '/')}/.claude/codebuddy/usage.json`, JSON.stringify({
108      text: `token usage · ${label}`, updatedAt: Date.now(),
109    }))
110  } catch {
111    $.ui.log('无法更新 CodeBuddy 状态栏缓存', { to: 'debug' })
112  }
113}
114
115const refresh = ($: EngineInterface): Promise<void> => {
116  if (pending) return pending
117  const current = generation
118  pending = (async () => {
119    try {
120      const nextLabel = await loadUsage($)
121      if (current !== generation) return
122      label = nextLabel
123      await publishUsage($)
124    } finally {
125      if (current === generation) pending = undefined
126    }
127  })()
128  return pending
129}
130
131export const register: Register = (on) => {
132  on('session.start', async ($, e, next) => {
133    timer?.cancel()
134    generation += 1
135    pending = undefined
136    label = '正在查询积分…'
137    $.ui.status(undefined)
138    await $.command.register({ name: 'codebuddy', description: '刷新 CodeBuddy 剩余积分、总积分和今日使用', immediate: true })
139    timer = $.clock.every(60_000, () => { void refresh($) })
140    void refresh($)
141    return next(e)
142  })
143
144  on('command.run', { command: 'codebuddy' }, async ($) => {
145    await refresh($)
146    return { text: `token usage · ${label}` }
147  })
148
149  on('ui.render', { component: 'InfoNotice' }, ($, e, next) => {
150    if (/^(?:●\s*)?codebuddy:\s*reloaded\s*\(\d+ hooks?:[^\r\n]*\)\s*$/.test(e.props.text.trim())) {
151      const { Box } = $.ui.resolve(e)
152      return <Box />
153    }
154    return next(e)
155  })
156
157  on('session.end', ($, e, next) => {
158    timer?.cancel()
159    timer = undefined
160    generation += 1
161    pending = undefined
162    $.ui.status(undefined)
163    return next(e)
164  })
165}
166