lens-view

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { LensViewCell, LensViewMeta, LensViewProps, LensViewSegment, LensViewState } from '../types'
5import { clamp, layout, quoted, shown } from './layout'
6
7const PANE = 'lens-view'
8const DEFAULT_K = 5
9const DEFAULTS: LensViewState = {
10  path: null,
11  cache: null,
12  meta: null,
13  error: null,
14  pos: 0,
15  k: DEFAULT_K,
16  lensIndex: 0,
17  topOffset: 0,
18  panelOffset: Number.MAX_SAFE_INTEGER, // start at the final layers
19  cell: null,
20}
21const view = atom({ plugin: 'lens-view', key: 'view' } as const, DEFAULTS)
22
23// The session keeps $.state across reloads, so a value an older version wrote may lack fields
24// this one reads: fill them from the defaults.
25function complete(v: Partial<LensViewState>): LensViewState {
26  const merged = { ...DEFAULTS, ...v }
27  for (const key of ['pos', 'k', 'lensIndex', 'topOffset', 'panelOffset'] as const)
28    if (!Number.isFinite(merged[key])) merged[key] = DEFAULTS[key]
29  return merged
30}
31
32const current = async ($: EngineInterface) => complete(await read($, view))
33const change = ($: EngineInterface, fn: (v: LensViewState) => LensViewState) => update($, view, v => fn(complete(v)))
34
35// The pane's geometry as last drawn: what the wheel and the keys move within. Recomputed on every
36// draw, so losing it on a reload costs nothing.
37let geometry = { topRows: 10, panelRows: 8, lineCount: 0, layerCount: 0, lineOf: [] as number[] }
38// What the pane last reported about itself, for the configure tool's answer.
39let seen = 'not drawn yet'
40let wheel = { events: 0, last: 'none' }
41// Rows to draw when the surface reports few (an inline pane grows to its content).
42const MIN_ROWS = 24
43
44// Python owns the trace formats: lens-view-json exports the trace to a cache directory of
45// meta.json and one pos-<i>.json per position, and the pane reads only what it shows.
46async function load($: EngineInterface, path: string, options: { lens?: string; k?: number; position?: number } = {}) {
47  const argv = ['uv', 'run', '--quiet', '--frozen', '--no-dev', '--project', $.plugin.root, 'lens-view-json', path]
48  let cache: string | null = null
49  let meta: LensViewMeta | null = null
50  let error: string | null = null
51  try {
52    const { stdout, stderr } = await $.process.run(argv, { timeoutMs: 300_000 })
53    const out = JSON.parse(stdout.trim() || '{}') as { cache?: string; error?: string }
54    if (out.cache) {
55      cache = out.cache
56      meta = JSON.parse(await $.fs.read(`${cache}/meta.json`)) as LensViewMeta
57    } else error = out.error ?? (stderr.trim() || 'lens-view-json printed nothing.')
58  } catch (exc) {
59    error = `Could not load the trace (is uv installed?): ${String(exc)}`
60  }
61  const lensIndex = meta && options.lens ? Math.max(0, meta.lenses.findIndex(l => l.name === options.lens)) : 0
62  await change($, v => ({
63    ...v,
64    path,
65    cache,
66    meta,
67    error,
68    lensIndex,
69    pos: 0,
70    topOffset: 0,
71    panelOffset: Number.MAX_SAFE_INTEGER,
72    cell: null,
73    k: options.k ? clamp(options.k, 1, meta?.k ?? 10) : v.k,
74  }))
75  if (meta) await select($, options.position ?? 0)
76  return meta
77    ? `Loaded ${path}: ${meta.tokens.length} tokens, lenses ${meta.lenses.map(l => l.name).join(', ')}, top ${meta.k} stored per cell.`
78    : (error ?? 'Nothing loaded.')
79}
80
81async function select($: EngineInterface, pos: number) {
82  const v = await current($)
83  if (!v.meta || !v.cache) return
84  const target = clamp(pos, 0, v.meta.tokens.length - 1)
85  let cell: LensViewCell | null = null
86  try {
87    cell = JSON.parse(await $.fs.read(`${v.cache}/pos-${target}.json`)) as LensViewCell
88  } catch (exc) {
89    await change($, s => ({ ...s, error: `Could not read position ${target}: ${String(exc)}` }))
90    return
91  }
92  // Keep the selected token in view.
93  const line = geometry.lineOf[target] ?? 0
94  let topOffset = v.topOffset
95  if (line < topOffset) topOffset = line
96  if (line >= topOffset + geometry.topRows) topOffset = line - geometry.topRows + 1
97  await change($, s => ({ ...s, pos: target, cell, topOffset, error: null }))
98}
99
100async function configure($: EngineInterface, options: { k?: number; lens?: string; position?: number }) {
101  const v = await current($)
102  const notes: string[] = []
103  if (options.k !== undefined) {
104    const k = clamp(Math.round(options.k), 1, v.meta?.k ?? 10)
105    await change($, s => ({ ...s, k }))
106    notes.push(`k = ${k}`)
107  }
108  if (options.lens !== undefined && v.meta) {
109    const i = v.meta.lenses.findIndex(l => l.name === options.lens)
110    if (i < 0) notes.push(`no lens "${options.lens}" (have ${v.meta.lenses.map(l => l.name).join(', ')})`)
111    else {
112      await change($, s => ({ ...s, lensIndex: i }))
113      notes.push(`lens = ${options.lens}`)
114    }
115  }
116  if (options.position !== undefined) {
117    await select($, options.position)
118    notes.push(`position = ${(await current($)).pos}`)
119  }
120  return `lens-view: ${notes.length ? notes.join(', ') : 'nothing changed'}. Pane: ${seen}; wheel events: ${wheel.events} (last ${wheel.last}); trace: ${v.path ?? 'none'}${v.error ? `; error: ${v.error}` : ''}.`
121}
122
123async function onKey($: EngineInterface, key: string) {
124  const v = await current($)
125  if (!v.meta) return
126  const last = v.meta.tokens.length - 1
127  const lineStart = (line: number) => geometry.lineOf.findIndex(l => l === line)
128  switch (key) {
129    case 'left':
130      return select($, v.pos - 1)
131    case 'right':
132    case ' ':
133      return select($, v.pos + 1)
134    case 'up':
135    case 'down': {
136      const line = (geometry.lineOf[v.pos] ?? 0) + (key === 'up' ? -1 : 1)
137      const start = lineStart(clamp(line, 0, geometry.lineCount - 1))
138      return start >= 0 ? select($, start) : undefined
139    }
140    case 'home':
141      return select($, 0)
142    case 'end':
143      return select($, last)
144    case 'pageup':
145    case 'pagedown':
146      return scrollPanel($, (key === 'pageup' ? -1 : 1) * geometry.panelRows)
147    case '+':
148    case '=':
149      return configure($, { k: v.k + 1 })
150    case '-':
151      return configure($, { k: v.k - 1 })
152    case 'l':
153    case 'tab':
154      return change($, s => ({ ...s, lensIndex: s.meta ? (s.lensIndex + 1) % s.meta.lenses.length : 0 }))
155    case 'r':
156      return v.path ? void (await load($, v.path, { lens: v.meta.lenses[v.lensIndex]?.name, k: v.k, position: v.pos })) : undefined
157  }
158}
159
160async function scrollPanel($: EngineInterface, by: number) {
161  const max = Math.max(0, geometry.layerCount - geometry.panelRows)
162  await change($, s => ({ ...s, panelOffset: clamp(Math.min(s.panelOffset, max) + by, 0, max) }))
163}
164
165async function scrollTranscript($: EngineInterface, by: number) {
166  const max = Math.max(0, geometry.lineCount - geometry.topRows)
167  await change($, s => ({ ...s, topOffset: clamp(Math.min(s.topOffset, max) + by, 0, max) }))
168}
169
170// bodyRows: the Client's rows (the title bar above it is the hook's own).
171function draw(v: LensViewState, width: number, bodyRows: number): LensViewProps {
172  const panelRows = Math.max(4, Math.floor((bodyRows - 2) * 0.38))
173  const topRows = Math.max(3, bodyRows - 2 - panelRows)
174  const hint = '←/→ token  ↑/↓ line  click select  wheel scroll  +/- k  l lens  r reload'
175  if (!v.meta) {
176    geometry = { topRows, panelRows, lineCount: 0, layerCount: 0, lineOf: [] }
177    const message = v.error ?? (v.path ? `Loading ${v.path}…` : 'No trace loaded. Run /lens-view <path>, or ask Claude to open one.')
178    return { width, title: 'lens-view', status: '', lines: [[{ text: message, pos: -1, style: 'odd' }]], topRows, selection: '', rows: [], panelRows, chipWidth: 12, hint }
179  }
180  const meta = v.meta
181  const lens = meta.lenses[v.lensIndex] ?? meta.lenses[0]!
182  const { lines, lineOf } = layout(meta.tokens, Math.max(10, width))
183  geometry = { topRows, panelRows, lineCount: lines.length, layerCount: lens.layers.length, lineOf }
184  const topOffset = clamp(v.topOffset, 0, Math.max(0, lines.length - topRows))
185  const styled: LensViewSegment[][] = lines.slice(topOffset, topOffset + topRows).map(line =>
186    line.map(piece => ({
187      text: piece.text,
188      pos: piece.pos,
189      style:
190        piece.pos === v.pos
191          ? 'selected'
192          : (meta.next_prob[piece.pos] ?? 1) < 0.05
193            ? 'surprise'
194            : piece.pos % 2
195              ? 'odd'
196              : 'even',
197    })),
198  )
199
200  const panelOffset = clamp(v.panelOffset, 0, Math.max(0, lens.layers.length - panelRows))
201  const next = meta.targets?.[v.pos] ?? null
202  const layers = v.cell?.[lens.name] ?? []
203  const rows = lens.layers.slice(panelOffset, panelOffset + panelRows).map((label, i) => ({
204    label,
205    chips: (layers[panelOffset + i] ?? []).slice(0, v.k).map(([text, prob]) => ({ text: quoted(text), prob, isNext: text === next })),
206  }))
207  const chipWidth = Math.max(8, Math.floor((width - 5) / v.k))
208  const nextProb = meta.next_prob[v.pos]
209  const selection =
210    ` ▸ ${v.pos}/${meta.tokens.length - 1}  "${quoted(shown(meta.tokens[v.pos] ?? ''))}"` +
211    (next !== null ? `  → next "${quoted(next)}"${nextProb != null ? ` (${(nextProb * 100).toFixed(1)}% at output)` : ''}` : '') +
212    `  ·  layers ${panelOffset}–${panelOffset + rows.length - 1} of ${lens.layers.length}`
213  const name = meta.path.split('/').pop() ?? meta.path
214  return {
215    width,
216    title: `${lens.name} · top ${v.k}`,
217    status: `${name} · ${meta.model}`,
218    lines: styled,
219    topRows,
220    selection,
221    rows,
222    panelRows,
223    chipWidth,
224    hint: v.error ?? hint,
225  }
226}
227
228const OPEN_SCHEMA = {
229  type: 'object',
230  properties: {
231    path: { type: 'string', description: 'Trace file: a lens-view.v0 JSON trace, or a saved grid (<id>.json.gz with vocab.json beside it). Absolute or relative to the working directory.' },
232    lens: { type: 'string', description: 'Lens to show first, e.g. jlens or logitlens' },
233    k: { type: 'number', description: 'How many top tokens to show per layer (default 5)' },
234    position: { type: 'number', description: 'Token position to select first (default 0)' },
235  },
236  required: ['path'],
237}
238
239const CONFIGURE_SCHEMA = {
240  type: 'object',
241  properties: {
242    k: { type: 'number', description: 'How many top tokens to show per layer' },
243    lens: { type: 'string', description: 'Lens to show, e.g. jlens or logitlens' },
244    position: { type: 'number', description: 'Token position to select' },
245  },
246}
247
248export const register: Register = on => {
249  on('session.start', async ($, e, next) => {
250    await $.command.register({
251      name: 'lens-view',
252      description: 'Lens traces in a pane: /lens-view <path> · /lens-view k <n> · /lens-view lens <name>',
253    })
254    await $.tool.register({
255      name: 'open',
256      description:
257        'Open a lens trace (J-lens, logit lens, tuned lens) in the lens-view pane beside the conversation: the ' +
258        'transcript on top, the selected token\'s top-k tokens per layer below. Use when the user asks to open ' +
259        'lens-view, the J-lens viewer, or a lens trace. Its contents are not returned to you.',
260      inputSchema: OPEN_SCHEMA,
261    })
262    await $.tool.register({
263      name: 'configure',
264      description:
265        'Change the open lens-view pane: k (how many top tokens per layer), the lens shown, or the selected token ' +
266        'position. Use when the user asks to show more or fewer top tokens, switch lens, or jump to a token.',
267      inputSchema: CONFIGURE_SCHEMA,
268    })
269    return next(e)
270  })
271
272  on('command.run', { command: 'lens-view' }, async ($, e) => {
273    const [word, value] = e.args.trim().split(/\s+/, 2)
274    let text = 'Opened lens-view.'
275    if (word === 'k' && value) text = await configure($, { k: Number(value) })
276    else if (word === 'lens' && value) text = await configure($, { lens: value })
277    else if (word) text = await load($, e.args.trim())
278    const v = await current($)
279    await $.ui.open({ id: PANE, title: v.path ? `lens-view · ${v.path.split('/').pop()}` : 'lens-view', focus: true })
280    return { text }
281  })
282
283  // A tool's arguments sit on the event itself, beside `tool`.
284  on('tool.call', { tool: 'mcp__lens-view__open' }, async ($, e) => {
285    const input = e as unknown as { path?: string; lens?: string; k?: number; position?: number }
286    if (!input.path) return { deny: 'path is required.' }
287    try {
288      const text = await load($, input.path, input)
289      // Not awaited: a pane opened unasked waits for a wide enough terminal (144 columns) to seat.
290      $.ui.open({ id: PANE, title: 'lens-view' }).catch(() => {})
291      return { result: `${text} The pane shows it to the user (it seats once the terminal is at least 144 columns wide, or when they run /lens-view); its contents were not returned to you.` }
292    } catch (exc) {
293      return { result: `lens-view failed: ${String(exc)}` }
294    }
295  })
296
297  on('tool.call', { tool: 'mcp__lens-view__configure' }, async ($, e) => {
298    const input = e as unknown as { k?: number; lens?: string; position?: number }
299    try {
300      return { result: await configure($, input) }
301    } catch (exc) {
302      return { result: `lens-view failed: ${String(exc)}` }
303    }
304  })
305
306  // The wheel: over the transcript it scrolls the transcript, over the top-k panel the layers.
307  // The pane draws exactly its body, so the engine has nothing of its own to scroll.
308  on('ui.scroll', { component: 'Pane', requestId: PANE }, async ($, e) => {
309    const row = e.pointer?.row
310    wheel = { events: wheel.events + 1, last: `by ${e.by}, row ${row ?? 'none (keys)'}, origin ${e.origin.kind}` }
311    if (row !== undefined && row > geometry.topRows + 1) await scrollPanel($, e.by) // title bar, transcript, selection bar
312    else await scrollTranscript($, e.by)
313    return {}
314  })
315
316  on('ui.message', { requestId: PANE }, async ($, e) => {
317    const data = e.data as { select?: number; key?: string }
318    if (typeof data.select === 'number') await select($, data.select)
319    else if (typeof data.key === 'string') await onKey($, data.key)
320    return {}
321  })
322
323  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
324    const v = await current($)
325    if (e.surface !== 'terminal' && e.surface !== 'desktop') {
326      const { Text } = $.ui.resolve(e)
327      return <Text>lens-view draws in the terminal and the desktop app. {v.path ?? ''}</Text>
328    }
329    const { Box, Text, Client } = $.ui.resolve(e)
330    const width = Math.max(40, e.props.bodyColumns)
331    const rows = Math.max(MIN_ROWS, e.props.scroll.bodyRows)
332    seen = `${e.surface}, ${e.props.placement}, body ${e.props.bodyColumns}×${e.props.scroll.bodyRows}, drawn ${width}×${rows}`
333    const props = draw(v, width, rows - 1)
334    return (
335      <Box flexDirection="column">
336        <Box flexDirection="row" justifyContent="space-between" height={1}>
337          <Text color="#7aa2f7" bold wrap="truncate">{props.title}</Text>
338          <Text color="#565f89" wrap="truncate">{props.status}</Text>
339        </Box>
340        <Client key="viewer" module="./viewer.tsx" props={props} width={width} height={rows - 1} />
341      </Box>
342    )
343  })
344}
345

