fix: AI discussions going back instead of preparing

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{