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