selection-comment

1import { atom, read, update } from 'claude-code'
2import type { Register, SessionMessage } from 'claude-code'
3
4import type { Pending, Stored, View } from '../types'
5import { collapseDraft, expandDraft, findOrigin, flat, tokenFor, tokenRuns, viewOf } from './comment'
6
7const pending = atom({ plugin: 'selection-comment', key: 'pending' } as const, null)
8// Comments added this session, so each one gets the next id.
9const count = atom({ plugin: 'selection-comment', key: 'count' } as const, 0)
10// What each token in the draft stands for, by id.
11const comments = atom({ plugin: 'selection-comment', key: 'comments' } as const, {})
12const view = atom({ plugin: 'selection-comment', key: 'view' } as const, null)
13
14const POLL_MS = 300
15const PREVIEW_CHARS = 60
16const INPUT_KEY = 'comment-input'
17
18function preview(text: string): string {
19  return flat(text, PREVIEW_CHARS)
20}
21
22export const register: Register = on => {
23  // The last selection the person commented on or dismissed, so the band
24  // stays down until they select something else.
25  let handled: string | undefined
26  let bandId: string | undefined
27
28  on('session.start', async ($, e, next) => {
29    $.clock.every(POLL_MS, () => {
30      void (async () => {
31        const current = await read($, pending)
32        if (current?.isEditing) {
33          return
34        }
35
36        const selection = await $.ui.selection()
37        const text = selection?.text.trim()
38        if (!text || text === handled) {
39          if (current) {
40            await update($, pending, () => null)
41          }
42          return
43        }
44
45        if (current?.text !== text) {
46          await update($, pending, (): Pending => ({ text, requestId: selection?.requestId, isEditing: false }))
47        }
48      })()
49    })
50
51    return next(e)
52  })
53
54  // Keeps the tokens painted and the band's toggle in step with the draft.
55  on('prompt.edit', async ($, e, next) => {
56    const result = await next(e)
57    try {
58      const known = await read($, comments)
59      if (Object.keys(known).length === 0) {
60        return result
61      }
62      const shown = viewOf(result.text, known)
63      if (shown !== (await read($, view))) {
64        await update($, view, (): View => shown)
65      }
66      return { ...result, decorations: [...(result.decorations ?? []), ...tokenRuns(result.text, known)] }
67    } catch {
68      return result
69    }
70  })
71
72  // The model reads each token as its full block, with the note as the token
73  // has it now.
74  on('prompt.submit', async ($, e, next) => {
75    let text = e.text
76    try {
77      text = expandDraft(e.text, await read($, comments))
78      await update($, pending, () => null)
79      await update($, view, (): View => null)
80    } catch {
81      // Sent as typed.
82    }
83
84    return next(text === e.text ? e : { ...e, text })
85  })
86
87  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
88    bandId = e.requestId
89    const current = await read($, pending)
90    const shown = await read($, view)
91
92    // Mobile draws no Input yet, so the band stays the engine's there.
93    if (e.props.hasSurvey || (current === null && shown === null) || e.surface === 'mobile') {
94      return next(e)
95    }
96
97    const { Box, Button, Input, Text } = $.ui.resolve(e)
98    const nextId = `C${(await read($, count)) + 1}`
99
100    const toggle = async () => {
101      const box = await $.prompt.read()
102      const known = await read($, comments)
103      if (shown === 'tokens') {
104        const text = expandDraft(box.text, known)
105        await $.prompt.fill({ text, mode: 'replace' })
106        await update($, view, (): View => viewOf(text, known))
107        return
108      }
109      const collapsed = collapseDraft(box.text, known)
110      await update($, comments, () => collapsed.comments)
111      await $.prompt.fill({ text: collapsed.text, mode: 'replace', decorations: tokenRuns(collapsed.text, collapsed.comments) })
112      await update($, view, (): View => viewOf(collapsed.text, collapsed.comments))
113    }
114
115    const toggleButton =
116      shown === null ? null : (
117        <Button key="toggle" label={shown === 'tokens' ? 'Expand comments' : 'Collapse comments'} onPress={toggle} />
118      )
119
120    if (current === null) {
121      return <Box>{toggleButton}</Box>
122    }
123
124    const dismiss = async () => {
125      handled = current.text
126      await update($, pending, () => null)
127    }
128
129    if (!current.isEditing) {
130      return (
131        <Box>
132          <Text color="claude" bold>
133            “{preview(current.text)}”{' '}
134          </Text>
135          <Button
136            key="comment"
137            label={`Comment as ${nextId}`}
138            onPress={async () => {
139              await update($, pending, (): Pending => ({ ...current, isEditing: true }))
140              if (bandId) {
141                void $.ui.focus({ requestId: bandId, key: INPUT_KEY }).catch(() => undefined)
142              }
143            }}
144          />
145          <Text> </Text>
146          <Button key="dismiss" label="Dismiss" onPress={dismiss} />
147          {toggleButton && <Text> </Text>}
148          {toggleButton}
149        </Box>
150      )
151    }
152
153    return (
154      <Box flexDirection="column">
155        <Text color="claude" bold>
156          “{preview(current.text)}”
157        </Text>
158        <Box borderStyle="round" borderColor="suggestion" paddingX={1}>
159          <Text color="suggestion" bold>
160            {nextId}:{' '}
161          </Text>
162          <Input
163            key={INPUT_KEY}
164            placeholder="type a comment, Enter to add to prompt"
165            submitLabel="add"
166            autoFocus
167            onSubmit={async (value: string) => {
168              let n = 0
169              await update($, count, c => (n = c + 1))
170              const found = await $.session.messages().catch(() => [])
171              const messages = Array.isArray(found) ? (found as SessionMessage[]) : []
172              const stored: Stored = { id: `C${n}`, text: current.text, ...findOrigin(messages, current.text, current.requestId) }
173              await update($, comments, known => ({ ...known, [stored.id]: stored }))
174
175              // A token holds a line of its own, so it reads back whole.
176              const box = await $.prompt.read()
177              const prefix = box.text.length > 0 && !box.text.endsWith('\n') ? '\n' : ''
178              const token = tokenFor({ ...stored, note: value })
179              await $.prompt.fill({
180                text: `${prefix}${token}\n`,
181                mode: 'append',
182                decorations: [{ start: prefix.length, end: prefix.length + token.length, color: 'suggestion' }],
183              })
184              await update($, view, (): View => 'tokens')
185              await dismiss()
186            }}
187          />
188          <Text> </Text>
189          <Button key="cancel" label="Cancel" onPress={dismiss} />
190        </Box>
191      </Box>
192    )
193  })
194}
195

