From 42e95718347fd219cc6f4895a537cf91feb19101 Mon Sep 17 00:00:00 2001
From: pablu <me@pablu.de>
Date: Mon, 3 Aug 2026 12:00:56 +0200
Subject: [PATCH] fix: AI discussions going back instead of preparing

---
 ai_tui.go   | 27 ++++++++++++++++++--------
 tui_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 version.go  |  2 +-
 3 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/ai_tui.go b/ai_tui.go
index fce3529..acc504d 100644
--- a/ai_tui.go
+++ b/ai_tui.go
@@ -26,8 +26,9 @@ const (
 )
 
 type aiPreparedMsg struct {
-	preview AIPreview
-	err     error
+	preview  AIPreview
+	threadID string
+	err      error
 }
 
 type aiCompletedMsg struct {
@@ -81,23 +82,24 @@ func (m *App) startAIDiscussion(threadID string) {
 func (m *App) beginAIPrepare(threadID, message string) tea.Cmd {
 	if m.ai == nil || !m.ai.config.Enabled {
 		m.err = errors.New("AI integration is disabled; set ai.enabled = true")
-		m.aiMode = aiMenu
+		m.returnFromAIPrepareFailure(threadID)
 		return nil
 	}
-	if m.loading {
+	if m.loading && threadID == "" {
 		m.err = errors.New("AI preparation is unavailable while PR data is refreshing")
-		m.aiMode = aiMenu
+		m.returnFromAIPrepareFailure(threadID)
 		return nil
 	}
 	if m.details.FromCache {
 		m.err = errors.New("AI preparation requires current live PR data, not a cached snapshot")
-		m.aiMode = aiMenu
+		m.returnFromAIPrepareFailure(threadID)
 		return nil
 	}
 	ctx, cancel := context.WithCancel(context.Background())
 	m.aiCancel = cancel
 	controller, details := m.ai, m.details
 	m.aiMode = aiPreparing
+	m.err = nil
 	m.aiSpinner = 0
 	started := time.Now()
 	summary := "Checking the provider and loading the authenticated GitHub diff"
@@ -110,11 +112,19 @@ func (m *App) beginAIPrepare(threadID, message string) tea.Cmd {
 	}
 	prepare := func() tea.Msg {
 		preview, err := controller.Prepare(ctx, details, threadID, message)
-		return aiPreparedMsg{preview: preview, err: err}
+		return aiPreparedMsg{preview: preview, threadID: threadID, err: err}
 	}
 	return tea.Batch(prepare, nextAIAnimationTick())
 }
 
+func (m *App) returnFromAIPrepareFailure(threadID string) {
+	if threadID != "" {
+		m.aiMode = aiDiscussion
+		return
+	}
+	m.aiMode = aiMenu
+}
+
 func (m *App) beginAIRun() tea.Cmd {
 	if m.details.HeadOID != m.aiPreview.HeadOID {
 		m.err = errors.New("PR head changed after preparation; prepare the AI review again")
@@ -210,7 +220,8 @@ func (m App) updateAI(msg tea.Msg) (tea.Model, tea.Cmd, bool) {
 		}
 		m.aiCancel = nil
 		if msg.err != nil {
-			m.err, m.aiMode = msg.err, aiMenu
+			m.err = msg.err
+			m.returnFromAIPrepareFailure(msg.threadID)
 		} else {
 			m.aiPreview, m.aiMode, m.aiPreviewScroll, m.err = msg.preview, aiConfirm, 0, nil
 			m.aiStatus = AIProviderStatus{
diff --git a/tui_test.go b/tui_test.go
index a35d76d..31ad0fe 100644
--- a/tui_test.go
+++ b/tui_test.go
@@ -856,6 +856,61 @@ func TestVimModeAppliesToEveryTextInput(t *testing.T) {
 	}
 }
 
+func TestVimAIDiscussionSubmitDuringRefreshPreservesDraft(t *testing.T) {
+	config := defaultAIConfig()
+	config.Enabled = true
+	settings := defaultAppSettings()
+	settings.AI = &AIController{config: config}
+	m := NewAppWithSettings(nil, "o", "r", false, 50, time.Second, settings)
+	m.screen, m.loading, m.width, m.height = threadScreen, true, 80, 24
+	m.details = PRDetails{
+		HeadOID: "head",
+		Threads: []ReviewThread{{ID: "thread", Path: "main.go"}},
+	}
+	m.aiMode, m.writeThreadID, m.aiInput = aiDiscussion, "thread", "Keep this question"
+	m.resetInputEditor(&m.aiInputEditor, m.aiInput)
+
+	updated, command, handled := m.updateAI(tea.KeyMsg{Type: tea.KeyCtrlS})
+	m = updated.(App)
+	if !handled || command == nil || m.aiMode != aiPreparing {
+		t.Fatalf("submit handled=%v command=%v mode=%v", handled, command, m.aiMode)
+	}
+	if m.aiInput != "Keep this question" || m.aiInputEditor.Text != m.aiInput {
+		t.Fatalf("AI discussion draft changed during submit: input=%q editor=%q",
+			m.aiInput, m.aiInputEditor.Text)
+	}
+}
+
+func TestAIDiscussionPreparationFailurePreservesDraft(t *testing.T) {
+	m := NewApp(nil, "o", "r", false, 50, time.Second)
+	m.aiMode, m.writeThreadID, m.aiInput = aiPreparing, "thread", "Keep this question"
+	m.resetInputEditor(&m.aiInputEditor, m.aiInput)
+
+	updated, command, handled := m.updateAI(aiPreparedMsg{
+		threadID: "thread", err: errors.New("prepare failed"),
+	})
+	m = updated.(App)
+	if !handled || command != nil || m.aiMode != aiDiscussion || m.err == nil {
+		t.Fatalf("failure handled=%v command=%v mode=%v err=%v",
+			handled, command, m.aiMode, m.err)
+	}
+	if m.aiInput != "Keep this question" || m.aiInputEditor.Text != m.aiInput {
+		t.Fatalf("AI discussion draft lost after failure: input=%q editor=%q",
+			m.aiInput, m.aiInputEditor.Text)
+	}
+}
+
+func TestPullRequestAIPreparationFailureIgnoresOldDiscussionThread(t *testing.T) {
+	m := NewApp(nil, "o", "r", false, 50, time.Second)
+	m.aiMode, m.writeThreadID = aiPreparing, "old-thread"
+
+	updated, _, _ := m.updateAI(aiPreparedMsg{err: errors.New("prepare failed")})
+	m = updated.(App)
+	if m.aiMode != aiMenu {
+		t.Fatalf("pull-request preparation failure returned to mode %v, want AI menu", m.aiMode)
+	}
+}
+
 func TestVimModeUsesGlobalFooterBar(t *testing.T) {
 	m := NewApp(nil, "o", "r", false, 50, time.Second)
 	m.width, m.height = 80, 10
diff --git a/version.go b/version.go
index a1b48dd..1719ec9 100644
--- a/version.go
+++ b/version.go
@@ -1,3 +1,3 @@
 package main
 
-const dipleVersion = "0.3.2"
+const dipleVersion = "0.3.3"