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 startLine int
79 startCol int
80 endLine int
81 endCol int
82 highlighter list.Highlighter
83}
84
85var _ list.Highlightable = (*highlightableMessageItem)(nil)
86
87// isHighlighted returns true if the item has a highlight range set.
88func (h *highlightableMessageItem) isHighlighted() bool {
89 return h.startLine != -1 || h.endLine != -1
90}
91
92// renderHighlighted highlights the content if necessary.
93func (h *highlightableMessageItem) renderHighlighted(content string, width, height int) string {
94 if !h.isHighlighted() {
95 return content
96 }
97 area := image.Rect(0, 0, width, height)
98 return list.Highlight(content, area, h.startLine, h.startCol, h.endLine, h.endCol, h.highlighter)
99}
100
101// SetHighlight implements list.Highlightable.
102func (h *highlightableMessageItem) SetHighlight(startLine int, startCol int, endLine int, endCol int) {
103 // Adjust columns for the style's left inset (border + padding) since we
104 // highlight the content only.
105 offset := MessageLeftPaddingTotal
106 h.startLine = startLine
107 h.startCol = max(0, startCol-offset)
108 h.endLine = endLine
109 if endCol >= 0 {
110 h.endCol = max(0, endCol-offset)
111 } else {
112 h.endCol = endCol
113 }
114}
115
116// Highlight implements list.Highlightable.
117func (h *highlightableMessageItem) Highlight() (startLine int, startCol int, endLine int, endCol int) {
118 return h.startLine, h.startCol, h.endLine, h.endCol
119}
120
121func defaultHighlighter(sty *styles.Styles) *highlightableMessageItem {
122 return &highlightableMessageItem{
123 startLine: -1,
124 startCol: -1,
125 endLine: -1,
126 endCol: -1,
127 highlighter: list.ToHighlighter(sty.TextSelection),
128 }
129}
130
131// cacheClearable is implemented by message items that cache rendered
132// output and can be asked to drop the cache.
133type cacheClearable interface {
134 clearCache()
135}
136
137// ClearItemCaches drops any cached rendered output on each item so the
138// next render uses the current styles.
139func ClearItemCaches(items []MessageItem) {
140 for _, item := range items {
141 if cc, ok := item.(cacheClearable); ok {
142 cc.clearCache()
143 }
144 }
145}
146
147// cachedMessageItem caches rendered message content to avoid re-rendering.
148//
149// This should be used by any message that can store a cached version of its render. e.x user,assistant... and so on
150//
151// THOUGHT(kujtim): we should consider if its efficient to store the render for different widths
152// the issue with that could be memory usage
153type cachedMessageItem struct {
154 // rendered is the cached rendered string
155 rendered string
156 // width and height are the dimensions of the cached render
157 width int
158 height int
159
160 // prefixedRendered caches the per-line-prefixed Render output (the
161 // result of splitting RawRender by newlines and prepending a focus
162 // or selection prefix to every line). Items rebuild this every
163 // frame today; caching it keyed by (prefixedWidth, prefixedKey)
164 // turns Render into a pointer return when item state is stable.
165 //
166 // Invalidation lives in clearCache; callers must additionally
167 // bypass this cache whenever the prefixed output would not be
168 // stable (spinner ticks, active highlight ranges) by not calling
169 // setCachedPrefixedRender for those frames.
170 prefixedRendered string
171 prefixedWidth int
172 prefixedKey uint64
173}
174
175// getCachedRender returns the cached render if it exists for the given width.
176func (c *cachedMessageItem) getCachedRender(width int) (string, int, bool) {
177 if c.width == width && c.rendered != "" {
178 return c.rendered, c.height, true
179 }
180 return "", 0, false
181}
182
183// setCachedRender sets the cached render.
184func (c *cachedMessageItem) setCachedRender(rendered string, width, height int) {
185 c.rendered = rendered
186 c.width = width
187 c.height = height
188}
189
190// getCachedPrefixedRender returns the cached prefixed render if it exists
191// for the given (width, key). The key encodes any state that changes the
192// per-line prefix (focused/blurred, compact, ...).
193func (c *cachedMessageItem) getCachedPrefixedRender(width int, key uint64) (string, bool) {
194 if c.prefixedRendered != "" && c.prefixedWidth == width && c.prefixedKey == key {
195 return c.prefixedRendered, true
196 }
197 return "", false
198}
199
200// setCachedPrefixedRender stores the cached prefixed render.
201func (c *cachedMessageItem) setCachedPrefixedRender(rendered string, width int, key uint64) {
202 c.prefixedRendered = rendered
203 c.prefixedWidth = width
204 c.prefixedKey = key
205}
206
207// clearCache clears the cached render.
208func (c *cachedMessageItem) clearCache() {
209 c.rendered = ""
210 c.width = 0
211 c.height = 0
212 c.prefixedRendered = ""
213 c.prefixedWidth = 0
214 c.prefixedKey = 0
215}
216
217// focusableMessageItem is a base struct for message items that can be focused.
218type focusableMessageItem struct {
219 focused bool
220}
221
222// SetFocused implements MessageItem.
223func (f *focusableMessageItem) SetFocused(focused bool) {
224 f.focused = focused
225}
226
227// AssistantInfoID returns a stable ID for assistant info items.
228func AssistantInfoID(messageID string) string {
229 return fmt.Sprintf("%s:assistant-info", messageID)
230}
231
232// AssistantInfoItem renders model info and response time after assistant completes.
233type AssistantInfoItem struct {
234 *cachedMessageItem
235
236 id string
237 message *message.Message
238 sty *styles.Styles
239 cfg *config.Config
240 lastUserMessageTime time.Time
241}
242
243// NewAssistantInfoItem creates a new AssistantInfoItem.
244func NewAssistantInfoItem(sty *styles.Styles, message *message.Message, cfg *config.Config, lastUserMessageTime time.Time) MessageItem {
245 return &AssistantInfoItem{
246 cachedMessageItem: &cachedMessageItem{},
247 id: AssistantInfoID(message.ID),
248 message: message,
249 sty: sty,
250 cfg: cfg,
251 lastUserMessageTime: lastUserMessageTime,
252 }
253}
254
255// ID implements MessageItem.
256func (a *AssistantInfoItem) ID() string {
257 return a.id
258}
259
260// RawRender implements MessageItem.
261func (a *AssistantInfoItem) RawRender(width int) string {
262 innerWidth := max(0, width-MessageLeftPaddingTotal)
263 content, _, ok := a.getCachedRender(innerWidth)
264 if !ok {
265 content = a.renderContent(innerWidth)
266 height := lipgloss.Height(content)
267 a.setCachedRender(content, innerWidth, height)
268 }
269 return content
270}
271
272// Render implements MessageItem.
273func (a *AssistantInfoItem) Render(width int) string {
274 // AssistantInfoItem uses a single, state-independent prefix; key 0
275 // is sufficient. The cache is invalidated whenever the underlying
276 // cachedMessageItem render is cleared.
277 if cached, ok := a.getCachedPrefixedRender(width, 0); ok {
278 return cached
279 }
280 prefix := a.sty.Messages.SectionHeader.Render()
281 lines := strings.Split(a.RawRender(width), "\n")
282 for i, line := range lines {
283 lines[i] = prefix + line
284 }
285 out := strings.Join(lines, "\n")
286 a.setCachedPrefixedRender(out, width, 0)
287 return out
288}
289
290func (a *AssistantInfoItem) renderContent(width int) string {
291 finishData := a.message.FinishPart()
292 if finishData == nil {
293 return ""
294 }
295 finishTime := time.Unix(finishData.Time, 0)
296 duration := finishTime.Sub(a.lastUserMessageTime)
297 infoMsg := a.sty.Messages.AssistantInfoDuration.Render(duration.String())
298 icon := a.sty.Messages.AssistantInfoIcon.Render(styles.ModelIcon)
299 model := a.cfg.GetModel(a.message.Provider, a.message.Model)
300 if model == nil {
301 model = &catwalk.Model{Name: "Unknown Model"}
302 }
303 modelFormatted := a.sty.Messages.AssistantInfoModel.Render(model.Name)
304 providerName := a.message.Provider
305 if providerConfig, ok := a.cfg.Providers.Get(a.message.Provider); ok {
306 providerName = providerConfig.Name
307 }
308 provider := a.sty.Messages.AssistantInfoProvider.Render(fmt.Sprintf("via %s", providerName))
309 assistant := fmt.Sprintf("%s %s %s %s", icon, modelFormatted, provider, infoMsg)
310 return common.Section(a.sty, assistant, width)
311}
312
313// cappedMessageWidth returns the maximum width for message content for readability.
314func cappedMessageWidth(availableWidth int) int {
315 return min(availableWidth-MessageLeftPaddingTotal, maxTextWidth)
316}
317
318// ExtractMessageItems extracts [MessageItem]s from a [message.Message]. It
319// returns all parts of the message as [MessageItem]s.
320//
321// For assistant messages with tool calls, pass a toolResults map to link results.
322// Use BuildToolResultMap to create this map from all messages in a session.
323func ExtractMessageItems(sty *styles.Styles, msg *message.Message, toolResults map[string]message.ToolResult) []MessageItem {
324 switch msg.Role {
325 case message.User:
326 r := attachments.NewRenderer(
327 sty.Attachments.Normal,
328 sty.Attachments.Deleting,
329 sty.Attachments.Image,
330 sty.Attachments.Text,
331 )
332 return []MessageItem{NewUserMessageItem(sty, msg, r)}
333 case message.Assistant:
334 var items []MessageItem
335 if ShouldRenderAssistantMessage(msg) {
336 items = append(items, NewAssistantMessageItem(sty, msg))
337 }
338 for _, tc := range msg.ToolCalls() {
339 var result *message.ToolResult
340 if tr, ok := toolResults[tc.ID]; ok {
341 result = &tr
342 }
343 items = append(items, NewToolMessageItem(
344 sty,
345 msg.ID,
346 tc,
347 result,
348 msg.FinishReason() == message.FinishReasonCanceled,
349 ))
350 }
351 return items
352 }
353 return []MessageItem{}
354}
355
356// ShouldRenderAssistantMessage determines if an assistant message should be rendered
357//
358// In some cases the assistant message only has tools so we do not want to render an
359// empty message.
360func ShouldRenderAssistantMessage(msg *message.Message) bool {
361 content := strings.TrimSpace(msg.Content().Text)
362 thinking := strings.TrimSpace(msg.ReasoningContent().Thinking)
363 isError := msg.FinishReason() == message.FinishReasonError
364 isCancelled := msg.FinishReason() == message.FinishReasonCanceled
365 hasToolCalls := len(msg.ToolCalls()) > 0
366 return !hasToolCalls || content != "" || thinking != "" || msg.IsThinking() || isError || isCancelled
367}
368
369// BuildToolResultMap creates a map of tool call IDs to their results from a list of messages.
370// Tool result messages (role == message.Tool) contain the results that should be linked
371// to tool calls in assistant messages.
372func BuildToolResultMap(messages []*message.Message) map[string]message.ToolResult {
373 resultMap := make(map[string]message.ToolResult)
374 for _, msg := range messages {
375 if msg.Role == message.Tool {
376 for _, result := range msg.ToolResults() {
377 if result.ToolCallID != "" {
378 resultMap[result.ToolCallID] = result
379 }
380 }
381 }
382 }
383 return resultMap
384}