Auto-install: fix/improve project extension requirement resolution

cli/azd/cmd/auto_install.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"path/filepath"
 	"slices"
 	"strconv"
 	"strings"
@@ -170,18 +171,35 @@ func checkForMatchingExtensions(
 func promptForExtensionChoice(
 	ctx context.Context,
 	console input.Console,
-	extensions []*extensions.ExtensionMetadata) (*extensions.ExtensionMetadata, error) {
+	matches []*extensions.ExtensionMetadata) (*extensions.ExtensionMetadata, error) {
 
-	if len(extensions) == 0 {
+	if len(matches) == 0 {
 		return nil, fmt.Errorf("no extensions to choose from")
 	}
 
-	if len(extensions) == 1 {
-		return extensions[0], nil
+	if len(matches) == 1 {
+		return matches[0], nil
 	}
 
-	options := make([]string, len(extensions))
-	for i, ext := range extensions {
+	// Under --no-prompt there is no basis for choosing between the matches, and guessing could
+	// install a different binary than the user expects. `azd extension install` refuses to pick a
+	// source for the same reason.
+	if console.IsNoPromptMode() {
+		choices := make([]string, 0, len(matches))
+		for _, ext := range matches {
+			choices = append(choices, fmt.Sprintf("%s (%s)", ext.Id, ext.Source))
+		}
+		slices.Sort(choices)
+
+		return nil, &internal.ErrorWithSuggestion{
+			Err: fmt.Errorf("more than one extension can be installed: %s", strings.Join(choices, ", ")),
+			Suggestion: "Run 'azd extension install <id> --source <source>' to select one, " +
+				"then run this command again.",
+		}
+	}
+
+	options := make([]string, len(matches))
+	for i, ext := range matches {
 		options[i] = fmt.Sprintf("%s (%s) - %s", ext.DisplayName, ext.Source, ext.Description)
 	}
 
@@ -193,7 +211,7 @@ func promptForExtensionChoice(
 		return nil, err
 	}
 
-	return extensions[choice], nil
+	return matches[choice], nil
 }
 
 // isBuiltInCommand checks if the given command is a built-in command by examining
@@ -338,14 +356,37 @@ func tryAutoInstallForPartialNamespace(
 func tryAutoInstallExtension(
 	ctx context.Context,
 	console input.Console,
-	extensionManager *extensions.Manager,
+	extensionManager extensionAutoInstallManager,
 	extension extensions.ExtensionMetadata) (bool, error) {
+	return tryAutoInstallExtensionVersion(ctx, console, extensionManager, extension, "")
+}
+
+type extensionAutoInstallManager interface {
+	FindExtensions(ctx context.Context, options *extensions.FilterOptions) ([]*extensions.ExtensionMetadata, error)
+	GetInstalled(options extensions.FilterOptions) (*extensions.Extension, error)
+	Install(
+		ctx context.Context,
+		extension *extensions.ExtensionMetadata,
+		versionPreference string,
+	) (*extensions.ExtensionVersion, error)
+	ListInstalled() (map[string]*extensions.Extension, error)
+}
 
+func tryAutoInstallExtensionVersion(
+	ctx context.Context,
+	console input.Console,
+	extensionManager extensionAutoInstallManager,
+	extension extensions.ExtensionMetadata,
+	versionPreference string,
+) (bool, error) {
 	// Check if the extension is already installed
-	_, err := extensionManager.GetInstalled(extensions.FilterOptions{
+	installedExtension, err := extensionManager.GetInstalled(extensions.FilterOptions{
 		Id: extension.Id,
 	})
 	if err == nil {
+		if err := validateInstalledExtensionVersion(installedExtension, versionPreference); err != nil {
+			return false, err
+		}
 		return false, nil
 	}
 
@@ -372,7 +413,7 @@ func tryAutoInstallExtension(
 		Message:      "Confirm installation",
 	})
 	if err != nil {
-		return false, nil
+		return false, err
 	}
 
 	if !shouldInstall {
@@ -381,7 +422,7 @@ func tryAutoInstallExtension(
 
 	// Install the extension
 	console.Message(ctx, fmt.Sprintf("Installing extension '%s'...\n", extension.Id))
-	_, err = extensionManager.Install(ctx, &extension, "")
+	_, err = extensionManager.Install(ctx, &extension, versionPreference)
 	if err != nil {
 		return false, fmt.Errorf("failed to install extension: %w", err)
 	}
@@ -452,6 +493,59 @@ type ExecuteResult struct {
 	LatestVersion <-chan *update.VersionInfo
 }
 
+// projectDirExists reports whether the directory azd will run in already exists. An empty cwd means
+// the caller's own directory. A --cwd that PersistentPreRunE still has to create holds no project.
+func projectDirExists(cwd string) bool {
+	if cwd == "" {
+		return true
+	}
+
+	_, err := os.Stat(cwd)
+	return err == nil
+}
+
+// newRootCmdForExecution builds the root command, constructing it from --cwd when one was supplied
+// so that cached AzdContext and ProjectConfig state resolves against the requested project. Cobra's
+// PersistentPreRunE performs the real directory change during execution, so the caller's directory
+// is restored before returning. globalOpts.Cwd is normalized to an absolute path.
+func newRootCmdForExecution(
+	rootContainer *ioc.NestedContainer,
+	globalOpts *internal.GlobalCommandOptions,
+) (cmd *cobra.Command, err error) {
+	if globalOpts.Cwd == "" {
+		return NewRootCmd(false, nil, rootContainer), nil
+	}
+
+	absoluteCwd, err := filepath.Abs(globalOpts.Cwd)
+	if err != nil {
+		return nil, fmt.Errorf("resolving cwd: %w", err)
+	}
+	globalOpts.Cwd = absoluteCwd
+
+	if _, statErr := os.Stat(absoluteCwd); os.IsNotExist(statErr) {
+		// PersistentPreRunE owns prompting for and creating a missing --cwd directory.
+		return NewRootCmd(false, nil, rootContainer), nil
+	} else if statErr != nil {
+		return nil, fmt.Errorf("checking cwd: %w", statErr)
+	}
+
+	previousCwd, err := os.Getwd()
+	if err != nil {
+		return nil, fmt.Errorf("getting current directory: %w", err)
+	}
+	if err := os.Chdir(absoluteCwd); err != nil {
+		return nil, fmt.Errorf("changing directory to %s: %w", absoluteCwd, err)
+	}
+	defer func() {
+		// Deferred so the process never keeps the temporary directory after a failure.
+		if restoreErr := os.Chdir(previousCwd); restoreErr != nil && err == nil {
+			cmd, err = nil, fmt.Errorf("restoring current directory: %w", restoreErr)
+		}
+	}()
+
+	return NewRootCmd(false, nil, rootContainer), nil
+}
+
 // ExecuteWithAutoInstall executes the command and handles auto-installation of extensions for unknown commands.
 func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContainer) *ExecuteResult {
 	result := &ExecuteResult{}
@@ -473,7 +567,12 @@ func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContai
 
 	// Creating the RootCmd takes care of registering common dependencies in rootContainer.
 	// The command tree will retrieve globalOpts from the container via its FlagsResolver.
-	rootCmd := NewRootCmd(false, nil, rootContainer)
+	rootCmd, err := newRootCmdForExecution(rootContainer, globalOpts)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, output.WithErrorFormat("ERROR: %s", err.Error()))
+		result.Err = err
+		return result
+	}
 
 	var extensionManager *extensions.Manager
 	var console input.Console
@@ -495,24 +594,61 @@ func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContai
 			result.LatestVersion = startUpdateCheck(ctx)
 		}
 
+		projectExtensions := projectExtensionResult{}
+		// A --cwd that does not exist yet holds no project, so resolving now would pick up the
+		// caller's unrelated project instead.
+		if projectDirExists(globalOpts.Cwd) {
+			projectExtensions, err = tryAutoInstallProjectExtensions(
+				ctx, rootContainer, foundCmd, originalArgs,
+			)
+			if err != nil {
+				if resolveErr := rootContainer.Resolve(&console); resolveErr != nil {
+					fmt.Fprintln(os.Stderr, output.WithErrorFormat("ERROR: %s", err.Error()))
+				} else {
+					displayAutoInstallError(ctx, console, err)
+				}
+				result.Err = err
+				return result
+			}
+
+			if projectExtensions.installed {
+				rootCmd = newRootCmdWithoutRegistration(rootContainer)
+				foundCmd, originalArgs, err = rootCmd.Find(os.Args[1:])
+				if err != nil {
+					result.Err = err
+					return result
+				}
+			}
+		}
+
 		// Check for partial namespace match (e.g., "ai" found but "ai.agent" not installed)
 		if installed := tryAutoInstallForPartialNamespace(
 			ctx, rootContainer, foundCmd, originalArgs,
 		); installed {
 			// Extension was installed, rebuild command tree and execute
-			rootCmd = NewRootCmd(false, nil, rootContainer)
+			rootCmd = newRootCmdWithoutRegistration(rootContainer)
 			result.Err = rootCmd.ExecuteContext(ctx)
 			return result
 		}
 
-		// Known command, proceed with normal execution
-		err := rootCmd.ExecuteContext(ctx)
+		// Known command, proceed with normal execution. The failure is held separately because the
+		// auto-install below declares its own err, and every path out of here has to report it.
+		commandErr := rootCmd.ExecuteContext(ctx)
 
 		// Only attempt service-host auto-install when the command failed with that specific error.
 		// Other command errors (for example, unsupported output formats) should be returned directly.
-		unsupportedErr, ok := errors.AsType[*project.UnsupportedServiceHostError](err)
+		unsupportedErr, ok := errors.AsType[*project.UnsupportedServiceHostError](commandErr)
 		if !ok {
-			result.Err = err
+			result.Err = commandErr
+			return result
+		}
+		if projectExtensions.handled {
+			if resolveErr := rootContainer.Resolve(&console); resolveErr != nil {
+				fmt.Fprintln(os.Stderr, unsupportedErr.ErrorMessage)
+			} else {
+				console.Message(ctx, unsupportedErr.ErrorMessage)
+			}
+			result.Err = commandErr
 			return result
 		}
 
@@ -529,17 +665,32 @@ func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContai
 			Provider:   requiredHost,
 		})
 		if err != nil {
-			// Do not fail if we couldn't check for extensions - just proceed to normal execution
+			// Do not fail if we couldn't check for extensions - just report the command's own failure
 			log.Println("Error: check for extensions. Skipping auto-install:", err)
 			console.Message(ctx, unsupportedErr.ErrorMessage)
+			result.Err = commandErr
 			return result
 		}
-		// Note: We don't need to filter or check which extensions are installed.
-		// If any of these extensions would be installed, the auto-install wouldn't have been triggered because
-		// there would be at least one extensions providing the capability and provider.
+		installedExtensions, err := extensionManager.ListInstalled()
+		if err != nil {
+			log.Println("Error: list installed extensions. Skipping auto-install:", err)
+			console.Message(ctx, unsupportedErr.ErrorMessage)
+			// Auto-install could not run, so the command's own failure stands.
+			result.Err = commandErr
+			return result
+		}
+		// Offer only the extensions whose selected version supplies the host and that are not
+		// already installed.
+		availableExtensionsForHost = filterExtensionsForProvider(
+			availableExtensionsForHost,
+			extensions.ServiceTargetProviderCapability,
+			requiredHost,
+		)
+		availableExtensionsForHost = uninstalledExtensionMatches(availableExtensionsForHost, installedExtensions)
 		if len(availableExtensionsForHost) == 0 {
-			// did not find an extension with the capability, just print the original error message
+			// Nothing can be installed to supply the host, so the command's failure stands.
 			console.Message(ctx, unsupportedErr.ErrorMessage)
+			result.Err = commandErr
 			return result
 		}
 
@@ -572,12 +723,13 @@ func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContai
 
 		if installed {
 			// Extension was installed, build command tree and execute
-			rootCmd := NewRootCmd(false, nil, rootContainer)
+			rootCmd := newRootCmdWithoutRegistration(rootContainer)
 			result.Err = rootCmd.ExecuteContext(ctx)
 			return result
 		}
 
-		result.Err = err
+		// The install was declined, so the command's failure stands.
+		result.Err = commandErr
 		return result
 	}
 
@@ -699,7 +851,7 @@ func ExecuteWithAutoInstall(ctx context.Context, rootContainer *ioc.NestedContai
 
 			if installed {
 				// Extension was installed, build command tree and execute
-				rootCmd := NewRootCmd(false, nil, rootContainer)
+				rootCmd := newRootCmdWithoutRegistration(rootContainer)
 				result.Err = rootCmd.ExecuteContext(ctx)
 				return result
 			}