From 1c7ad75fd52cc87e085235790de375ae6804261e Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Fri, 15 Mar 2024 06:44:42 +0000
Subject: [PATCH] tui: fixed response cancelling

---
 pkg/tui/tui.go | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index 843eea6..f62d888 100644
--- a/pkg/tui/tui.go
+++ b/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)
@@ -375,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),
 	}
@@ -551,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)
 	}
 }