fix: AI discussions going back instead of preparing

ai_tui.go

@@ -27,6 +27,7 @@ const (
 
 type aiPreparedMsg struct {
 	preview  AIPreview
+	threadID string
 	err      error
 }
 
@@ -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{

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

version.go

@@ -1,3 +1,3 @@
 package main
 
-const dipleVersion = "0.3.2"
+const dipleVersion = "0.3.3"