chat.go

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