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).
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 if shouldShowMessage {
308 m.listCmp.UpdateItem(
309 assistantIndex,
310 messages.NewMessageCmp(
311 msg,
312 ),
313 )
314
315 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {
316 m.listCmp.AppendItem(
317 messages.NewAssistantSection(
318 msg,
319 time.Unix(m.lastUserMessageTime, 0),
320 ),
321 )
322 }
323 } else if hasToolCallsOnly {
324 m.listCmp.DeleteItem(assistantIndex)
325 }
326
327 return nil
328}
329
330// shouldShowAssistantMessage determines if an assistant message should be displayed.
331func (m *messageListCmp) shouldShowAssistantMessage(msg message.Message) bool {
332 return len(msg.ToolCalls()) == 0 || msg.Content().Text != "" || msg.IsThinking()
333}
334
335// updateToolCalls handles updates to tool calls, updating existing ones and adding new ones.
336func (m *messageListCmp) updateToolCalls(msg message.Message, existingToolCalls map[int]messages.ToolCallCmp) tea.Cmd {
337 var cmds []tea.Cmd
338
339 for _, tc := range msg.ToolCalls() {
340 if cmd := m.updateOrAddToolCall(msg, tc, existingToolCalls); cmd != nil {
341 cmds = append(cmds, cmd)
342 }
343 }
344
345 return tea.Batch(cmds...)
346}
347
348// updateOrAddToolCall updates an existing tool call or adds a new one.
349func (m *messageListCmp) updateOrAddToolCall(msg message.Message, tc message.ToolCall, existingToolCalls map[int]messages.ToolCallCmp) tea.Cmd {
350 // Try to find existing tool call
351 for index, existingTC := range existingToolCalls {
352 if tc.ID == existingTC.GetToolCall().ID {
353 existingTC.SetToolCall(tc)
354 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonCanceled {
355 existingTC.SetCancelled()
356 }
357 m.listCmp.UpdateItem(index, existingTC)
358 return nil
359 }
360 }
361
362 // Add new tool call if not found
363 return m.listCmp.AppendItem(messages.NewToolCallCmp(msg.ID, tc))
364}
365
366// handleNewAssistantMessage processes new assistant messages and their tool calls.
367func (m *messageListCmp) handleNewAssistantMessage(msg message.Message) tea.Cmd {
368 var cmds []tea.Cmd
369
370 // Add assistant message if it should be displayed
371 if m.shouldShowAssistantMessage(msg) {
372 cmd := m.listCmp.AppendItem(
373 messages.NewMessageCmp(
374 msg,
375 ),
376 )
377 cmds = append(cmds, cmd)
378 }
379
380 // Add tool calls
381 for _, tc := range msg.ToolCalls() {
382 cmd := m.listCmp.AppendItem(messages.NewToolCallCmp(msg.ID, tc))
383 cmds = append(cmds, cmd)
384 }
385
386 return tea.Batch(cmds...)
387}
388
389// SetSession loads and displays messages for a new session.
390func (m *messageListCmp) SetSession(session session.Session) tea.Cmd {
391 if m.session.ID == session.ID {
392 return nil
393 }
394
395 m.session = session
396 sessionMessages, err := m.app.Messages.List(context.Background(), session.ID)
397 if err != nil {
398 return util.ReportError(err)
399 }
400
401 if len(sessionMessages) == 0 {
402 return m.listCmp.SetItems([]util.Model{})
403 }
404
405 // Initialize with first message timestamp
406 m.lastUserMessageTime = sessionMessages[0].CreatedAt
407
408 // Build tool result map for efficient lookup
409 toolResultMap := m.buildToolResultMap(sessionMessages)
410
411 // Convert messages to UI components
412 uiMessages := m.convertMessagesToUI(sessionMessages, toolResultMap)
413
414 return m.listCmp.SetItems(uiMessages)
415}
416
417// buildToolResultMap creates a map of tool call ID to tool result for efficient lookup.
418func (m *messageListCmp) buildToolResultMap(messages []message.Message) map[string]message.ToolResult {
419 toolResultMap := make(map[string]message.ToolResult)
420 for _, msg := range messages {
421 for _, tr := range msg.ToolResults() {
422 toolResultMap[tr.ToolCallID] = tr
423 }
424 }
425 return toolResultMap
426}
427
428// convertMessagesToUI converts database messages to UI components.
429func (m *messageListCmp) convertMessagesToUI(sessionMessages []message.Message, toolResultMap map[string]message.ToolResult) []util.Model {
430 uiMessages := make([]util.Model, 0)
431
432 for _, msg := range sessionMessages {
433 switch msg.Role {
434 case message.User:
435 m.lastUserMessageTime = msg.CreatedAt
436 uiMessages = append(uiMessages, messages.NewMessageCmp(msg))
437 case message.Assistant:
438 uiMessages = append(uiMessages, m.convertAssistantMessage(msg, toolResultMap)...)
439 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {
440 uiMessages = append(uiMessages, messages.NewAssistantSection(msg, time.Unix(m.lastUserMessageTime, 0)))
441 }
442 }
443 }
444
445 return uiMessages
446}
447
448// convertAssistantMessage converts an assistant message and its tool calls to UI components.
449func (m *messageListCmp) convertAssistantMessage(msg message.Message, toolResultMap map[string]message.ToolResult) []util.Model {
450 var uiMessages []util.Model
451
452 // Add assistant message if it should be displayed
453 if m.shouldShowAssistantMessage(msg) {
454 uiMessages = append(
455 uiMessages,
456 messages.NewMessageCmp(
457 msg,
458 ),
459 )
460 }
461
462 // Add tool calls with their results and status
463 for _, tc := range msg.ToolCalls() {
464 options := m.buildToolCallOptions(tc, msg, toolResultMap)
465 uiMessages = append(uiMessages, messages.NewToolCallCmp(msg.ID, tc, options...))
466 // If this tool call is the agent tool, fetch nested tool calls
467 if tc.Name == agent.AgentToolName {
468 nestedMessages, _ := m.app.Messages.List(context.Background(), tc.ID)
469 nestedUIMessages := m.convertMessagesToUI(nestedMessages, make(map[string]message.ToolResult))
470 nestedToolCalls := make([]messages.ToolCallCmp, 0, len(nestedUIMessages))
471 for _, nestedMsg := range nestedUIMessages {
472 if toolCall, ok := nestedMsg.(messages.ToolCallCmp); ok {
473 toolCall.SetIsNested(true)
474 nestedToolCalls = append(nestedToolCalls, toolCall)
475 }
476 }
477 uiMessages[len(uiMessages)-1].(messages.ToolCallCmp).SetNestedToolCalls(nestedToolCalls)
478 }
479 }
480
481 return uiMessages
482}
483
484// buildToolCallOptions creates options for tool call components based on results and status.
485func (m *messageListCmp) buildToolCallOptions(tc message.ToolCall, msg message.Message, toolResultMap map[string]message.ToolResult) []messages.ToolCallOption {
486 var options []messages.ToolCallOption
487
488 // Add tool result if available
489 if tr, ok := toolResultMap[tc.ID]; ok {
490 options = append(options, messages.WithToolCallResult(tr))
491 }
492
493 // Add cancelled status if applicable
494 if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonCanceled {
495 options = append(options, messages.WithToolCallCancelled())
496 }
497
498 return options
499}
500
501// GetSize returns the current width and height of the component.
502func (m *messageListCmp) GetSize() (int, int) {
503 return m.width, m.height
504}
505
506// SetSize updates the component dimensions and propagates to the list component.
507func (m *messageListCmp) SetSize(width int, height int) tea.Cmd {
508 m.width = width
509 m.height = height
510 return m.listCmp.SetSize(width-2, height-2) // for padding
511}
512
513// Blur implements MessageListCmp.
514func (m *messageListCmp) Blur() tea.Cmd {
515 return m.listCmp.Blur()
516}
517
518// Focus implements MessageListCmp.
519func (m *messageListCmp) Focus() tea.Cmd {
520 return m.listCmp.Focus()
521}
522
523// IsFocused implements MessageListCmp.
524func (m *messageListCmp) IsFocused() bool {
525 return m.listCmp.IsFocused()
526}
527
528func (m *messageListCmp) Bindings() []key.Binding {
529 return m.defaultListKeyMap.KeyBindings()
530}