implement three solutions prompt on error handling flow (#6313)

* tmc/langchaingo 0.1.14

* add three solutions

* lll

* add enhancement

* add test

* skip error with suggestion test in CI/CD

* address UX suggestions

* make sure prompt do not run ongoing azd commands

* improve prompt

* address comment

* lll

* improve prompt

* fix bug

* fix tests

Author: hemarina

cli/azd/cmd/middleware/error.go

@@ -5,14 +5,14 @@ package middleware
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"strings"
 
 	"github.com/azure/azure-dev/cli/azd/cmd/actions"
 	"github.com/azure/azure-dev/cli/azd/internal"
 	"github.com/azure/azure-dev/cli/azd/internal/agent"
-	"github.com/azure/azure-dev/cli/azd/internal/agent/feedback"
 	"github.com/azure/azure-dev/cli/azd/internal/tracing"
 	"github.com/azure/azure-dev/cli/azd/internal/tracing/events"
 	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
@@ -53,19 +53,35 @@ func NewErrorMiddleware(
 	}
 }
 
-func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
-	actionResult, err := next(ctx)
+func (e *ErrorMiddleware) displayAgentResponse(ctx context.Context, response string, disclaimer string) {
+	if response != "" {
+		e.console.Message(ctx, disclaimer)
+		e.console.Message(ctx, "")
+		e.console.Message(ctx, fmt.Sprintf("%s:", output.AzdAgentLabel()))
+		e.console.Message(ctx, output.WithMarkdown(response))
+		e.console.Message(ctx, "")
+	}
+}
 
+func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionResult, error) {
 	// Short-circuit agentic error handling in non-interactive scenarios:
 	// - LLM feature is disabled
 	// - User specified --no-prompt (non-interactive mode)
 	// - Running in CI/CD environment where user interaction is not possible
 	if !e.featuresManager.IsEnabled(llm.FeatureLlm) || e.global.NoPrompt || resource.IsRunningOnCI() {
-		return actionResult, err
+		return next(ctx)
 	}
 
+	// Preserve a non-cancellable parent context BEFORE making the first attempt
+	// This ensures that if the context gets cancelled during the first attempt,
+	// retries can use a fresh context
+	parentCtx := context.WithoutCancel(ctx)
+
+	actionResult, err := next(parentCtx)
+
 	// Stop the spinner always to un-hide cursor
 	e.console.StopSpinner(ctx, "", input.Step)
+
 	if err == nil || e.options.IsChildAction(ctx) {
 		return actionResult, err
 	}
@@ -92,6 +108,7 @@ func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.Action
 
 	// Warn user that this is an alpha feature
 	e.console.WarnForFeature(ctx, llm.FeatureLlm)
+
 	ctx, span := tracing.Start(ctx, events.AgentTroubleshootEvent)
 	defer span.End()
 
@@ -160,31 +177,25 @@ func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.Action
 				`Steps to follow:
 			1. Use available tool including azd_error_troubleshooting tool to identify and explain the error.
 			Diagnose its root cause when running azd command.
-			2. Provide actionable troubleshooting steps. Do not perform any file changes.
+			2. Provide actionable troubleshooting steps in natural language format with clear sections.
+			DO NOT return JSON. Use readable narrative text with markdown formatting.
+			Do not perform any file changes.
 			Error details: %s`, errorInput))
 
 			if err != nil {
-				if agentOutput != "" {
-					e.console.Message(ctx, AIDisclaimer)
-					e.console.Message(ctx, output.WithMarkdown(agentOutput))
-				}
-
+				e.displayAgentResponse(ctx, agentOutput, AIDisclaimer)
 				span.SetStatus(codes.Error, "agent.send_message.failed")
 				return nil, err
 			}
 
-			e.console.Message(ctx, AIDisclaimer)
-			e.console.Message(ctx, "")
-			e.console.Message(ctx, fmt.Sprintf("%s:", output.AzdAgentLabel()))
-			e.console.Message(ctx, output.WithMarkdown(agentOutput))
-			e.console.Message(ctx, "")
+			e.displayAgentResponse(ctx, agentOutput, AIDisclaimer)
 		}
 
