From aeeb7bb7f79ccc5a1380b47a54e4b33559381101 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Tue, 7 May 2024 08:07:48 +0000
Subject: [PATCH] tui: Add --system-prompt handling

And some state handling changes
---
 pkg/tui/chat.go          | 33 +++++++++++++++++++++------------
 pkg/tui/conversations.go |  2 +-
 pkg/tui/tui.go           | 18 +++++++++++-------
 3 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/pkg/tui/chat.go b/pkg/tui/chat.go
index f2be4e5..250e2f5 100644
--- a/pkg/tui/chat.go
+++ b/pkg/tui/chat.go
@@ -115,6 +115,14 @@ func newChatModel(tui *model) chatModel {
 		)),
 	}
 
+	system := tui.ctx.GetSystemPrompt()
+	if system != "" {
+		m.messages = []models.Message{{
+			Role:    models.MessageRoleSystem,
+			Content: system,
+		}}
+	}
+
 	m.input.Focus()
 	m.input.MaxHeight = 0
 	m.input.CharLimit = 0
@@ -177,16 +185,17 @@ func (m *chatModel) handleInput(msg tea.KeyMsg) (bool, tea.Cmd) {
 
 	switch msg.String() {
 	case "esc":
-		return true, func() tea.Msg {
-			return msgChangeState(stateConversations)
-		}
-	case "ctrl+c":
 		if m.waitingForReply {
 			m.stopSignal <- struct{}{}
 			return true, nil
 		}
 		return true, func() tea.Msg {
-			return msgChangeState(stateConversations)
+			return msgStateChange(stateConversations)
+		}
+	case "ctrl+c":
+		if m.waitingForReply {
+			m.stopSignal <- struct{}{}
+			return true, nil
 		}
 	case "ctrl+p":
 		m.persistence = !m.persistence
@@ -227,10 +236,12 @@ func (m *chatModel) handleResize(width, height int) {
 func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 	var cmds []tea.Cmd
 	switch msg := msg.(type) {
-	case msgChangeState:
+	case msgStateEnter:
 		if m.opts.convShortname != "" && m.conversation.ShortName.String != m.opts.convShortname {
 			cmds = append(cmds, m.loadConversation(m.opts.convShortname))
 		}
+		m.rebuildMessageCache()
+		m.updateContent()
 	case tea.WindowSizeMsg:
 		m.handleResize(msg.Width, msg.Height)
 	case msgTempfileEditorClosed:
@@ -254,7 +265,8 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		cmds = append(cmds, m.loadMessages(m.conversation))
 	case msgMessagesLoaded:
 		m.selectedMessage = len(msg) - 1
-		m.setMessages(msg)
+		m.messages = msg
+		m.rebuildMessageCache()
 		m.updateContent()
 		m.content.GotoBottom()
 	case msgResponseChunk:
@@ -358,10 +370,12 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		fixedHeight := height(m.views.header) + height(m.views.error) + height(m.views.footer)
 
 		// calculate clamped input height to accomodate input text
+		// minimum 4 lines, maximum half of content area
 		newHeight := max(4, min((m.height-fixedHeight-1)/2, m.input.LineCount()))
 		m.input.SetHeight(newHeight)
 		m.views.input = m.input.View()
 
+		// remaining height towards content
 		m.content.Height = m.height - fixedHeight - height(m.views.input)
 		m.views.content = m.content.View()
 	}
@@ -701,11 +715,6 @@ func (m *chatModel) footerView() string {
 	return footerStyle.Width(m.width).Render(footer)
 }
 
-func (m *chatModel) setMessages(messages []models.Message) {
-	m.messages = messages
-	m.rebuildMessageCache()
-}
-
 func (m *chatModel) setMessage(i int, msg models.Message) {
 	if i >= len(m.messages) {
 		panic("i out of range")
diff --git a/pkg/tui/conversations.go b/pkg/tui/conversations.go
index 7933718..109f86b 100644
--- a/pkg/tui/conversations.go
+++ b/pkg/tui/conversations.go
@@ -115,7 +115,7 @@ func (m *conversationsModel) handleResize(width, height int) {
 func (m conversationsModel) Update(msg tea.Msg) (conversationsModel, tea.Cmd) {
 	var cmds []tea.Cmd
 	switch msg := msg.(type) {
-	case msgChangeState:
+	case msgStateChange:
 		cmds = append(cmds, m.loadConversations())
 		m.content.SetContent(m.renderConversationList())
 	case tea.WindowSizeMsg:
diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index 0b36bc1..f7f111c 100644
--- a/pkg/tui/tui.go
+++ b/pkg/tui/tui.go
@@ -35,8 +35,10 @@ type views struct {
 }
 
 type (
-	// send to change the current app state
-	msgChangeState state
+	// send to change the current state
+	msgStateChange state
+	// sent to a state when it is entered
+	msgStateEnter struct{}
 	// sent when an error occurs
 	msgError error
 )
@@ -81,7 +83,7 @@ func (m model) Init() tea.Cmd {
 		m.conversations.Init(),
 		m.chat.Init(),
 		func() tea.Msg {
-			return msgChangeState(m.state)
+			return msgStateChange(m.state)
 		},
 	)
 }
@@ -124,18 +126,20 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		if handled {
 			return m, cmd
 		}
-	case msgChangeState:
-		switch msg {
+	case msgStateChange:
+		m.state = state(msg)
+		switch m.state {
 		case stateChat:
 			m.chat.handleResize(m.width, m.height)
 		case stateConversations:
 			m.conversations.handleResize(m.width, m.height)
 		}
-		m.state = state(msg)
+		return m, func() tea.Msg { return msgStateEnter(struct{}{}) }
 	case msgConversationSelected:
+		// passed up through conversation list model
 		m.opts.convShortname = msg.ShortName.String
 		cmds = append(cmds, func() tea.Msg {
-			return msgChangeState(stateChat)
+			return msgStateChange(stateChat)
 		})
 	case tea.WindowSizeMsg:
 		m.width, m.height = msg.Width, msg.Height