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