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