chat.go

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