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