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