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