1package chat
2
3import (
4 "context"
5 "time"
6
7 "github.com/charmbracelet/bubbles/v2/key"
8 tea "github.com/charmbracelet/bubbletea/v2"
9 "github.com/charmbracelet/crush/internal/app"
10 "github.com/charmbracelet/crush/internal/llm/agent"
11 "github.com/charmbracelet/crush/internal/message"
12 "github.com/charmbracelet/crush/internal/pubsub"
13 "github.com/charmbracelet/crush/internal/session"
14 "github.com/charmbracelet/crush/internal/tui/components/chat/messages"
15 "github.com/charmbracelet/crush/internal/tui/components/core/layout"
16 "github.com/charmbracelet/crush/internal/tui/exp/list"
17 "github.com/charmbracelet/crush/internal/tui/styles"
18 "github.com/charmbracelet/crush/internal/tui/util"
19)
20
21type SendMsg struct {
22 Text string
23 Attachments []message.Attachment
24}
25
26type SessionSelectedMsg = session.Session
27
28type SessionClearedMsg struct{}
29
30const (
31 NotFound = -1
32)
33
34// MessageListCmp represents a component that displays a list of chat messages
35// with support for real-time updates and session management.
36type MessageListCmp interface {
37 util.Model
38 layout.Sizeable
39 layout.Focusable
40 layout.Help
41
42 SetSession(session.Session) tea.Cmd
43}
44
45// messageListCmp implements MessageListCmp, providing a virtualized list
46// of chat messages with support for tool calls, real-time updates, and
47// session switching.
48type messageListCmp struct {
49 app *app.App
50 width, height int
51 session session.Session
52 listCmp list.List[list.Item]
53 previousSelected string // Last selected item index for restoring focus
54
55 lastUserMessageTime int64
56 defaultListKeyMap list.KeyMap
57}
58
59// New creates a new message list component with custom keybindings
60// and reverse ordering (newest messages at bottom).
61func New(app *app.App) MessageListCmp {
62 defaultListKeyMap := list.DefaultKeyMap()
63 listCmp := list.New(
64 []list.Item{},
65 list.WithGap(1),
66 list.WithDirectionBackward(),
67 list.WithKeyMap(defaultListKeyMap),
68 )
69 return &messageListCmp{
70 app: app,
71 listCmp: listCmp,
72 previousSelected: "",
73 defaultListKeyMap: defaultListKeyMap,
74 }
75}
76
77// Init initializes the component.
78func (m *messageListCmp) Init() tea.Cmd {
79 return tea.Sequence(m.listCmp.Init(), m.listCmp.Blur())
80}
81
82// Update handles incoming messages and updates the component state.
83func (m *messageListCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
84 switch msg := msg.(type) {
85 case SessionSelectedMsg:
86 if msg.ID != m.session.ID {
87 cmd := m.SetSession(msg)
88 return m, cmd
89 }
90 return m, nil
91 case SessionClearedMsg:
92 m.session = session.Session{}
93 return m, m.listCmp.SetItems([]list.Item{})
94
95 case pubsub.Event[message.Message]:
96 cmd := m.handleMessageEvent(msg)
97 return m, cmd
98 default:
99 var cmds []tea.Cmd
100 u, cmd := m.listCmp.Update(msg)
101 m.listCmp = u.(list.List[list.Item])
102 cmds = append(cmds, cmd)
103 return m, tea.Batch(cmds...)
104 }
105}
106
107// View renders the message list or an initial screen if empty.
108func (m *messageListCmp) View() string {
109 t := styles.CurrentTheme()
110 return t.S().Base.
111 Padding(1, 1, 0, 1).
112 Width(m.width).
113 Height(m.height).
114 Render(
115 m.listCmp.View(),
116 )
117}
118
119// handleChildSession handles messages from child sessions (agent tools).
120func (m *messageListCmp) handleChildSession(event pubsub.Event[message.Message]) tea.Cmd {
121 var cmds []tea.Cmd
122 if len(event.Payload.ToolCalls()) == 0 && len(event.Payload.ToolResults()) == 0 {
123 return nil
124 }
125 items := m.listCmp.Items()
126 toolCallInx := NotFound
127 var toolCall messages.ToolCallCmp
128 for i := len(items) - 1; i >= 0; i-- {
129 if msg, ok := items[i].(messages.ToolCallCmp); ok {
130 if msg.GetToolCall().ID == event.Payload.SessionID {
131 toolCallInx = i
132 toolCall = msg
133 }
134 }
135 }
136 if toolCallInx == NotFound {
137 return nil
138 }
139 nestedToolCalls := toolCall.GetNestedToolCalls()
140 for _, tc := range event.Payload.ToolCalls() {
141 found := false
142 for existingInx, existingTC := range nestedToolCalls {
143 if existingTC.GetToolCall().ID == tc.ID {
144 nestedToolCalls[existingInx].SetToolCall(tc)
145 found = true
146 break
147 }
148 }
149 if !found {
150 nestedCall := messages.NewToolCallCmp(
151 event.Payload.ID,
152 tc,
153 messages.WithToolCallNested(true),
154 )
155 cmds = append(cmds, nestedCall.Init())
156 nestedToolCalls = append(
157 nestedToolCalls,
158 nestedCall,
159 )
160 }
161 }
162 for _, tr := range event.Payload.ToolResults() {
163 for nestedInx, nestedTC := range nestedToolCalls {
164 if nestedTC.GetToolCall().ID == tr.ToolCallID {
165 nestedToolCalls[nestedInx].SetToolResult(tr)
166 break
167 }
168 }
169 }
170
171 toolCall.SetNestedToolCalls(nestedToolCalls)
172 m.listCmp.UpdateItem(
173 toolCall.ID(),
174 toolCall,
175 )
176 return tea.Batch(cmds...)
177}
178
179// handleMessageEvent processes different types of message events (created/updated).
180func (m *messageListCmp) handleMessageEvent(event pubsub.Event[message.Message]) tea.Cmd {
181 switch event.Type {
182 case pubsub.CreatedEvent:
183 if event.Payload.SessionID != m.session.ID {
184 return m.handleChildSession(event)
185 }
186 if m.messageExists(event.Payload.ID) {
187 return nil
188 }
189 return m.handleNewMessage(event.Payload)
190 case pubsub.UpdatedEvent:
191 if event.Payload.SessionID != m.session.ID {
192 return m.handleChildSession(event)
193 }
194 return m.handleUpdateAssistantMessage(event.Payload)
195 }
196 return nil
197}
198
199// messageExists checks if a message with the given ID already exists in the list.
200func (m *messageListCmp) messageExists(messageID string) bool {
201 items := m.listCmp.Items()
202 // Search backwards as new messages are more likely to be at the end
203 for i := len(items) - 1; i >= 0; i-- {
204 if msg, ok := items[i].(messages.MessageCmp); ok && msg.GetMessage().ID == messageID {
205 return true
206 }
207 }
208 return false
209}
210
211// handleNewMessage routes new messages to appropriate handlers based on role.
212func (m *messageListCmp) handleNewMessage(msg message.Message) tea.Cmd {
213 switch msg.Role {
214 case message.User:
215 return m.handleNewUserMessage(msg)
216 case message.Assistant:
217 return m.handleNewAssistantMessage(msg)
218 case message.Tool:
219 return m.handleToolMessage(msg)
220 }
221 return nil
222}
223
224// handleNewUserMessage adds a new user message to the list and updates the timestamp.
225func (m *messageListCmp) handleNewUserMessage(msg message.Message) tea.Cmd {
226 m.lastUserMessageTime = msg.CreatedAt
227 return m.listCmp.AppendItem(messages.NewMessageCmp(msg))
228}
229
230// handleToolMessage updates existing tool calls with their results.
231func (m *messageListCmp) handleToolMessage(msg message.Message) tea.Cmd {
232 items := m.listCmp.Items()
233 for _, tr := range msg.ToolResults() {
234 if toolCallIndex := m.findToolCallByID(items, tr.ToolCallID); toolCallIndex != NotFound {
235 toolCall := items[toolCallIndex].(messages.ToolCallCmp)
236 toolCall.SetToolResult(tr)
237 m.listCmp.UpdateItem(toolCall.ID(), toolCall)
238 }
239 }
240 return nil
241}
242
243// findToolCallByID searches for a tool call with the specified ID.
244// Returns the index if found, NotFound otherwise.
245func (m *messageListCmp) findToolCallByID(items []list.Item, toolCallID string) int {
246 // Search backwards as tool calls are more likely to be recent
247 for i := len(items) - 1; i >= 0; i-- {
248 if toolCall, ok := items[i].(messages.ToolCallCmp); ok && toolCall.GetToolCall().ID == toolCallID {
249 return i
250 }
251 }
252 return NotFound
253}
254
255// handleUpdateAssistantMessage processes updates to assistant messages,
256// managing both message content and associated tool calls.
257func (m *messageListCmp) handleUpdateAssistantMessage(msg message.Message) tea.Cmd {
258 var cmds []tea.Cmd
259 items := m.listCmp.Items()
260
261 // Find existing assistant message and tool calls for this message
262 assistantIndex, existingToolCalls := m.findAssistantMessageAndToolCalls(items, msg.ID)
263
264 // Handle assistant message content
265 if cmd := m.updateAssistantMessageContent(msg, assistantIndex); cmd != nil {
266 cmds = append(cmds, cmd)
267 }
268
269 // Handle tool calls
270 if cmd := m.updateToolCalls(msg, existingToolCalls); cmd != nil {
271 cmds = append(cmds, cmd)
272 }
273
274 return tea.Batch(cmds...)
275}
276
277// findAssistantMessageAndToolCalls locates the assistant message and its tool calls.
278func (m *messageListCmp) findAssistantMessageAndToolCalls(items []list.Item, messageID string) (int, map[int]messages.ToolCallCmp) {
279 assistantIndex := NotFound
280 toolCalls := make(map[int]messages.ToolCallCmp)
281
282 // Search backwards as messages are more likely to be at the end
283 for i := len(items) - 1; i >= 0; i-- {
284 item := items[i]
285 if asMsg, ok := item.(messages.MessageCmp); ok {
286 if asMsg.GetMessage().ID == messageID {
287 assistantIndex = i
288 }
289 } else if tc, ok := item.(messages.ToolCallCmp); ok {
290 if tc.ParentMessageID() == messageID {
291 toolCalls[i] = tc
292 }
293 }
294 }
295
296 return assistantIndex, toolCalls
297}
298
299// updateAssistantMessageContent updates or removes the assistant message based on content.
300func (m *messageListCmp) updateAssistantMessageContent(msg message.Message, assistantIndex int) tea.Cmd {
301 if assistantIndex == NotFound {
302 return nil
303 }
304
305 shouldShowMessage := m.shouldShowAssistantMessage(msg)
306 hasToolCallsOnly := len(msg.ToolCalls()) > 0 && msg.Content().Text == ""
307
308 var cmd tea.Cmd
309 if shouldShowMessage {
310 items := m.listCmp.Items()
311 uiMsg := items[assistantIndex].(messages.MessageCmp)
312 uiMsg.SetMessage(msg)
313 m.listCmp.UpdateItem(
314 items[assistantIndex].ID(),
315 uiMsg,
316 )
317 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {
318 m.listCmp.AppendItem(
319 messages.NewAssistantSection(
320 msg,
321 time.Unix(m.lastUserMessageTime, 0),
322 ),
323 )
324 }
325 } else if hasToolCallsOnly {
326 items := m.listCmp.Items()
327 m.listCmp.DeleteItem(items[assistantIndex].ID())
328 }
329
330 return cmd
331}
332
333// shouldShowAssistantMessage determines if an assistant message should be displayed.
334func (m *messageListCmp) shouldShowAssistantMessage(msg message.Message) bool {
335 return len(msg.ToolCalls()) == 0 || msg.Content().Text != "" || msg.ReasoningContent().Thinking != "" || msg.IsThinking()
336}
337
338// updateToolCalls handles updates to tool calls, updating existing ones and adding new ones.
339func (m *messageListCmp) updateToolCalls(msg message.Message, existingToolCalls map[int]messages.ToolCallCmp) tea.Cmd {
340 var cmds []tea.Cmd
341
342 for _, tc := range msg.ToolCalls() {
343 if cmd := m.updateOrAddToolCall(msg, tc, existingToolCalls); cmd != nil {
344 cmds = append(cmds, cmd)
345 }
346 }
347
348 return tea.Batch(cmds...)
349}
350
351// updateOrAddToolCall updates an existing tool call or adds a new one.
352func (m *messageListCmp) updateOrAddToolCall(msg message.Message, tc message.ToolCall, existingToolCalls map[int]messages.ToolCallCmp) tea.Cmd {
353 // Try to find existing tool call
354 for _, existingTC := range existingToolCalls {
355 if tc.ID == existingTC.GetToolCall().ID {
356 existingTC.SetToolCall(tc)
357 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonCanceled {
358 existingTC.SetCancelled()
359 }
360 m.listCmp.UpdateItem(tc.ID, existingTC)
361 return nil
362 }
363 }
364
365 // Add new tool call if not found
366 return m.listCmp.AppendItem(messages.NewToolCallCmp(msg.ID, tc))
367}
368
369// handleNewAssistantMessage processes new assistant messages and their tool calls.
370func (m *messageListCmp) handleNewAssistantMessage(msg message.Message) tea.Cmd {
371 var cmds []tea.Cmd
372
373 // Add assistant message if it should be displayed
374 if m.shouldShowAssistantMessage(msg) {
375 cmd := m.listCmp.AppendItem(
376 messages.NewMessageCmp(
377 msg,
378 ),
379 )
380 cmds = append(cmds, cmd)
381 }
382
383 // Add tool calls
384 for _, tc := range msg.ToolCalls() {
385 cmd := m.listCmp.AppendItem(messages.NewToolCallCmp(msg.ID, tc))
386 cmds = append(cmds, cmd)
387 }
388
389 return tea.Batch(cmds...)
390}
391
392// SetSession loads and displays messages for a new session.
393func (m *messageListCmp) SetSession(session session.Session) tea.Cmd {
394 if m.session.ID == session.ID {
395 return nil
396 }
397
398 m.session = session
399 sessionMessages, err := m.app.Messages.List(context.Background(), session.ID)
400 if err != nil {
401 return util.ReportError(err)
402 }
403
404 if len(sessionMessages) == 0 {
405 return m.listCmp.SetItems([]list.Item{})
406 }
407
408 // Initialize with first message timestamp
409 m.lastUserMessageTime = sessionMessages[0].CreatedAt
410
411 // Build tool result map for efficient lookup
412 toolResultMap := m.buildToolResultMap(sessionMessages)
413
414 // Convert messages to UI components
415 uiMessages := m.convertMessagesToUI(sessionMessages, toolResultMap)
416
417 return m.listCmp.SetItems(uiMessages)
418}
419
420// buildToolResultMap creates a map of tool call ID to tool result for efficient lookup.
421func (m *messageListCmp) buildToolResultMap(messages []message.Message) map[string]message.ToolResult {
422 toolResultMap := make(map[string]message.ToolResult)
423 for _, msg := range messages {
424 for _, tr := range msg.ToolResults() {
425 toolResultMap[tr.ToolCallID] = tr
426 }
427 }
428 return toolResultMap
429}
430
431// convertMessagesToUI converts database messages to UI components.
432func (m *messageListCmp) convertMessagesToUI(sessionMessages []message.Message, toolResultMap map[string]message.ToolResult) []list.Item {
433 uiMessages := make([]list.Item, 0)
434
435 for _, msg := range sessionMessages {
436 switch msg.Role {
437 case message.User:
438 m.lastUserMessageTime = msg.CreatedAt
439 uiMessages = append(uiMessages, messages.NewMessageCmp(msg))
440 case message.Assistant:
441 uiMessages = append(uiMessages, m.convertAssistantMessage(msg, toolResultMap)...)
442 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {
443 uiMessages = append(uiMessages, messages.NewAssistantSection(msg, time.Unix(m.lastUserMessageTime, 0)))
444 }
445 }
446 }
447
448 return uiMessages
449}
450
451// convertAssistantMessage converts an assistant message and its tool calls to UI components.
452func (m *messageListCmp) convertAssistantMessage(msg message.Message, toolResultMap map[string]message.ToolResult) []list.Item {
453 var uiMessages []list.Item
454
455 // Add assistant message if it should be displayed
456 if m.shouldShowAssistantMessage(msg) {
457 uiMessages = append(
458 uiMessages,
459 messages.NewMessageCmp(
460 msg,
461 ),
462 )
463 }
464
465 // Add tool calls with their results and status
466 for _, tc := range msg.ToolCalls() {
467 options := m.buildToolCallOptions(tc, msg, toolResultMap)
468 uiMessages = append(uiMessages, messages.NewToolCallCmp(msg.ID, tc, options...))
469 // If this tool call is the agent tool, fetch nested tool calls
470 if tc.Name == agent.AgentToolName {
471 nestedMessages, _ := m.app.Messages.List(context.Background(), tc.ID)
472 nestedUIMessages := m.convertMessagesToUI(nestedMessages, make(map[string]message.ToolResult))
473 nestedToolCalls := make([]messages.ToolCallCmp, 0, len(nestedUIMessages))
474 for _, nestedMsg := range nestedUIMessages {
475 if toolCall, ok := nestedMsg.(messages.ToolCallCmp); ok {
476 toolCall.SetIsNested(true)
477 nestedToolCalls = append(nestedToolCalls, toolCall)
478 }
479 }
480 uiMessages[len(uiMessages)-1].(messages.ToolCallCmp).SetNestedToolCalls(nestedToolCalls)
481 }
482 }
483
484 return uiMessages
485}
486
487// buildToolCallOptions creates options for tool call components based on results and status.
488func (m *messageListCmp) buildToolCallOptions(tc message.ToolCall, msg message.Message, toolResultMap map[string]message.ToolResult) []messages.ToolCallOption {
489 var options []messages.ToolCallOption
490
491 // Add tool result if available
492 if tr, ok := toolResultMap[tc.ID]; ok {
493 options = append(options, messages.WithToolCallResult(tr))
494 }
495
496 // Add cancelled status if applicable
497 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonCanceled {
498 options = append(options, messages.WithToolCallCancelled())
499 }
500
501 return options
502}
503
504// GetSize returns the current width and height of the component.
505func (m *messageListCmp) GetSize() (int, int) {
506 return m.width, m.height
507}
508
509// SetSize updates the component dimensions and propagates to the list component.
510func (m *messageListCmp) SetSize(width int, height int) tea.Cmd {
511 m.width = width
512 m.height = height
513 return m.listCmp.SetSize(width-2, height-1) // for padding
514}
515
516// Blur implements MessageListCmp.
517func (m *messageListCmp) Blur() tea.Cmd {
518 return m.listCmp.Blur()
519}
520
521// Focus implements MessageListCmp.
522func (m *messageListCmp) Focus() tea.Cmd {
523 return m.listCmp.Focus()
524}
525
526// IsFocused implements MessageListCmp.
527func (m *messageListCmp) IsFocused() bool {
528 return m.listCmp.IsFocused()
529}
530
531func (m *messageListCmp) Bindings() []key.Binding {
532 return m.defaultListKeyMap.KeyBindings()
533}