tool-timing-badge

1/**
2 * tool-timing-badge — Claude Mod
3 *
4 * Measures how long every tool call takes and draws a colored duration badge
5 * beside the engine's own ToolUse rendering, on the terminal, Desktop and
6 * mobile alike. Two hooks: `tool.call` (timing, "after": it awaits `next`) and
7 * `ui.render` on the ToolUse component (drawing: wraps what `next` drew).
8 *
9 * Elements come from the surface's table (`$.ui.resolve(e)`), never from
10 * globals; JSX compiles against the environment's `h`. A ToolUse is drawn
11 * once per props change, so the badge appears on the redraw that follows the
12 * call settling (`isRunning` flips to false).
13 *
14 * Needs Claude Code >= 2.1.287. Typed
15 * against Anthropic's declarations: https://github.com/anthropics/claude-code/tree/main/mods
16 *
17 * Options:
18 *   slowMs: number  calls at or above this are red, above a quarter of it yellow (default 5000)
19 */
20import type { Register } from 'claude-code'
21
22const durations = new Map<string, number>()
23
24function badgeColor(ms: number, slowMs: number): string {
25  if (ms >= slowMs) return 'red'
26  if (ms >= slowMs / 4) return 'yellow'
27  return 'green'
28}
29
30function format(ms: number): string {
31  return ms < 1000 ? `${ms} ms` : `${(ms / 1000).toFixed(1)} s`
32}
33
34export const register: Register = (on, options) => {
35  const slowMs = typeof options.slowMs === 'number' ? options.slowMs : 5000
36
37  // 1. Time every tool call, the hooks beneath and the tool itself included.
38  on('tool.call', async ($, e, next) => {
39    const startedAt = Date.now()
40    try {
41      return await next(e)
42    } finally {
43      const ms = Date.now() - startedAt
44      if (e.tool_use_id) durations.set(e.tool_use_id, ms)
45      if (ms >= slowMs) $.ui.log(`[tool-timing-badge] slow ${e.tool}: ${format(ms)}`)
46    }
47  })
48
49  // 2. Wrap the engine's ToolUse drawing with a badge.
50  on('ui.render', { component: 'ToolUse' }, async ($, e, next) => {
51    const { Box, Text } = $.ui.resolve(e)
52    const drawn = await next(e)
53
54    const ms = durations.get(e.props.tool_use_id)
55    if (e.props.isRunning || ms === undefined) return drawn
56
57    return (
58      <Box flexDirection="row" gap={1}>
59        {drawn}
60        <Text color={badgeColor(ms, slowMs)}>{format(ms)}</Text>
61      </Box>
62    )
63  })
64}
65