1package chat
2
3import (
4 "fmt"
5 "image"
6 "strings"
7 "time"
8
9 tea "charm.land/bubbletea/v2"
10 "charm.land/catwalk/pkg/catwalk"
11 "charm.land/lipgloss/v2"
12 "github.com/charmbracelet/crush/internal/config"
13 "github.com/charmbracelet/crush/internal/message"
14 "github.com/charmbracelet/crush/internal/ui/anim"
15 "github.com/charmbracelet/crush/internal/ui/attachments"
16 "github.com/charmbracelet/crush/internal/ui/common"
17 "github.com/charmbracelet/crush/internal/ui/list"
18 "github.com/charmbracelet/crush/internal/ui/styles"
19)
20
21// MessageLeftPaddingTotal is the total width that is taken up by the border +
22// padding. We also cap the width so text is readable to the maxTextWidth(120).
23const MessageLeftPaddingTotal = 2
24
25// maxTextWidth is the maximum width text messages can be
26const maxTextWidth = 120
27
28// Identifiable is an interface for items that can provide a unique identifier.
29type Identifiable interface {
30 ID() string
31}
32
33// Animatable is an interface for items that support animation.
34type Animatable interface {
35 StartAnimation() tea.Cmd
36 Animate(msg anim.StepMsg) tea.Cmd
37}
38
39// Expandable is an interface for items that can be expanded or collapsed.
40type Expandable interface {
41 // ToggleExpanded toggles the expanded state of the item. It returns
42 // whether the item is now expanded.
43 ToggleExpanded() bool
44}
45
46// KeyEventHandler is an interface for items that can handle key events.
47type KeyEventHandler interface {
48 HandleKeyEvent(key tea.KeyMsg) (bool, tea.Cmd)
49}
50
51// MessageItem represents a [message.Message] item that can be displayed in the
52// UI and be part of a [list.List] identifiable by a unique ID.
53type MessageItem interface {
54 list.Item
55 list.RawRenderable
56 Identifiable
57}
58
59// HighlightableMessageItem is a message item that supports highlighting.
60type HighlightableMessageItem interface {
61 MessageItem
62 list.Highlightable
63}
64
65// FocusableMessageItem is a message item that supports focus.
66type FocusableMessageItem interface {
67 MessageItem
68 list.Focusable
69}
70
71// SendMsg represents a message to send a chat message.
72type SendMsg struct {
73 Text string
74 Attachments []message.Attachment
75}
76
77type highlightableMessageItem struct {
78 // version is the parent item's version counter. SetHighlight
79 // bumps it on every observable change so the F6 list memo and
80 // any frozen entry get invalidated when a selection drag enters
81 // or leaves the item.
82 version *list.Versioned
83
84 startLine int
85 startCol int
86 endLine int
87 endCol int
88 highlighter list.Highlighter
89}
90
91var _ list.Highlightable = (*highlightableMessageItem)(nil)
92
93// isHighlighted returns true if the item has a highlight range set.
94func (h *highlightableMessageItem) isHighlighted() bool {
95 return h.startLine != -1 || h.endLine != -1
96}
97
98// renderHighlighted highlights the content if necessary.
99func (h *highlightableMessageItem) renderHighlighted(content string, width, height int) string {
100 if !h.isHighlighted() {
101 return content
102 }
103 area := image.Rect(0, 0, width, height)
104 return list.Highlight(content, area, h.startLine, h.startCol, h.endLine, h.endCol, h.highlighter)
105}
106
107// SetHighlight implements list.Highlightable.
108func (h *highlightableMessageItem) SetHighlight(startLine int, startCol int, endLine int, endCol int) {
109 // Adjust columns for the style's left inset (border + padding) since we
110 // highlight the content only.
111 offset := MessageLeftPaddingTotal
112 newStartCol := max(0, startCol-offset)
113 newEndCol := endCol
114 if endCol >= 0 {
115 newEndCol = max(0, endCol-offset)
116 }
117 if h.startLine == startLine && h.startCol == newStartCol && h.endLine == endLine && h.endCol == newEndCol {
118 return
119 }
120 h.startLine = startLine
121 h.startCol = newStartCol
122 h.endLine = endLine
123 h.endCol = newEndCol
124 if h.version != nil {
125 h.version.Bump()
126 }
127}
128
129// Highlight implements list.Highlightable.
130func (h *highlightableMessageItem) Highlight() (startLine int, startCol int, endLine int, endCol int) {
131 return h.startLine, h.startCol, h.endLine, h.endCol
132}
133
134func defaultHighlighter(sty *styles.Styles, v *list.Versioned) *highlightableMessageItem {
135 return &highlightableMessageItem{
136 version: v,
137 startLine: -1,
138 startCol: -1,
139 endLine: -1,
140 endCol: -1,
141 highlighter: list.ToHighlighter(sty.TextSelection),
142 }
143}
144
145// cacheClearable is implemented by message items that cache rendered
146// output and can be asked to drop the cache.
147type cacheClearable interface {
148 clearCache()
149}
150
151// ClearItemCaches drops any cached rendered output on each item so the
152// next render uses the current styles. It also bumps each item's
153// version so the F6 list-level memo invalidates frozen entries on
154// the next render.
155func ClearItemCaches(items []MessageItem) {
156 for _, item := range items {
157 if cc, ok := item.(cacheClearable); ok {
158 cc.clearCache()
159 }
160 if v, ok := item.(interface{ Bump() }); ok {
161 v.Bump()
162 }
163 }
164}
165
166// cachedMessageItem caches rendered message content to avoid re-rendering.
167//
168// This should be used by any message that can store a cached version of its render. e.x user,assistant... and so on
169//
170// THOUGHT(kujtim): we should consider if its efficient to store the render for different widths
171// the issue with that could be memory usage
172type cachedMessageItem struct {
173 // rendered is the cached rendered string
174 rendered string
175 // width and height are the dimensions of the cached render
176 width int
177 height int
178
179 // prefixedRendered caches the per-line-prefixed Render output (the
180 // result of splitting RawRender by newlines and prepending a focus
181 // or selection prefix to every line). Items rebuild this every
182 // frame today; caching it keyed by (prefixedWidth, prefixedKey)
183 // turns Render into a pointer return when item state is stable.
184 //
185 // Invalidation lives in clearCache; callers must additionally
186 // bypass this cache whenever the prefixed output would not be
187 // stable (spinner ticks, active highlight ranges) by not calling
188 // setCachedPrefixedRender for those frames.
189 prefixedRendered string
190 prefixedWidth int
191 prefixedKey uint64
192}
193
194// getCachedRender returns the cached render if it exists for the given width.
195func (c *cachedMessageItem) getCachedRender(width int) (string, int, bool) {
196 if c.width == width && c.rendered != "" {
197 return c.rendered, c.height, true
198 }
199 return "", 0, false
200}
201
202// setCachedRender sets the cached render.
203func (c *cachedMessageItem) setCachedRender(rendered string, width, height int) {
204 c.rendered = rendered
205 c.width = width
206 c.height = height
207}
208
209// getCachedPrefixedRender returns the cached prefixed render if it exists
210// for the given (width, key). The key encodes any state that changes the
211// per-line prefix (focused/blurred, compact, ...).
212func (c *cachedMessageItem) getCachedPrefixedRender(width int, key uint64) (string, bool) {
213 if c.prefixedRendered != "" && c.prefixedWidth == width && c.prefixedKey == key {
214 return c.prefixedRendered, true
215 }
216 return "", false
217}
218
219// setCachedPrefixedRender stores the cached prefixed render.
220func (c *cachedMessageItem) setCachedPrefixedRender(rendered string, width int, key uint64) {
221 c.prefixedRendered = rendered
222 c.prefixedWidth = width
223 c.prefixedKey = key
224}
225
226// clearCache clears the cached render.
227func (c *cachedMessageItem) clearCache() {
228 c.rendered = ""
229 c.width = 0
230 c.height = 0
231 c.prefixedRendered = ""
232 c.prefixedWidth = 0
233 c.prefixedKey = 0
234}
235
236// focusableMessageItem is a base struct for message items that can be focused.
237type focusableMessageItem struct {
238 // version is the parent item's version counter. SetFocused
239 // bumps it whenever focus actually flips so the F6 list memo
240 // invalidates the per-line focus prefix.
241 version *list.Versioned
242 focused bool
243}
244
245// newFocusableMessageItem returns a focusableMessageItem wired to the
246// shared version counter.
247func newFocusableMessageItem(v *list.Versioned) *focusableMessageItem {
248 return &focusableMessageItem{version: v}
249}
250
251// SetFocused implements MessageItem.
252func (f *focusableMessageItem) SetFocused(focused bool) {
253 if f.focused == focused {
254 return
255 }
256 f.focused = focused
257 if f.version != nil {
258 f.version.Bump()
259 }
260}
261
262// AssistantInfoID returns a stable ID for assistant info items.
263func AssistantInfoID(messageID string) string {
264 return fmt.Sprintf("%s:assistant-info", messageID)
265}
266
267// AssistantInfoItem renders model info and response time after assistant completes.
268type AssistantInfoItem struct {
269 *list.Versioned
270 *cachedMessageItem
271
272 id string
273 message *message.Message
274 sty *styles.Styles
275 cfg *config.Config
276 lastUserMessageTime time.Time
277}
278
279// NewAssistantInfoItem creates a new AssistantInfoItem.
280func NewAssistantInfoItem(sty *styles.Styles, message *message.Message, cfg *config.Config, lastUserMessageTime time.Time) MessageItem {
281 return &AssistantInfoItem{
282 Versioned: list.NewVersioned(),
283 cachedMessageItem: &cachedMessageItem{},
284 id: AssistantInfoID(message.ID),
285 message: message,
286 sty: sty,
287 cfg: cfg,
288 lastUserMessageTime: lastUserMessageTime,
289 }
290}
291
292// Finished implements list.Item. Assistant info blocks render a fixed
293// model/duration footer once the assistant turn finishes; the data
294// is immutable after construction so the entry is safe to freeze.
295func (a *AssistantInfoItem) Finished() bool {
296 return true
297}
298
299// ID implements MessageItem.
300func (a *AssistantInfoItem) ID() string {
301 return a.id
302}
303
304// RawRender implements MessageItem.
305func (a *AssistantInfoItem) RawRender(width int) string {
306 innerWidth := max(0, width-MessageLeftPaddingTotal)
307 content, _, ok := a.getCachedRender(innerWidth)
308 if !ok {
309 content = a.renderContent(innerWidth)
310 height := lipgloss.Height(content)
311 a.setCachedRender(content, innerWidth, height)
312 }
313 return content
314}
315
316// Render implements MessageItem.
317func (a *AssistantInfoItem) Render(width int) string {
318 // AssistantInfoItem uses a single, state-independent prefix; key 0
319 // is sufficient. The cache is invalidated whenever the underlying
320 // cachedMessageItem render is cleared.
321 if cached, ok := a.getCachedPrefixedRender(width, 0); ok {
322 return cached
323 }
324 prefix := a.sty.Messages.SectionHeader.Render()
325 lines := strings.Split(a.RawRender(width), "\n")
326 for i, line := range lines {
327 lines[i] = prefix + line
328 }
329 out := strings.Join(lines, "\n")
330 a.setCachedPrefixedRender(out, width, 0)
331 return out
332}
333
334func (a *AssistantInfoItem) renderContent(width int) string {
335 finishData := a.message.FinishPart()
336 if finishData == nil {
337 return ""
338 }
339 finishTime := time.Unix(finishData.Time, 0)
340 duration := finishTime.Sub(a.lastUserMessageTime)
341 infoMsg := a.sty.Messages.AssistantInfoDuration.Render(duration.String())
342 icon := a.sty.Messages.AssistantInfoIcon.Render(styles.ModelIcon)
343 model := a.cfg.GetModel(a.message.Provider, a.message.Model)
344 if model == nil {
345 model = &catwalk.Model{Name: "Unknown Model"}
346 }
347 modelFormatted := a.sty.Messages.AssistantInfoModel.Render(model.Name)
348 providerName := a.message.Provider
349 if providerConfig, ok := a.cfg.Providers.Get(a.message.Provider); ok {
350 providerName = providerConfig.Name
351 }
352 provider := a.sty.Messages.AssistantInfoProvider.Render(fmt.Sprintf("via %s", providerName))
353 assistant := fmt.Sprintf("%s %s %s %s", icon, modelFormatted, provider, infoMsg)
354 return common.Section(a.sty, assistant, width)
355}
356
357// cappedMessageWidth returns the maximum width for message content for readability.
358func cappedMessageWidth(availableWidth int) int {
359 return min(availableWidth-MessageLeftPaddingTotal, maxTextWidth)
360}
361
362// ExtractMessageItems extracts [MessageItem]s from a [message.Message]. It
363// returns all parts of the message as [MessageItem]s.
364//
365// For assistant messages with tool calls, pass a toolResults map to link results.
366// Use BuildToolResultMap to create this map from all messages in a session.
367func ExtractMessageItems(sty *styles.Styles, msg *message.Message, toolResults map[string]message.ToolResult) []MessageItem {
368 switch msg.Role {
369 case message.User:
370 r := attachments.NewRenderer(
371 sty.Attachments.Normal,
372 sty.Attachments.Deleting,
373 sty.Attachments.Image,
374 sty.Attachments.Text,
375 sty.Attachments.Skill,
376 )
377 return []MessageItem{NewUserMessageItem(sty, msg, r)}
378 case message.Assistant:
379 var items []MessageItem
380 if ShouldRenderAssistantMessage(msg) {
381 items = append(items, NewAssistantMessageItem(sty, msg))
382 }
383 for _, tc := range msg.ToolCalls() {
384 var result *message.ToolResult
385 if tr, ok := toolResults[tc.ID]; ok {
386 result = &tr
387 }
388 items = append(items, NewToolMessageItem(
389 sty,
390 msg.ID,
391 tc,
392 result,
393 msg.FinishReason() == message.FinishReasonCanceled,
394 ))
395 }
396 return items
397 }
398 return []MessageItem{}
399}
400
401// ShouldRenderAssistantMessage determines if an assistant message should be rendered
402//
403// In some cases the assistant message only has tools so we do not want to render an
404// empty message.
405func ShouldRenderAssistantMessage(msg *message.Message) bool {
406 content := strings.TrimSpace(msg.Content().Text)
407 thinking := strings.TrimSpace(msg.ReasoningContent().Thinking)
408 isError := msg.FinishReason() == message.FinishReasonError
409 isCancelled := msg.FinishReason() == message.FinishReasonCanceled
410 hasToolCalls := len(msg.ToolCalls()) > 0
411 return !hasToolCalls || content != "" || thinking != "" || msg.IsThinking() || isError || isCancelled
412}
413
414// BuildToolResultMap creates a map of tool call IDs to their results from a list of messages.
415// Tool result messages (role == message.Tool) contain the results that should be linked
416// to tool calls in assistant messages.
417func BuildToolResultMap(messages []*message.Message) map[string]message.ToolResult {
418 resultMap := make(map[string]message.ToolResult)
419 for _, msg := range messages {
420 if msg.Role == message.Tool {
421 for _, result := range msg.ToolResults() {
422 if result.ToolCallID != "" {
423 resultMap[result.ToolCallID] = result
424 }
425 }
426 }
427 }
428 return resultMap
429}