1import type { PromptDecoration, SessionMessage, ToolUseSummary } from 'claude-code'
2
3import type { Stored, View } from '../types'
4
5const SOURCE_CHARS = 60
6const QUOTE_MAX_LINES = 20
7// Longer selections carry their own context.
8const CONTEXT_MAX_LINES = 5
9// The tool input that best names a call, in order of preference.
10const SOURCE_KEYS = ['file_path', 'command', 'path', 'pattern', 'url', 'description']
11
12export type Comment = Stored & { note: string }
13
14export type Origin = { source?: string; context?: string }
15
16function quoteOf(text: string): string {
17  const lines = text.split('\n')
18
19  return lines.length > QUOTE_MAX_LINES
20    ? [...lines.slice(0, QUOTE_MAX_LINES), `[… ${lines.length - QUOTE_MAX_LINES} more lines]`].join('\n')
21    : text
22}
23
24export function formatComment(comment: Comment): string {
25  const quote = quoteOf(comment.text)
26  const source = comment.source ? ` source="${comment.source.replace(/"/g, "'")}"` : ''
27  const parts = [`<comment id="${comment.id}"${source}>`, '<quote>', quote, '</quote>']
28
29  if (comment.context) {
30    parts.push('<context>', comment.context, '</context>')
31  }
32  const note = comment.note.trim()
33  if (note.length > 0) {
34    parts.push(`<note>${note}</note>`)
35  }
36  parts.push('</comment>')
37
38  return `${parts.join('\n')}\n`
39}
40
41// The screen shows rendered markdown, so compare without its markers and
42// without the wrapping.
43function normalize(text: string): string {
44  return text.replace(/[*`_#>]/g, '').replace(/\s+/g, ' ').trim()
45}
46
47function locate(lines: string[], selected: string[]): { start: number; end: number } | undefined {
48  const first = normalize(selected[0] ?? '')
49  const last = normalize(selected.at(-1) ?? '')
50  if (normalize(selected.join(' ')).length < 3 || first.length === 0 || last.length === 0) {
51    return undefined
52  }
53
54  for (let start = 0; start < lines.length; start++) {
55    if (!normalize(lines[start]!).includes(first)) {
56      continue
57    }
58    // Rendering can split or join lines, so allow some slack for the end.
59    const limit = Math.min(lines.length, start + selected.length * 2 + 5)
60    for (let end = start; end < limit; end++) {
61      if (normalize(lines[end]!).includes(last)) {
62        return { start, end }
63      }
64    }
65  }
66
67  return undefined
68}
69
70function surround(lines: string[], hit: { start: number; end: number }, selectedLines: number): string | undefined {
71  if (selectedLines > CONTEXT_MAX_LINES) {
72    return undefined
73  }
74  const before = lines
75    .slice(0, hit.start)
76    .reverse()
77    .find(line => line.trim().length > 0)
78  const after = lines.slice(hit.end + 1).find(line => line.trim().length > 0)
79  if (before === undefined && after === undefined) {
80    return undefined
81  }
82
83  return [before?.trim(), '[SELECTION]', after?.trim()].filter(line => line !== undefined).join('\n')
84}
85
86export function flat(text: string, max: number): string {
87  const line = text.replace(/\s+/g, ' ').trim()
88
89  return line.length > max ? `${line.slice(0, max)}…` : line
90}
91
92function toolLabel(use: ToolUseSummary): string {
93  const key = SOURCE_KEYS.find(k => typeof use.input[k] === 'string' && (use.input[k] as string).trim().length > 0)
94
95  return key ? `${use.tool} · ${flat(use.input[key] as string, SOURCE_CHARS)}` : use.tool
96}
97
98function toolLines(use: ToolUseSummary): string[] {
99  const inputs = Object.values(use.input).filter((v): v is string => typeof v === 'string')
100
101  return [...inputs, use.text ?? ''].flatMap(text => text.split('\n'))
102}
103
104// A turn opens with each prompt the person typed; tool results ride on user
105// messages too but open none.
106function turnNumbers(messages: readonly SessionMessage[]): number[] {
107  let turn = 0
108
109  return messages.map(message => {
110    if (message.role === 'user' && !message.toolResults?.length && message.text.trim().length > 0) {
111      turn++
112    }
113    return Math.max(turn, 1)
114  })
115}
116
117/**
118 * Where the selection came from: the tool call its row belongs to, or else
119 * the newest message or tool output holding it, plus the lines around it.
120 */
121export function findOrigin(messages: readonly SessionMessage[], text: string, requestId?: string): Origin {
122  const selected = text
123    .split('\n')
124    .map(line => line.trim())
125    .filter(line => line.length > 0)
126  const turns = turnNumbers(messages)
127  const match = (lines: string[], source: string): Origin | undefined => {
128    const hit = locate(lines, selected)
129    return hit ? { source, context: surround(lines, hit, selected.length) } : undefined
130  }
131
132  if (requestId) {
133    for (let i = messages.length - 1; i >= 0; i--) {
134      const use = messages[i]!.toolUses.find(u => u.tool_use_id === requestId)
135      if (use) {
136        const source = `${toolLabel(use)} · turn ${turns[i]}`
137        return match(toolLines(use), source) ?? { source }
138      }
139    }
140  }
141
142  for (let i = messages.length - 1; i >= 0; i--) {
143    const message = messages[i]!
144    const found =
145      match(message.text.split('\n'), `${message.role} · turn ${turns[i]}`) ??
146      message.toolUses.map(use => match(toolLines(use), `${toolLabel(use)} · turn ${turns[i]}`)).find(Boolean)
147    if (found) {
148      return found
149    }
150  }
151
152  return {}
153}
154
155const TOKEN_QUOTE_CHARS = 40
156const NOTE_MARK = ' → '
157// A token sits on a line of its own: `[C1 · source: "quote…" → note]`.
158const TOKEN = /^\[(C\d+)[ :][^\n]*\]$/gm
159const BLOCK = /<comment id="(C\d+)"[^>\n]*>\n[\s\S]*?\n<\/comment>/g
160
161// Keeps the token's own marks out of what it shows.
162function clean(text: string): string {
163  return text.replace(/\[/g, '(').replace(/\]/g, ')').replace(/→/g, '->').replace(/"/g, "'")
164}
165
166export function tokenFor(comment: Comment): string {
167  const head = clean([comment.id, comment.source].filter(Boolean).join(' · '))
168  const note = flat(comment.note, Infinity)
169
170  return `[${head}: "${clean(flat(comment.text, TOKEN_QUOTE_CHARS))}"${note ? `${NOTE_MARK}${note}` : ''}]`
171}
172
173// The note is everything after the first arrow, up to the closing bracket, so
174// the person may type anything there, brackets included.
175function noteOf(token: string): string {
176  const at = token.indexOf(NOTE_MARK)
177
178  return at < 0 ? '' : token.slice(at + NOTE_MARK.length, -1).trim()
179}
180
181/** Each token of a known comment, replaced by its full block. */
182export function expandDraft(text: string, comments: Record<string, Stored>): string {
183  return text.replace(TOKEN, (token, id: string) => {
184    const stored = comments[id]
185    return stored ? formatComment({ ...stored, note: noteOf(token) }).trimEnd() : token
186  })
187}
188
189function parseBlock(block: string, stored: Stored): Comment | undefined {
190  const quote = /<quote>\n([\s\S]*?)\n<\/quote>/.exec(block)?.[1]
191  if (quote === undefined) {
192    return undefined
193  }
194
195  return {
196    id: stored.id,
197    // A cut quote stands for the whole one while it is left as it was.
198    text: quote === quoteOf(stored.text) ? stored.text : quote,
199    source: /^<comment [^>\n]*source="([^"]*)"/.exec(block)?.[1],
200    context: /<context>\n([\s\S]*?)\n<\/context>/.exec(block)?.[1],
201    note: /<note>([\s\S]*?)<\/note>/.exec(block)?.[1]?.trim() ?? '',
202  }
203}
204
205/**
206 * Each full block of a known comment, read back with the person's edits and
207 * replaced by its token; a block too mangled to read stays as it is.
208 */
209export function collapseDraft(
210  text: string,
211  comments: Record<string, Stored>,
212): { text: string; comments: Record<string, Stored> } {
213  const next = { ...comments }
214  const collapsed = text.replace(BLOCK, (block, id: string) => {
215    const stored = comments[id]
216    const comment = stored && parseBlock(block, stored)
217    if (!comment) {
218      return block
219    }
220    const { note: _, ...rest } = comment
221    next[id] = rest
222    return tokenFor(comment)
223  })
224
225  return { text: collapsed, comments: next }
226}
227
228function known(text: string, pattern: RegExp, comments: Record<string, Stored>): RegExpMatchArray[] {
229  return [...text.matchAll(pattern)].filter(match => comments[match[1]!] !== undefined)
230}
231
232/** What the draft shows of the comments; tokens win when it holds both. */
233export function viewOf(text: string, comments: Record<string, Stored>): View {
234  if (known(text, TOKEN, comments).length > 0) {
235    return 'tokens'
236  }
237
238  return known(text, BLOCK, comments).length > 0 ? 'blocks' : null
239}
240
241/** A run painting each token of a known comment as a chip. */
242export function tokenRuns(text: string, comments: Record<string, Stored>): PromptDecoration[] {
243  return known(text, TOKEN, comments).map(match => ({
244    start: match.index!,
245    end: match.index! + match[0].length,
246    color: 'suggestion',
247  }))
248}
249

1export type Pending = { text: string; requestId?: string; isEditing: boolean }
2
3// A comment as the prompt box's short token stands for it; its note lives in
4// the token, so the person can edit it in place.
5export type Stored = { id: string; text: string; source?: string; context?: string }
6
7// What the draft shows of the comments: short tokens, full blocks, or none.
8export type View = 'tokens' | 'blocks' | null
9
10declare module 'claude-code' {
11  interface PluginState {
12    'selection-comment': {
13      pending: Pending | null
14      count: number
15      comments: Record<string, Stored>
16      view: View
17    }
18  }
19}
20