-		// Ask user if they want to let AI fix the
+		// Ask user if they want to let AI fix the error
 		confirmFix, err := e.checkErrorHandlingConsent(
 			ctx,
 			"mcp.errorHandling.fix",
-			fmt.Sprintf("Fix this error using %s?", agentName),
+			fmt.Sprintf("Brainstorm solutions using %s?", agentName),
 			fmt.Sprintf("This action will run AI tools to help fix the error."+
 				" Edit permissions for AI tools anytime by running %s.",
 				output.WithHighLightFormat("azd mcp consent")),
@@ -207,55 +218,90 @@ func (e *ErrorMiddleware) Run(ctx context.Context, next NextFn) (*actions.Action
 		previousError = originalError
 		agentOutput, err := azdAgent.SendMessage(ctx, fmt.Sprintf(
 			`Steps to follow:
-			1. Use available tool to identify, explain and diagnose this error when running azd command and its root cause.
-			2. Resolve the error by making the minimal, targeted change required to the code or configuration.
-			Avoid unnecessary modifications and focus only on what is essential to restore correct functionality.
-			3. Remove any changes that were created solely for validation and are not part of the actual error fix.
-			Error details: %s`, errorInput))
+            1. Use available tools to identify, explain and diagnose this error when running azd command and its root cause.
+            2. Only return a JSON object in the following format:
+            {
+              "analysis": "Brief explanation of the error and its root cause",
+              "solutions": [
+                "Solution 1 Short description (one sentence)",
+                "Solution 2 Short description (one sentence)",
+                "Solution 3 Short description (one sentence)"
+              ]
+            }
+            Provide 1-3 solutions. Each solution must be concise (one sentence).
+            Error details: %s`, errorInput))
+
+		// Extract solutions from agent output even if there's a parsing error
+		// The agent may return valid content
+		solutions := extractSuggestedSolutions(agentOutput)
+
+		// Only fail if we got an error AND couldn't extract any solutions
+		if err != nil && len(solutions) == 0 {
+			e.displayAgentResponse(ctx, agentOutput, AIDisclaimer)
+			span.SetStatus(codes.Error, "agent.send_message.failed")
+			return nil, fmt.Errorf("failed to generate solutions: %w", err)
+		}
 
+		e.console.Message(ctx, "")
+		selectedSolution, continueWithFix, err := promptUserForSolution(ctx, solutions, agentName)
 		if err != nil {
-			if agentOutput != "" {
-				e.console.Message(ctx, AIDisclaimer)
-				e.console.Message(ctx, output.WithMarkdown(agentOutput))
-			}
-
-			span.SetStatus(codes.Error, "agent.send_message.failed")
-			return nil, err
+			return nil, fmt.Errorf("prompting for solution selection: %w", err)
 		}
 
-		// Ask the user to add feedback
-		if err := e.collectAndApplyFeedback(ctx, azdAgent, AIDisclaimer); err != nil {
-			span.SetStatus(codes.Error, "agent.collect_feedback.failed")
-			return nil, err
+		if continueWithFix {
+			agentOutput, err := azdAgent.SendMessage(ctx, fmt.Sprintf(
+				`Steps to follow:
+            1. Use available tools to identify, explain and diagnose this error when running azd command and its root cause.
+            2. Resolve the error by making the minimal, targeted change required to the code or configuration.
+            Avoid unnecessary modifications and focus only on what is essential to restore correct functionality.
+            3. Remove any changes that were created solely for validation and are not part of the actual error fix.
+			4. You are currently in the middle of executing '%s'. Never run this command.
+            Error details: %s`, e.options.CommandPath, errorInput))
+
+			if err != nil {
+				e.displayAgentResponse(ctx, agentOutput, AIDisclaimer)
+				span.SetStatus(codes.Error, "agent.send_message.failed")
+				return nil, err
+			}
+
+			span.SetStatus(codes.Ok, "agent.fix.agent")
+		} else {
+			if selectedSolution != "" {
+				// User selected a solution
+				agentOutput, err = azdAgent.SendMessage(ctx, fmt.Sprintf(
+					`Steps to follow:
+						1. Perform the following actions to resolve the error: %s. 
+						During this, make minimal changes and avoid unnecessary modifications.
+						2. Remove any changes that were created solely for validation and
+						 are not part of the actual error fix.
+						3. You are currently in the middle of executing '%s'. Never run this command.
+						Error details: %s`, selectedSolution, e.options.CommandPath, errorInput))
+
+				if err != nil {
+					e.displayAgentResponse(ctx, agentOutput, AIDisclaimer)
+					span.SetStatus(codes.Error, "agent.send_message.failed")
+					return nil, err
+				}
+				span.SetStatus(codes.Ok, "agent.fix.solution")
+			} else {
+				// User selected cancel
+				span.SetStatus(codes.Error, "agent.fix.cancelled")
+				return actionResult, originalError
+			}
 		}
 
+		// Use a fresh child context per retry to avoid reusing a canceled ctx
+		attemptCtx, cancel := context.WithCancel(parentCtx)
 		// Clear check cache to prevent skip of tool related error
-		ctx = tools.WithInstalledCheckCache(ctx)
-
-		actionResult, err = next(ctx)
+		attemptCtx = tools.WithInstalledCheckCache(attemptCtx)
+		actionResult, err = next(attemptCtx)
+		cancel()
 		originalError = err
 	}
 
 	return actionResult, err
 }
 
-// collectAndApplyFeedback prompts for user feedback and applies it using the agent
-func (e *ErrorMiddleware) collectAndApplyFeedback(
-	ctx context.Context,
-	azdAgent agent.Agent,
-	AIDisclaimer string,
-) error {
-	collector := feedback.NewFeedbackCollector(e.console, feedback.FeedbackCollectorOptions{
-		EnableLoop:      false,
-		FeedbackPrompt:  "Any changes you'd like to make?",
-		FeedbackHint:    "Describe your changes or press enter to skip.",
-		RequireFeedback: false,
-		AIDisclaimer:    AIDisclaimer,
-	})
-
-	return collector.CollectFeedbackAndApply(ctx, azdAgent, AIDisclaimer)
-}
-
 func (e *ErrorMiddleware) checkErrorHandlingConsent(
 	ctx context.Context,
 	promptName string,
@@ -326,7 +372,7 @@ func promptForErrorHandlingConsent(
 		HelpMessage:     helpMessage,
 		Choices:         choices,
 		EnableFiltering: uxlib.Ptr(false),
-		DisplayCount:    5,
+		DisplayCount:    len(choices),
 	})
 
 	choiceIndex, err := selector.Ask(ctx)
@@ -340,3 +386,76 @@ func promptForErrorHandlingConsent(
 
 	return choices[*choiceIndex].Value, nil
 }
+
+// AgentResponse represents the structured JSON response from the LLM agent
+type AgentResponse struct {
+	Analysis  string   `json:"analysis"`
+	Solutions []string `json:"solutions"`
+}
+
+// extractSuggestedSolutions extracts solutions from the LLM response.
+// It expects a JSON response with the structure: {"analysis": "...", "solutions": ["...", "...", "..."]}
+// If JSON parsing fails, it returns an empty slice.
+func extractSuggestedSolutions(llmResponse string) []string {
+	var response AgentResponse
+	if err := json.Unmarshal([]byte(llmResponse), &response); err != nil {
+		return []string{}
+	}
+
+	return response.Solutions
+}
+
+// promptUserForSolution displays extracted solutions to the user and prompts them to select which solution to try.
+// Returns the selected solution text, a flag indicating if user wants to continue with AI fix, and error if any.
+func promptUserForSolution(ctx context.Context, solutions []string, agentName string) (string, bool, error) {
+	choices := make([]*uxlib.SelectChoice, len(solutions)+2)
+
+	if len(solutions) > 0 {
+		// Add the three solutions
+		for i, solution := range solutions {
+			choices[i] = &uxlib.SelectChoice{
+				Value: solution,
+				Label: "Yes. " + solution,
+			}
+		}
+	}
+
+	choices[len(solutions)] = &uxlib.SelectChoice{
+		Value: "continue",
+		Label: fmt.Sprintf("Yes, let %s choose the best approach", agentName),
+	}
+
+	choices[len(solutions)+1] = &uxlib.SelectChoice{
+		Value: "cancel",
+		Label: "No, cancel",
+	}
+
+	selector := uxlib.NewSelect(&uxlib.SelectOptions{
+		Message:         fmt.Sprintf("Allow %s to fix the error?", agentName),
+		HelpMessage:     "Select a suggested fix, or let AI decide",
+		Choices:         choices,
+		EnableFiltering: uxlib.Ptr(false),
+		DisplayCount:    len(choices),
+	})
+
+	choiceIndex, err := selector.Ask(ctx)
+	if err != nil {
+		return "", false, err
+	}
+
+	if choiceIndex == nil || *choiceIndex < 0 || *choiceIndex >= len(choices) {
+		return "", false, fmt.Errorf("invalid choice selected")
+	}
+
+	selectedValue := choices[*choiceIndex].Value
+
+	// Handle different selections
+	switch selectedValue {
+	case "continue":
+		return "", true, nil // Continue to AI fix
+	case "cancel":
+		return "", false, nil // Cancel and return error
+	default:
+		return selectedValue, false, nil // User selected a solution
+	}
+}