chat.go

  1package chat
  2
  3import (
  4	"context"
  5
  6	"github.com/charmbracelet/bubbles/v2/key"
  7	tea "github.com/charmbracelet/bubbletea/v2"
  8	"github.com/charmbracelet/crush/internal/app"
  9	"github.com/charmbracelet/crush/internal/config"
 10	"github.com/charmbracelet/crush/internal/llm/models"
 11	"github.com/charmbracelet/crush/internal/message"
 12	"github.com/charmbracelet/crush/internal/session"
 13	"github.com/charmbracelet/crush/internal/tui/components/chat"
 14	"github.com/charmbracelet/crush/internal/tui/components/chat/editor"
 15	"github.com/charmbracelet/crush/internal/tui/components/chat/sidebar"
 16	"github.com/charmbracelet/crush/internal/tui/components/core/layout"
 17	"github.com/charmbracelet/crush/internal/tui/components/dialogs/commands"
 18	"github.com/charmbracelet/crush/internal/tui/page"
 19	"github.com/charmbracelet/crush/internal/tui/util"
 20)
 21
 22var ChatPageID page.PageID = "chat"
 23
 24type (
 25	OpenFilePickerMsg struct{}
 26	ChatFocusedMsg    struct {
 27		Focused bool // True if the chat input is focused, false otherwise
 28	}
 29)
 30
 31type ChatPage interface {
 32	util.Model
 33	layout.Help
 34}
 35
 36type chatPage struct {
 37	app *app.App
 38
 39	layout layout.SplitPaneLayout
 40
 41	session session.Session
 42
 43	keyMap KeyMap
 44
 45	chatFocused bool
 46}
 47
 48func (p *chatPage) Init() tea.Cmd {
 49	cmd := p.layout.Init()
 50	return tea.Batch(
 51		cmd,
 52		p.layout.FocusPanel(layout.BottomPanel), // Focus on the bottom panel (editor)
 53	)
 54}
 55
 56func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 57	var cmds []tea.Cmd
 58	switch msg := msg.(type) {
 59	case tea.WindowSizeMsg:
 60		cmd := p.layout.SetSize(msg.Width, msg.Height)
 61		cmds = append(cmds, cmd)
 62	case chat.SendMsg:
 63		cmd := p.sendMessage(msg.Text, msg.Attachments)
 64		if cmd != nil {
 65			return p, cmd
 66		}
 67	case commands.CommandRunCustomMsg:
 68		// Check if the agent is busy before executing custom commands
 69		if p.app.CoderAgent.IsBusy() {
 70			return p, util.ReportWarn("Agent is busy, please wait before executing a command...")
 71		}
 72
 73		// Handle custom command execution
 74		cmd := p.sendMessage(msg.Content, nil)
 75		if cmd != nil {
 76			return p, cmd
 77		}
 78	case chat.SessionSelectedMsg:
 79		if p.session.ID == "" {
 80			cmd := p.setMessages()
 81			if cmd != nil {
 82				cmds = append(cmds, cmd)
 83			}
 84		}
 85		p.session = msg
 86	case tea.KeyPressMsg:
 87		switch {
 88		case key.Matches(msg, p.keyMap.NewSession):
 89			p.session = session.Session{}
 90			return p, tea.Batch(
 91				p.clearMessages(),
 92				util.CmdHandler(chat.SessionClearedMsg{}),
 93			)
 94
 95		case key.Matches(msg, p.keyMap.AddAttachment):
 96			cfg := config.Get()
 97			agentCfg := cfg.Agents[config.AgentCoder]
 98			selectedModelID := agentCfg.Model
 99			model := models.SupportedModels[selectedModelID]
100			if model.SupportsAttachments {
101				return p, util.CmdHandler(OpenFilePickerMsg{})
102			} else {
103				return p, util.ReportWarn("File attachments are not supported by the current model: " + string(selectedModelID))
104			}
105		case key.Matches(msg, p.keyMap.Tab):
106			if p.session.ID == "" {
107				return p, nil
108			}
109			p.chatFocused = !p.chatFocused
110			if p.chatFocused {
111				cmds = append(cmds, p.layout.FocusPanel(layout.LeftPanel))
112				cmds = append(cmds, util.CmdHandler(ChatFocusedMsg{Focused: true}))
113			} else {
114				cmds = append(cmds, p.layout.FocusPanel(layout.BottomPanel))
115				cmds = append(cmds, util.CmdHandler(ChatFocusedMsg{Focused: false}))
116			}
117			return p, tea.Batch(cmds...)
118		case key.Matches(msg, p.keyMap.Cancel):
119			if p.session.ID != "" {
120				// Cancel the current session's generation process
121				// This allows users to interrupt long-running operations
122				p.app.CoderAgent.Cancel(p.session.ID)
123				return p, nil
124			}
125		}
126	}
127	u, cmd := p.layout.Update(msg)
128	cmds = append(cmds, cmd)
129	p.layout = u.(layout.SplitPaneLayout)
130
131	return p, tea.Batch(cmds...)
132}
133
134func (p *chatPage) setMessages() tea.Cmd {
135	messagesContainer := layout.NewContainer(
136		chat.NewMessagesListCmp(p.app),
137		layout.WithPadding(1, 1, 0, 1),
138	)
139	return tea.Batch(p.layout.SetLeftPanel(messagesContainer), messagesContainer.Init())
140}
141
142func (p *chatPage) clearMessages() tea.Cmd {
143	return p.layout.ClearLeftPanel()
144}
145
146func (p *chatPage) sendMessage(text string, attachments []message.Attachment) tea.Cmd {
147	var cmds []tea.Cmd
148	if p.session.ID == "" {
149		session, err := p.app.Sessions.Create(context.Background(), "New Session")
150		if err != nil {
151			return util.ReportError(err)
152		}
153
154		p.session = session
155		cmd := p.setMessages()
156		if cmd != nil {
157			cmds = append(cmds, cmd)
158		}
159		cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(session)))
160	}
161
162	_, err := p.app.CoderAgent.Run(context.Background(), p.session.ID, text, attachments...)
163	if err != nil {
164		return util.ReportError(err)
165	}
166	return tea.Batch(cmds...)
167}
168
169func (p *chatPage) SetSize(width, height int) tea.Cmd {
170	return p.layout.SetSize(width, height)
171}
172
173func (p *chatPage) GetSize() (int, int) {
174	return p.layout.GetSize()
175}
176
177func (p *chatPage) View() tea.View {
178	return p.layout.View()
179}
180
181func (p *chatPage) Bindings() []key.Binding {
182	bindings := []key.Binding{
183		p.keyMap.NewSession,
184		p.keyMap.AddAttachment,
185	}
186	if p.app.CoderAgent.IsBusy() {
187		bindings = append([]key.Binding{p.keyMap.Cancel}, bindings...)
188	}
189
190	if p.chatFocused {
191		bindings = append([]key.Binding{
192			key.NewBinding(
193				key.WithKeys("tab"),
194				key.WithHelp("tab", "focus editor"),
195			),
196		}, bindings...)
197	} else {
198		bindings = append([]key.Binding{
199			key.NewBinding(
200				key.WithKeys("tab"),
201				key.WithHelp("tab", "focus chat"),
202			),
203		}, bindings...)
204	}
205
206	bindings = append(bindings, p.layout.Bindings()...)
207	return bindings
208}
209
210func NewChatPage(app *app.App) ChatPage {
211	sidebarContainer := layout.NewContainer(
212		sidebar.NewSidebarCmp(app.History, app.LSPClients),
213		layout.WithPadding(1, 1, 1, 1),
214	)
215	editorContainer := layout.NewContainer(
216		editor.NewEditorCmp(app),
217	)
218	return &chatPage{
219		app: app,
220		layout: layout.NewSplitPane(
221			layout.WithRightPanel(sidebarContainer),
222			layout.WithBottomPanel(editorContainer),
223			layout.WithFixedBottomHeight(5),
224			layout.WithFixedRightWidth(31),
225		),
226		keyMap: DefaultKeyMap(),
227	}
228}