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