pkg/tui/tui.go

@@ -44,9 +44,9 @@ type model struct {
 	conversation    *models.Conversation
 	messages        []models.Message
 	waitingForReply bool
-	stopSignal      chan interface{}
 	replyChan       chan models.Message
 	replyChunkChan  chan string
+	replyCancelFunc context.CancelFunc
 	err             error
 	persistence     bool // whether we will save new messages in the conversation
 
@@ -124,7 +124,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		switch msg.String() {
 		case "ctrl+c":
 			if m.waitingForReply {
-				m.stopSignal <- "stahp!"
+				m.replyCancelFunc()
 			} else {
 				return m, tea.Quit
 			}
@@ -204,9 +204,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.updateContent()
 		cmds = append(cmds, m.waitForReply())
 	case msgResponseEnd:
+		m.replyCancelFunc = nil
 		m.waitingForReply = false
 		m.status = "Press ctrl+s to send"
 	case msgResponseError:
+		m.replyCancelFunc = nil
 		m.waitingForReply = false
 		m.status = "Press ctrl+s to send"
 		m.err = error(msg)
@@ -223,6 +225,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.err = error(msg)
 	}
 
+
 	var cmd tea.Cmd
 	m.spinner, cmd = m.spinner.Update(msg)
 	if cmd != nil {
@@ -247,17 +250,23 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 }
 
 func (m model) View() string {
-	if m.width == 0 {
+	if m.content.Width == 0 {
 		// this is the case upon initial startup, but it's also a safe bet that
 		// we can just skip rendering if the terminal is really 0 width...
-		// without this, the m.*View() functions may crash
+		// without this, the below view functions may do weird things
 		return ""
 	}
 
+	m.content.Height = m.height - m.getFixedComponentHeight()
+
 	sections := make([]string, 0, 6)
-	sections = append(sections, m.headerView())
-	sections = append(sections, m.contentView())
 	error := m.errorView()
+	scrollbar := m.scrollbarView()
+	sections = append(sections, m.headerView())
+	if scrollbar != "" {
+		sections = append(sections, scrollbar)
+	}
+	sections = append(sections, m.contentView())
 	if error != "" {
 		sections = append(sections, error)
 	}
@@ -270,13 +279,15 @@ func (m model) View() string {
 	)
 }
 
-// returns the total height of "fixed" components, which are those which don't
-// change height dependent on window size.
 func (m *model) getFixedComponentHeight() int {
 	h := 0
 	h += m.input.Height()
 	h += lipgloss.Height(m.headerView())
 	h += lipgloss.Height(m.footerView())
+	scrollbar := m.scrollbarView()
+	if scrollbar != "" {
+		h += lipgloss.Height(scrollbar)
+	}
 	errorView := m.errorView()
 	if errorView != "" {
 		h += lipgloss.Height(errorView)
@@ -316,6 +327,20 @@ func (m *model) errorView() string {
 		Render(fmt.Sprintf("%s", m.err))
 }
 
+func (m *model) scrollbarView() string {
+	if m.content.AtTop() {
+		return ""
+	}
+
+	count := int(m.content.ScrollPercent() * float64(m.width-2))
+	fill := strings.Repeat("-", count)
+	return lipgloss.NewStyle().
+		Width(m.width).
+		PaddingLeft(1).
+		PaddingRight(1).
+		Render(fill)
+}
+
 func (m *model) inputView() string {
 	return m.input.View()
 }
@@ -324,12 +349,11 @@ func (m *model) footerView() string {
 	segmentStyle := lipgloss.NewStyle().PaddingLeft(1).PaddingRight(1).Faint(true)
 	segmentSeparator := "|"
 
-	savingStyle := segmentStyle.Copy().Bold(true)
 	saving := ""
 	if m.persistence {
-		saving = savingStyle.Foreground(lipgloss.Color("2")).Render("✅💾")
+		saving = segmentStyle.Copy().Bold(true).Foreground(lipgloss.Color("2")).Render("✅💾")
 	} else {
-		saving = savingStyle.Foreground(lipgloss.Color("1")).Render("❌💾")
+		saving = segmentStyle.Copy().Bold(true).Foreground(lipgloss.Color("1")).Render("❌💾")
 	}
 
 	status := m.status
@@ -359,11 +383,9 @@ func (m *model) footerView() string {
 	footer := left + padding + right
 	if remaining < 0 {
 		ellipses := "... "
-		// this doesn't work very well, due to trying to trim a string with
+		footer = footer[:m.width-len(ellipses)] + ellipses
-		// ansii chars already in it
-		footer = footer[:(len(footer)+remaining)-len(ellipses)-3] + ellipses
 	}
-	return footerStyle.Width(m.width).Render(footer)
+	return footerStyle.Render(footer)
 }
 
 func initialModel(ctx *lmcli.Context, convShortname string) model {
@@ -373,7 +395,6 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 		conversation:  &models.Conversation{},
 		persistence:   true,
 
-		stopSignal:     make(chan interface{}),
 		replyChan:      make(chan models.Message),
 		replyChunkChan: make(chan string),
 	}
@@ -427,10 +448,6 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 			return nil
 		}
 
-		if len(m.messages) > 0 && m.messages[len(m.messages)-1].Role == models.MessageRoleUser {
-			return wrapError(fmt.Errorf("Can't reply to a user message"))
-		}
-
 		reply := models.Message{
 			Role:    models.MessageRoleUser,
 			Content: userInput,
@@ -465,17 +482,8 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 		m.updateContent()
 		m.content.GotoBottom()
 
-		return m.promptLLM()
+		m.waitingForReply = true
-	case "ctrl+r":
+		m.status = "Press ctrl+c to cancel"
-		if len(m.messages) == 0 {
-			return nil
-		}
-		// TODO: retry from selected message
-		if m.messages[len(m.messages)-1].Role != models.MessageRoleUser {
-			m.messages = m.messages[:len(m.messages)-1]
-			m.updateContent()
-		}
-		m.content.GotoBottom()
 		return m.promptLLM()
 	}
 	return nil
@@ -530,9 +538,6 @@ func (m *model) generateConversationTitle() tea.Cmd {
 }
 
 func (m *model) promptLLM() tea.Cmd {
-	m.waitingForReply = true
-	m.status = "Press ctrl+c to cancel"
-
 	return func() tea.Msg {
 		completionProvider, err := m.ctx.GetCompletionProvider(*m.ctx.Config.Defaults.Model)
 		if err != nil {
@@ -550,25 +555,16 @@ func (m *model) promptLLM() tea.Cmd {
 			m.replyChan <- msg
 		}
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, replyCancelFunc := context.WithCancel(context.Background())
-
+		m.replyCancelFunc = replyCancelFunc
-		canceled := false
-		go func() {
-			select {
-			case <-m.stopSignal:
-				canceled = true
-				cancel()
-			}
-		}()
 
 		resp, err := completionProvider.CreateChatCompletionStream(
 			ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
 		)
 
-		if err != nil && !canceled {
+		if err != nil {
 			return msgResponseError(err)
 		}
-
 		return msgResponseEnd(resp)
 	}
 }
@@ -636,7 +632,8 @@ func (m *model) updateContent() {
 		case models.MessageRoleAssistant:
 			icon = ""
 			style = assistantStyle
-		case models.MessageRoleToolCall, models.MessageRoleToolResult:
+		case models.MessageRoleToolCall:
+		case models.MessageRoleToolResult:
 			icon = "🔧"
 		}