From eca120cde65840c3ef59463184e34f0c5bbba0de Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Tue, 12 Mar 2024 18:33:57 +0000
Subject: [PATCH] tui: ability to cancel request in flight

---
 pkg/tui/tui.go | 37 +++++++++++++++++++++++--------------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index abca398..e5eeb6b 100644
--- a/pkg/tui/tui.go
+++ b/pkg/tui/tui.go
@@ -30,15 +30,16 @@ type model struct {
 	convShortname string
 
 	// application state
-	conversation *models.Conversation
-	messages     []models.Message
-	replyChan    chan string
-	err          error
+	conversation    *models.Conversation
+	messages        []models.Message
+	waitingForReply bool
+	replyChan       chan string
+	replyCancelFunc context.CancelFunc
+	err             error
 
 	// ui state
-	focus     focusState
-	isWaiting bool
-	status    string // a general status message
+	focus  focusState
+	status string // a general status message
 
 	// ui elements
 	content viewport.Model
@@ -90,7 +91,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case tea.KeyMsg:
 		switch msg.String() {
 		case "ctrl+c":
-			return m, tea.Quit
+			if m.waitingForReply {
+				m.replyCancelFunc()
+			} else {
+				return m, tea.Quit
+			}
 		case "q":
 			if m.focus != focusInput {
 				return m, tea.Quit
@@ -135,7 +140,8 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 		cmd = waitForChunk(m.replyChan) // wait for the next chunk
 	case msgResponseEnd:
-		m.isWaiting = false
+		m.replyCancelFunc = nil
+		m.waitingForReply = false
 		m.status = "Press ctrl+s to send"
 	}
 
@@ -184,7 +190,7 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 
 	m.updateContent()
 
-	m.isWaiting = false
+	m.waitingForReply = false
 	m.status = "Press ctrl+s to send"
 	return m
 }
@@ -217,8 +223,8 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 		m.updateContent()
 		m.content.GotoBottom()
 
-		m.isWaiting = true
-		m.status = "Waiting for response... (Press 's' to stop)"
+		m.waitingForReply = true
+		m.status = "Waiting for response, press ctrl+c to cancel..."
 		return m.promptLLM()
 	}
 	return nil
@@ -278,9 +284,12 @@ func (m *model) promptLLM() tea.Cmd {
 			ToolBag:     toolBag,
 		}
 
+		ctx, replyCancelFunc := context.WithCancel(context.Background())
+		m.replyCancelFunc = replyCancelFunc
+
 		// TODO: supply a reply callback and handle error
 		resp, _ := completionProvider.CreateChatCompletionStream(
-			context.Background(), requestParams, m.messages, nil, m.replyChan,
+			ctx, requestParams, m.messages, nil, m.replyChan,
 		)
 
 		return msgResponseEnd(resp)
@@ -311,7 +320,7 @@ func (m *model) updateContent() {
 
 func (m model) inputView() string {
 	var inputView string
-	if m.isWaiting {
+	if m.waitingForReply {
 		inputView = inputStyle.Faint(true).Render(m.input.View())
 	} else {
 		inputView = inputStyle.Render(m.input.View())