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