tui: fixed response cancelling

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.replyCancelFunc()
+				m.stopSignal <- "stahp!"
 			} else {
 				return m, tea.Quit
 			}
@@ -204,11 +204,9 @@ 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)
@@ -225,7 +223,6 @@ 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 {
@@ -250,23 +247,17 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 }
 
 func (m model) View() string {
-	if m.content.Width == 0 {
+	if m.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 below view functions may do weird things
+		// without this, the m.*View() functions may crash
 		return ""
 	}
 
-	m.content.Height = m.height - m.getFixedComponentHeight()
-
 	sections := make([]string, 0, 6)
-	error := m.errorView()
-	scrollbar := m.scrollbarView()
 	sections = append(sections, m.headerView())
-	if scrollbar != "" {
-		sections = append(sections, scrollbar)
-	}
 	sections = append(sections, m.contentView())
+	error := m.errorView()
 	if error != "" {
 		sections = append(sections, error)
 	}
@@ -279,15 +270,13 @@ 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)
@@ -327,20 +316,6 @@ 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()
 }
@@ -349,11 +324,12 @@ 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 = segmentStyle.Copy().Bold(true).Foreground(lipgloss.Color("2")).Render("✅💾")
+		saving = savingStyle.Foreground(lipgloss.Color("2")).Render("✅💾")
 	} else {
-		saving = segmentStyle.Copy().Bold(true).Foreground(lipgloss.Color("1")).Render("❌💾")
+		saving = savingStyle.Foreground(lipgloss.Color("1")).Render("❌💾")
 	}
 
 	status := m.status
@@ -383,9 +359,11 @@ func (m *model) footerView() string {
 	footer := left + padding + right
 	if remaining < 0 {
 		ellipses := "... "
-		footer = footer[:m.width-len(ellipses)] + ellipses
+		// this doesn't work very well, due to trying to trim a string with
+		// ansii chars already in it
+		footer = footer[:(len(footer)+remaining)-len(ellipses)-3] + ellipses
 	}
-	return footerStyle.Render(footer)
+	return footerStyle.Width(m.width).Render(footer)
 }
 
 func initialModel(ctx *lmcli.Context, convShortname string) model {
@@ -395,6 +373,7 @@ 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),
 	}
@@ -448,6 +427,10 @@ 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,
@@ -482,8 +465,17 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 		m.updateContent()
 		m.content.GotoBottom()
 
-		m.waitingForReply = true
-		m.status = "Press ctrl+c to cancel"
+		return m.promptLLM()
+	case "ctrl+r":
+		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
@@ -538,6 +530,9 @@ 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 {
@@ -555,16 +550,25 @@ func (m *model) promptLLM() tea.Cmd {
 			m.replyChan <- msg
 		}
 
-		ctx, replyCancelFunc := context.WithCancel(context.Background())
-		m.replyCancelFunc = replyCancelFunc
+		ctx, cancel := context.WithCancel(context.Background())
+
+		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 {
+		if err != nil && !canceled {
 			return msgResponseError(err)
 		}
+
 		return msgResponseEnd(resp)
 	}
 }
@@ -632,8 +636,7 @@ func (m *model) updateContent() {
 		case models.MessageRoleAssistant:
 			icon = ""
 			style = assistantStyle
-		case models.MessageRoleToolCall:
-		case models.MessageRoleToolResult:
+		case models.MessageRoleToolCall, models.MessageRoleToolResult:
 			icon = "🔧"
 		}