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