assistant.go

  1package chat
  2
  3import (
  4	"encoding/binary"
  5	"fmt"
  6	"hash/fnv"
  7	"strings"
  8
  9	tea "charm.land/bubbletea/v2"
 10	"charm.land/lipgloss/v2"
 11	"github.com/charmbracelet/crush/internal/message"
 12	"github.com/charmbracelet/crush/internal/ui/anim"
 13	"github.com/charmbracelet/crush/internal/ui/common"
 14	"github.com/charmbracelet/crush/internal/ui/styles"
 15	"github.com/charmbracelet/x/ansi"
 16)
 17
 18// assistantMessageTruncateFormat is the text shown when an assistant message is
 19// truncated.
 20const assistantMessageTruncateFormat = "… (%d lines hidden) [click or space to expand]"
 21
 22// maxCollapsedThinkingHeight defines the maximum height of the thinking
 23const maxCollapsedThinkingHeight = 10
 24
 25// assistantSection is a per-section render cache for AssistantMessageItem.
 26// Each section (thinking, content, error) carries its own keys so that
 27// streaming a section does not invalidate a different — often more
 28// expensive — section's cached render. srcHash is an FNV-64 of the
 29// section's source text; extra captures any other state that changes
 30// the rendered output (e.g. thinkingExpanded, the thinking footer
 31// inputs). valid disambiguates a real cache hit from the zero value
 32// when both source text and extras hash to zero. aux carries any
 33// per-section side data that the caller needs to recover on a hit
 34// (e.g. the thinking box height for click detection).
 35type assistantSection struct {
 36	width   int
 37	srcHash uint64
 38	extra   uint64
 39	out     string
 40	h       int
 41	aux     int
 42	valid   bool
 43}
 44
 45// hit reports whether the cache entry matches the requested key.
 46func (s *assistantSection) hit(width int, srcHash, extra uint64) bool {
 47	return s.valid && s.width == width && s.srcHash == srcHash && s.extra == extra
 48}
 49
 50// store records the rendered output under the given key.
 51func (s *assistantSection) store(width int, srcHash, extra uint64, out string, aux int) {
 52	s.width = width
 53	s.srcHash = srcHash
 54	s.extra = extra
 55	s.out = out
 56	s.h = lipgloss.Height(out)
 57	s.aux = aux
 58	s.valid = true
 59}
 60
 61// reset drops the cached output.
 62func (s *assistantSection) reset() {
 63	*s = assistantSection{}
 64}
 65
 66// fnv64 hashes a single string with FNV-64.
 67func fnv64(s string) uint64 {
 68	h := fnv.New64a()
 69	_, _ = h.Write([]byte(s))
 70	return h.Sum64()
 71}
 72
 73// fnvFields hashes a list of byte fields with length-prefix framing
 74// so that no concatenation collision can occur between distinct
 75// field tuples (a NUL inside one field cannot impersonate a
 76// boundary between two fields). Each field is preceded by its
 77// length encoded as 8 bytes little-endian.
 78func fnvFields(fields ...[]byte) uint64 {
 79	h := fnv.New64a()
 80	var lenBuf [8]byte
 81	for _, f := range fields {
 82		binary.LittleEndian.PutUint64(lenBuf[:], uint64(len(f)))
 83		_, _ = h.Write(lenBuf[:])
 84		_, _ = h.Write(f)
 85	}
 86	return h.Sum64()
 87}
 88
 89// AssistantMessageItem represents an assistant message in the chat UI.
 90//
 91// This item includes thinking, and the content but does not include the tool calls.
 92type AssistantMessageItem struct {
 93	*highlightableMessageItem
 94	*cachedMessageItem
 95	*focusableMessageItem
 96
 97	message           *message.Message
 98	sty               *styles.Styles
 99	anim              *anim.Anim
100	thinkingExpanded  bool
101	thinkingBoxHeight int // Tracks the rendered thinking box height for click detection.
102
103	// Per-section render caches. Splitting these out means content
104	// streaming does not invalidate the (often expensive) thinking
105	// render, and vice versa.
106	thinkingSec assistantSection
107	contentSec  assistantSection
108	errorSec    assistantSection
109}
110
111var _ Expandable = (*AssistantMessageItem)(nil)
112
113// NewAssistantMessageItem creates a new AssistantMessageItem.
114func NewAssistantMessageItem(sty *styles.Styles, message *message.Message) MessageItem {
115	a := &AssistantMessageItem{
116		highlightableMessageItem: defaultHighlighter(sty),
117		cachedMessageItem:        &cachedMessageItem{},
118		focusableMessageItem:     &focusableMessageItem{},
119		message:                  message,
120		sty:                      sty,
121	}
122
123	a.anim = anim.New(anim.Settings{
124		ID:          a.ID(),
125		Size:        15,
126		GradColorA:  sty.WorkingGradFromColor,
127		GradColorB:  sty.WorkingGradToColor,
128		LabelColor:  sty.WorkingLabelColor,
129		CycleColors: true,
130	})
131	return a
132}
133
134// StartAnimation starts the assistant message animation if it should be spinning.
135func (a *AssistantMessageItem) StartAnimation() tea.Cmd {
136	if !a.isSpinning() {
137		return nil
138	}
139	return a.anim.Start()
140}
141
142// Animate progresses the assistant message animation if it should be spinning.
143func (a *AssistantMessageItem) Animate(msg anim.StepMsg) tea.Cmd {
144	if !a.isSpinning() {
145		return nil
146	}
147	return a.anim.Animate(msg)
148}
149
150// ID implements MessageItem.
151func (a *AssistantMessageItem) ID() string {
152	return a.message.ID
153}
154
155// RawRender implements [MessageItem].
156func (a *AssistantMessageItem) RawRender(width int) string {
157	cappedWidth := cappedMessageWidth(width)
158
159	var spinner string
160	if a.isSpinning() {
161		spinner = a.renderSpinning()
162	}
163
164	content, height := a.renderMessageContent(cappedWidth)
165	highlightedContent := a.renderHighlighted(content, cappedWidth, height)
166	if spinner != "" {
167		if highlightedContent != "" {
168			highlightedContent += "\n\n"
169		}
170		return highlightedContent + spinner
171	}
172
173	return highlightedContent
174}
175
176// Render implements MessageItem.
177func (a *AssistantMessageItem) Render(width int) string {
178	// XXX: Here, we're manually applying the focused/blurred styles because
179	// using lipgloss.Render can degrade performance for long messages due to
180	// it's wrapping logic.
181	// We already know that the content is wrapped to the correct width in
182	// RawRender, so we can just apply the styles directly to each line.
183	//
184	// The split + per-line prefix loop is O(L); cache the result keyed
185	// by (width, focused, sectionsFingerprint) so steady-state Render
186	// becomes a pointer return. The sectionsFingerprint folds in the
187	// per-section srcHash/extra so that any sub-cache change
188	// invalidates this prefix cache without requiring an explicit
189	// drop. Bypass the cache while spinning (RawRender's spinner
190	// suffix changes every animation frame) or while a highlight
191	// range is active (selection drag).
192	useCache := !a.isSpinning() && !a.isHighlighted()
193	cappedWidth := cappedMessageWidth(width)
194	key := a.prefixCacheKey(cappedWidth)
195	if useCache {
196		if cached, ok := a.getCachedPrefixedRender(width, key); ok {
197			return cached
198		}
199	}
200	focused := a.sty.Messages.AssistantFocused.Render()
201	blurred := a.sty.Messages.AssistantBlurred.Render()
202	rendered := a.RawRender(width)
203	lines := strings.Split(rendered, "\n")
204	for i, line := range lines {
205		if a.focused {
206			lines[i] = focused + line
207		} else {
208			lines[i] = blurred + line
209		}
210	}
211	out := strings.Join(lines, "\n")
212	if useCache {
213		a.setCachedPrefixedRender(out, width, key)
214	}
215	return out
216}
217
218// prefixCacheKey builds the F3 prefixed-render cache key. We pack the
219// focus bit into bit 0 and a fingerprint of the section caches into
220// the upper bits, so any change to a sub-section's source text or
221// extras forces the prefix cache to miss without needing an explicit
222// drop. cappedWidth is included so a cached prefix never survives a
223// section-cache miss caused by a width change. The finish reason is
224// folded in too because it controls the composition of
225// renderMessageContent (e.g. appending the constant "Canceled"
226// string) — that decision lives outside any section's own hash.
227func (a *AssistantMessageItem) prefixCacheKey(cappedWidth int) uint64 {
228	thinkSrc, thinkExtra := a.thinkingKey()
229	contentSrc, contentExtra := a.contentKey()
230	errSrc, errExtra := a.errorKey()
231	h := fnv.New64a()
232	var buf [8]byte
233	writeU64 := func(v uint64) {
234		for i := range 8 {
235			buf[i] = byte(v >> (8 * i))
236		}
237		_, _ = h.Write(buf[:])
238	}
239	writeU64(uint64(cappedWidth))
240	writeU64(thinkSrc)
241	writeU64(thinkExtra)
242	writeU64(contentSrc)
243	writeU64(contentExtra)
244	writeU64(errSrc)
245	writeU64(errExtra)
246	writeU64(a.compositionKey())
247	fingerprint := h.Sum64()
248	var focusBit uint64
249	if a.focused {
250		focusBit = 1
251	}
252	return (fingerprint &^ 1) | focusBit
253}
254
255// compositionKey hashes the inputs to renderMessageContent's structural
256// decisions (which sections to include, whether to append the
257// constant "Canceled" footer) so that flipping IsFinished or the
258// finish reason invalidates the prefix cache even when no section's
259// own source text changed.
260func (a *AssistantMessageItem) compositionKey() uint64 {
261	var finishedFlag byte
262	var reason string
263	if a.message.IsFinished() {
264		finishedFlag = 1
265		reason = string(a.message.FinishReason())
266	}
267	// Length-prefixed framing keeps the finished flag and the reason
268	// string from blending into one another.
269	return fnvFields([]byte{finishedFlag}, []byte(reason))
270}
271
272// renderMessageContent renders the message content including thinking, main
273// content, and finish reason. Each section is served from its own cache;
274// only the section whose source text or extras changed since the last
275// render is recomputed.
276func (a *AssistantMessageItem) renderMessageContent(width int) (string, int) {
277	var messageParts []string
278	thinking := strings.TrimSpace(a.message.ReasoningContent().Thinking)
279	content := strings.TrimSpace(a.message.Content().Text)
280
281	if thinking != "" {
282		messageParts = append(messageParts, a.cachedThinking(width))
283	}
284
285	if content != "" {
286		if thinking != "" {
287			messageParts = append(messageParts, "")
288		}
289		messageParts = append(messageParts, a.cachedContent(width))
290	}
291
292	if a.message.IsFinished() {
293		switch a.message.FinishReason() {
294		case message.FinishReasonCanceled:
295			messageParts = append(messageParts, a.sty.Messages.AssistantCanceled.Render("Canceled"))
296		case message.FinishReasonError:
297			messageParts = append(messageParts, a.cachedError(width))
298		}
299	}
300
301	out := strings.Join(messageParts, "\n")
302	return out, lipgloss.Height(out)
303}
304
305// thinkingKey returns the (srcHash, extra) cache key components for the
306// thinking section. extra folds in everything other than the raw
307// thinking text that affects the rendered output: the expanded flag
308// and the footer state (which depends on IsThinking, ToolCalls, and
309// ThinkingDuration).
310func (a *AssistantMessageItem) thinkingKey() (uint64, uint64) {
311	thinking := a.message.ReasoningContent().Thinking
312	srcHash := fnv64(thinking)
313
314	showFooter := !a.message.IsThinking() || len(a.message.ToolCalls()) > 0
315	var durationStr string
316	if showFooter {
317		duration := a.message.ThinkingDuration()
318		if duration.String() != "0s" {
319			durationStr = duration.String()
320		}
321	}
322	var expanded byte
323	if a.thinkingExpanded {
324		expanded = 1
325	}
326	var footer byte
327	if showFooter {
328		footer = 1
329	}
330	// Length-prefixed framing avoids any delimiter collision between
331	// the flag bytes and the duration string.
332	extra := fnvFields([]byte{expanded, footer}, []byte(durationStr))
333	return srcHash, extra
334}
335
336// contentKey returns the (srcHash, extra) cache key components for the
337// main content section.
338func (a *AssistantMessageItem) contentKey() (uint64, uint64) {
339	return fnv64(a.message.Content().Text), 0
340}
341
342// errorKey returns the (srcHash, extra) cache key components for the
343// error section. Returns (0, 0) when no error is present so the cache
344// stays a no-op for non-error messages.
345func (a *AssistantMessageItem) errorKey() (uint64, uint64) {
346	if !a.message.IsFinished() || a.message.FinishReason() != message.FinishReasonError {
347		return 0, 0
348	}
349	finishPart := a.message.FinishPart()
350	if finishPart == nil {
351		return 0, 0
352	}
353	// Length-prefixed framing prevents Message+Details collisions
354	// between distinct (Message, Details) tuples that would
355	// otherwise concatenate to the same byte sequence.
356	return fnvFields([]byte(finishPart.Message), []byte(finishPart.Details)), 0
357}
358
359// cachedThinking returns the rendered thinking section, computing and
360// caching it on miss. The thinking-box height (used for click target
361// detection) is preserved across hits via assistantSection.aux so the
362// cached path never desyncs click detection.
363func (a *AssistantMessageItem) cachedThinking(width int) string {
364	srcHash, extra := a.thinkingKey()
365	if a.thinkingSec.hit(width, srcHash, extra) {
366		a.thinkingBoxHeight = a.thinkingSec.aux
367		return a.thinkingSec.out
368	}
369	out := a.renderThinking(a.message.ReasoningContent().Thinking, width)
370	a.thinkingSec.store(width, srcHash, extra, out, a.thinkingBoxHeight)
371	return out
372}
373
374// cachedContent returns the rendered content section.
375func (a *AssistantMessageItem) cachedContent(width int) string {
376	srcHash, extra := a.contentKey()
377	if a.contentSec.hit(width, srcHash, extra) {
378		return a.contentSec.out
379	}
380	out := a.renderMarkdown(a.message.Content().Text, width)
381	a.contentSec.store(width, srcHash, extra, out, 0)
382	return out
383}
384
385// cachedError returns the rendered error section.
386func (a *AssistantMessageItem) cachedError(width int) string {
387	srcHash, extra := a.errorKey()
388	if a.errorSec.hit(width, srcHash, extra) {
389		return a.errorSec.out
390	}
391	out := a.renderError(width)
392	a.errorSec.store(width, srcHash, extra, out, 0)
393	return out
394}
395
396// renderThinking renders the thinking/reasoning content with footer.
397func (a *AssistantMessageItem) renderThinking(thinking string, width int) string {
398	renderer := common.QuietMarkdownRenderer(a.sty, width)
399	rendered, err := renderer.Render(thinking)
400	if err != nil {
401		rendered = thinking
402	}
403	rendered = strings.TrimSpace(rendered)
404
405	lines := strings.Split(rendered, "\n")
406	totalLines := len(lines)
407
408	isTruncated := totalLines > maxCollapsedThinkingHeight
409	if !a.thinkingExpanded && isTruncated {
410		lines = lines[totalLines-maxCollapsedThinkingHeight:]
411		hint := a.sty.Messages.ThinkingTruncationHint.Render(
412			fmt.Sprintf(assistantMessageTruncateFormat, totalLines-maxCollapsedThinkingHeight),
413		)
414		lines = append([]string{hint, ""}, lines...)
415	}
416
417	thinkingStyle := a.sty.Messages.ThinkingBox.Width(width)
418	result := thinkingStyle.Render(strings.Join(lines, "\n"))
419	a.thinkingBoxHeight = lipgloss.Height(result)
420
421	var footer string
422	// if thinking is done add the thought for footer
423	if !a.message.IsThinking() || len(a.message.ToolCalls()) > 0 {
424		duration := a.message.ThinkingDuration()
425		if duration.String() != "0s" {
426			footer = a.sty.Messages.ThinkingFooterTitle.Render("Thought for ") +
427				a.sty.Messages.ThinkingFooterDuration.Render(duration.String())
428		}
429	}
430
431	if footer != "" {
432		result += "\n\n" + footer
433	}
434
435	return result
436}
437
438// renderMarkdown renders content as markdown.
439func (a *AssistantMessageItem) renderMarkdown(content string, width int) string {
440	renderer := common.MarkdownRenderer(a.sty, width)
441	result, err := renderer.Render(content)
442	if err != nil {
443		return content
444	}
445	return strings.TrimSuffix(result, "\n")
446}
447
448func (a *AssistantMessageItem) renderSpinning() string {
449	if a.message.IsThinking() {
450		a.anim.SetLabel("Thinking")
451	} else if a.message.IsSummaryMessage {
452		a.anim.SetLabel("Summarizing")
453	}
454	return a.anim.Render()
455}
456
457// renderError renders an error message.
458func (a *AssistantMessageItem) renderError(width int) string {
459	finishPart := a.message.FinishPart()
460	errTag := a.sty.Messages.ErrorTag.Render("ERROR")
461	truncated := ansi.Truncate(finishPart.Message, width-2-lipgloss.Width(errTag), "...")
462	title := fmt.Sprintf("%s %s", errTag, a.sty.Messages.ErrorTitle.Render(truncated))
463	details := a.sty.Messages.ErrorDetails.Width(width - 2).Render(finishPart.Details)
464	return fmt.Sprintf("%s\n\n%s", title, details)
465}
466
467// isSpinning returns true if the assistant message is still generating.
468func (a *AssistantMessageItem) isSpinning() bool {
469	isThinking := a.message.IsThinking()
470	isFinished := a.message.IsFinished()
471	hasContent := strings.TrimSpace(a.message.Content().Text) != ""
472	hasToolCalls := len(a.message.ToolCalls()) > 0
473	return (isThinking || !isFinished) && !hasContent && !hasToolCalls
474}
475
476// SetMessage is used to update the underlying message. Only the
477// sub-section caches whose source text or extras changed are
478// invalidated; the others survive and serve cache hits on the next
479// RawRender.
480func (a *AssistantMessageItem) SetMessage(msg *message.Message) tea.Cmd {
481	wasSpinning := a.isSpinning()
482	a.message = msg
483	// The prefix cache is keyed by a fingerprint that includes every
484	// section's source hash, so an unchanged section keeps its prefix
485	// cache valid while a changed section forces a miss naturally.
486	// Section caches themselves are content-keyed, so they do not
487	// need an explicit drop here either.
488	if !wasSpinning && a.isSpinning() {
489		return a.StartAnimation()
490	}
491	return nil
492}
493
494// clearCache drops every cached render for this item, including the
495// per-section caches. Shadows the embedded cachedMessageItem.clearCache
496// so ClearItemCaches (style change) wipes the section caches too.
497func (a *AssistantMessageItem) clearCache() {
498	a.cachedMessageItem.clearCache()
499	a.thinkingSec.reset()
500	a.contentSec.reset()
501	a.errorSec.reset()
502}
503
504// ToggleExpanded toggles the expanded state of the thinking box and returns
505// whether the item is now expanded. Both the thinking section cache and
506// the F3 prefix cache key fold in thinkingExpanded (via the section's
507// extra hash and the prefix cache fingerprint respectively), so no
508// explicit invalidation is required.
509func (a *AssistantMessageItem) ToggleExpanded() bool {
510	a.thinkingExpanded = !a.thinkingExpanded
511	return a.thinkingExpanded
512}
513
514// HandleMouseClick implements MouseClickable. It signals (via a true return)
515// that the click lies on the thinking box so the caller can invoke
516// [AssistantMessageItem.ToggleExpanded] through the generic [Expandable]
517// path. Toggling here directly would double-toggle because the caller always
518// runs the generic path after a handled click.
519func (a *AssistantMessageItem) HandleMouseClick(btn ansi.MouseButton, x, y int) bool {
520	if btn != ansi.MouseLeft {
521		return false
522	}
523	// Only the thinking box is clickable; other regions of the assistant
524	// message should not trigger expansion.
525	return a.thinkingBoxHeight > 0 && y < a.thinkingBoxHeight
526}
527
528// HandleKeyEvent implements KeyEventHandler.
529func (a *AssistantMessageItem) HandleKeyEvent(key tea.KeyMsg) (bool, tea.Cmd) {
530	if k := key.String(); k == "c" || k == "y" {
531		text := a.message.Content().Text
532		return true, common.CopyToClipboard(text, "Message copied to clipboard")
533	}
534	return false, nil
535}