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