1// Transcript layout: tokens wrapped into lines of at most `width` cells.
2
3export type Piece = { text: string; pos: number }
4
5// Control characters the surface refuses to draw: C0 (but newline, which layout breaks on), DEL and C1.
6const CONTROL = /[\x00-\x09\x0b-\x1f\x7f-\x9f]/g
7const hex = (c: string) => `\\x${c.charCodeAt(0).toString(16).padStart(2, '0')}`
8
9// How a token reads on screen: newlines and tabs made visible, other control characters escaped, empty tokens marked.
10export function shown(token: string): string {
11  if (token === '') return '∅'
12  return token.replace(/\t/g, '⇥').replace(/\r/g, '␍').replace(CONTROL, hex)
13}
14
15// A token as the body of a quoted string: JSON's escapes, plus DEL and C1, which JSON leaves raw.
16export function quoted(token: string): string {
17  return JSON.stringify(token).slice(1, -1).replace(/[\x7f-\x9f]/g, hex)
18}
19
20export function layout(tokens: string[], width: number): { lines: Piece[][]; lineOf: number[] } {
21  const lines: Piece[][] = [[]]
22  const lineOf: number[] = []
23  let used = 0
24  const newLine = () => {
25    lines.push([])
26    used = 0
27  }
28  tokens.forEach((token, pos) => {
29    const parts = shown(token).split('\n')
30    lineOf[pos] = lines.length - 1
31    parts.forEach((part, i) => {
32      // A newline inside the token: mark it, then break.
33      let text = i < parts.length - 1 ? part + '⏎' : part
34      if (text === '' && i > 0) return
35      if (used > 0 && used + text.length > width && text.length <= width) newLine()
36      if (i === 0) lineOf[pos] = lines.length - 1
37      while (text.length > 0) {
38        const room = width - used
39        if (room <= 0) {
40          newLine()
41          continue
42        }
43        lines[lines.length - 1]!.push({ text: text.slice(0, room), pos })
44        used += Math.min(room, text.length)
45        text = text.slice(room)
46      }
47      if (i < parts.length - 1) newLine()
48    })
49  })
50  return { lines, lineOf }
51}
52
53export function clamp(value: number, low: number, high: number): number {
54  return Math.max(low, Math.min(high, value))
55}
56

