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