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