1// Draws the lens-view pane below its title bar: the transcript (top), the selected token's top-k per layer
2// (bottom). Clicks and keys go to the hooks module as posts; the wheel arrives there as ui.scroll.
3import type { ClientModule } from 'claude-code'
4
5import type { LensViewProps, LensViewSegment } from '../types'
6
7const INK = { text: '#c0caf5', alt: '#9aa5ce', dim: '#565f89', accent: '#7aa2f7', next: '#9ece6a', warn: '#e0af68' }
8const HEAT = ['#1f2335', '#283457', '#2e4482', '#3d59a1', '#4f6fc4']
9
10function heat(prob: number): string {
11  return HEAT[prob >= 0.5 ? 4 : prob >= 0.2 ? 3 : prob >= 0.05 ? 2 : prob >= 0.01 ? 1 : 0]!
12}
13
14function percent(prob: number): string {
15  if (prob >= 0.995) return '100%'
16  if (prob >= 0.1) return `${Math.round(prob * 100)}%`
17  if (prob >= 0.001) return `${(prob * 100).toFixed(1)}%`
18  return '<.1%'
19}
20
21function fit(text: string, width: number): string {
22  if (width <= 0) return ''
23  return text.length > width ? text.slice(0, Math.max(0, width - 1)) + '…' : text.padEnd(width)
24}
25
26const Viewer: ClientModule<LensViewProps> = (props, surface) => {
27  const { Box, Text } = surface.elements
28  const { width, lines, topRows, rows, panelRows, chipWidth } = props
29
30  surface.onPointer(event => {
31    if (event.type !== 'down' || event.button !== 'left') return
32    const line = lines[event.y]
33    if (event.y < 0 || event.y >= topRows || !line) return
34    let x = 0
35    for (const segment of line) {
36      if (event.x < x + segment.text.length) return surface.post({ select: segment.pos })
37      x += segment.text.length
38    }
39  })
40  surface.onKey(event => surface.post({ key: event.key, shift: event.shift === true }))
41
42  const segment = (s: LensViewSegment, i: number) =>
43    s.style === 'selected' ? (
44      <Text key={String(i)} backgroundColor={INK.accent} color="#1a1b26" bold>{s.text}</Text>
45    ) : (
46      <Text key={String(i)} color={s.style === 'surprise' ? INK.warn : s.style === 'odd' ? INK.alt : INK.text}>{s.text}</Text>
47    )
48
49  const transcript = Array.from({ length: topRows }, (_, row) => (
50    <Box key={`t${row}`} flexDirection="row" height={1}>
51      {(lines[row] ?? []).map(segment)}
52    </Box>
53  ))
54
55  const panel = Array.from({ length: panelRows }, (_, row) => {
56    const r = rows[row]
57    if (!r) return <Box key={`p${row}`} height={1} />
58    return (
59      <Box key={`p${row}`} flexDirection="row" height={1}>
60        <Text color={INK.dim}>{r.label.padStart(4)} </Text>
61        {r.chips.map((chip, i) => (
62          <Text key={String(i)} backgroundColor={heat(chip.prob)} color={chip.isNext ? INK.next : INK.text} bold={chip.isNext}>
63            {fit(chip.text, chipWidth - 6) + percent(chip.prob).padStart(5) + ' '}
64          </Text>
65        ))}
66      </Box>
67    )
68  })
69
70  return (
71    <Box flexDirection="column" width={width}>
72      {transcript}
73      <Box height={1} backgroundColor="#1f2335">
74        <Text color={INK.text} bold wrap="truncate">{fit(props.selection, width)}</Text>
75      </Box>
76      {panel}
77      <Box height={1}>
78        <Text color={INK.dim} wrap="truncate">{props.hint}</Text>
79      </Box>
80    </Box>
81  )
82}
83
84export default Viewer
85

