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