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