1// What `lens-view-json` writes to its cache directory (src/lens_view/trace.py, export).
2export type LensViewMeta = {
3  format: string
4  path: string
5  model: string
6  prompt: string
7  tokens: string[]
8  targets: (string | null)[] | null
9  lenses: { name: string; layers: string[] }[]
10  k: number // top-k stored per cell
11  next_prob: (number | null)[] // final-layer probability of the actual next token, per position
12}
13
14// pos-<i>.json: lens name -> [layer] -> [token, prob], highest first.
15export type LensViewCell = Record<string, [string, number][][]>
16
17export type LensViewState = {
18  path: string | null
19  cache: string | null
20  meta: LensViewMeta | null
21  error: string | null
22  pos: number
23  k: number
24  lensIndex: number
25  topOffset: number // first transcript line shown
26  panelOffset: number // first layer row shown
27  cell: LensViewCell | null // the selected position's top-k
28}
29
30// Props of the Client that draws the pane (hooks/viewer.tsx).
31export type LensViewSegment = { text: string; pos: number; style: 'even' | 'odd' | 'selected' | 'surprise' }
32export type LensViewChip = { text: string; prob: number; isNext: boolean }
33export type LensViewProps = {
34  width: number
35  title: string
36  status: string
37  lines: LensViewSegment[][]
38  topRows: number
39  selection: string
40  rows: { label: string; chips: LensViewChip[] }[]
41  panelRows: number
42  chipWidth: number
43  hint: string
44}
45
46declare module 'claude-code' {
47  interface PluginState {
48    'lens-view': { view: LensViewState }
49  }
50}
51