tracing: add usage attributes for extension installs and runs

cli/azd/cmd/actions/action_descriptor.go

@@ -187,6 +187,8 @@ type ActionDescriptorOptions struct {
 	DefaultFormat output.Format
 	// Whether or not telemetry should be disabled for the current action
 	DisableTelemetry bool
+	// Whether or not troubleshooting should be disabled for the current action
+	DisableTroubleshooting bool
 	// The logic that produces the command help
 	HelpOptions ActionHelpOptions
 	// Defines grouping options for the command

cli/azd/cmd/cobra_builder.go

@@ -122,6 +122,7 @@ func (cb *CobraBuilder) configureActionResolver(cmd *cobra.Command, descriptor *
 			CommandPath: cmd.CommandPath(),
 			Aliases:     cmd.Aliases,
 			Flags:       cmd.Flags(),
+			Annotations: cmd.Annotations,
 			Args:        args,
 		}
 

cli/azd/cmd/extensions.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/azure/azure-dev/cli/azd/cmd/actions"
 	"github.com/azure/azure-dev/cli/azd/internal/grpcserver"
+	"github.com/azure/azure-dev/cli/azd/internal/tracing"
+	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
 	"github.com/azure/azure-dev/cli/azd/pkg/environment"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
 	"github.com/azure/azure-dev/cli/azd/pkg/extensions"
@@ -78,8 +80,9 @@ func bindExtension(
 	}
 
 	current.Add(lastPart, &actions.ActionDescriptorOptions{
-		Command:        cmd,
-		ActionResolver: newExtensionAction,
+		Command:                cmd,
+		ActionResolver:         newExtensionAction,
+		DisableTroubleshooting: true,
 		GroupingOptions: actions.CommandGroupOptions{
 			RootLevelHelp: actions.CmdGroupExtensions,
 		},
@@ -132,6 +135,10 @@ func (a *extensionAction) Run(ctx context.Context) (*actions.ActionResult, error
 		return nil, fmt.Errorf("failed to get extension %s: %w", extensionId, err)
 	}
 
+	tracing.SetUsageAttributes(
+		fields.ExtensionId.String(extension.Id),
+		fields.ExtensionVersion.String(extension.Version))
+
 	allEnv := []string{}
 	allEnv = append(allEnv, os.Environ()...)
 
@@ -179,7 +186,7 @@ func (a *extensionAction) Run(ctx context.Context) (*actions.ActionResult, error
 
 	_, err = a.extensionRunner.Invoke(ctx, extension, options)
 	if err != nil {
-		os.Exit(1)
+		return nil, err
 	}
 
 	return nil, nil

cli/azd/cmd/middleware/middleware.go

@@ -36,6 +36,7 @@ type Options struct {
 	Aliases     []string
 	Flags       *pflag.FlagSet
 	Args        []string
+	Annotations map[string]string
 }
 
 func (o *Options) IsChildAction(ctx context.Context) bool {

cli/azd/cmd/middleware/telemetry.go

@@ -41,6 +41,12 @@ func (m *TelemetryMiddleware) Run(ctx context.Context, next NextFn) (*actions.Ac
 	// Note: CommandPath is constructed using the Use member on each command up to the root.
 	// It does not contain user input, and is safe for telemetry emission.
 	cmdPath := events.GetCommandEventName(m.options.CommandPath)
+
+	extensionId := m.options.Annotations["extension.id"]
+	if extensionId != "" {
+		cmdPath = events.ExtensionRunEvent
+	}
+
 	spanCtx, span := tracing.Start(ctx, cmdPath)
 
 	log.Printf("TraceID: %s", span.SpanContext().TraceID())

cli/azd/cmd/middleware/ux.go

@@ -11,6 +11,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/cmd/actions"
 	"github.com/azure/azure-dev/cli/azd/internal"
 	"github.com/azure/azure-dev/cli/azd/pkg/alpha"
+	"github.com/azure/azure-dev/cli/azd/pkg/extensions"
 	"github.com/azure/azure-dev/cli/azd/pkg/input"
 	"github.com/azure/azure-dev/cli/azd/pkg/output"
 	"github.com/azure/azure-dev/cli/azd/pkg/output/ux"
@@ -57,11 +58,14 @@ func (m *UxMiddleware) Run(ctx context.Context, next NextFn) (*actions.ActionRes
 			errorMessage.WriteString("\n" + suggestionErr.Suggestion)
 		}
 
-		// UnsupportedServiceHostError is a special error which needs to float up without printing output here yet.
-		// The error is bubble up for the caller to decide to show it or not
-		var unsupportedErr *project.UnsupportedServiceHostError
 		errMessage := errorMessage.String()
-		if errors.As(err, &unsupportedErr) {
+
+		// For specific errors, we silent the output display here and let the caller handle it
+		var unsupportedErr *project.UnsupportedServiceHostError
+		var extensionRunErr *extensions.ExtensionRunError
+		if errors.As(err, &extensionRunErr) {
+			return actionResult, err
+		} else if errors.As(err, &unsupportedErr) {
 			// set the error message so the caller can use it if needed
 			unsupportedErr.ErrorMessage = errMessage
 			return actionResult, err

cli/azd/cmd/root.go

@@ -379,7 +379,9 @@ func NewRootCmd(
 	root.
 		UseMiddleware("debug", middleware.NewDebugMiddleware).
 		UseMiddleware("ux", middleware.NewUxMiddleware).
-		UseMiddleware("error", middleware.NewErrorMiddleware).
+		UseMiddlewareWhen("error", middleware.NewErrorMiddleware, func(descriptor *actions.ActionDescriptor) bool {
+			return !descriptor.Options.DisableTroubleshooting
+		}).
 		UseMiddlewareWhen("telemetry", middleware.NewTelemetryMiddleware, func(descriptor *actions.ActionDescriptor) bool {
 			return !descriptor.Options.DisableTelemetry
 		}).

cli/azd/internal/cmd/errors.go

@@ -18,6 +18,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/pkg/auth"
 	"github.com/azure/azure-dev/cli/azd/pkg/azapi"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
+	"github.com/azure/azure-dev/cli/azd/pkg/extensions"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 )
@@ -31,6 +32,7 @@ func MapError(err error, span tracing.Span) {
 	var armDeployErr *azapi.AzureDeploymentError
 	var toolExecErr *exec.ExitError
 	var authFailedErr *auth.AuthFailedError
+	var extensionRunErr *extensions.ExtensionRunError
 	if errors.As(err, &respErr) {
 		serviceName := "other"
 		statusCode := -1
@@ -80,6 +82,8 @@ func MapError(err error, span tracing.Span) {
 		}
 
 		errCode = "service.arm.deployment.failed"
+	} else if errors.As(err, &extensionRunErr) {
+		errCode = "ext.run.failed"
 	} else if errors.As(err, &toolExecErr) {
 		toolName := "other"
 		cmdName := cmdAsName(toolExecErr.Cmd)

cli/azd/internal/tracing/events/events.go

@@ -20,3 +20,9 @@ const PackBuildEvent = "tools.pack.build"
 
 // AgentTroubleshootEvent is the name of the event which tracks agent troubleshoot operations.
 const AgentTroubleshootEvent = "agent.troubleshoot"
+
+// Extension related events.
+const (
+	ExtensionRunEvent     = "ext.run"
+	ExtensionInstallEvent = "ext.install"
+)

cli/azd/internal/tracing/fields/fields.go

@@ -287,3 +287,11 @@ const (
 	// Number of auto-fix.attempts
 	AgentFixAttempts = attribute.Key("agent.fix.attempts")
 )
+
+// Extension related fields
+const (
+	// The identifier of the extension.
+	ExtensionId = attribute.Key("extension.id")
+	// The version of the extension.
+	ExtensionVersion = attribute.Key("extension.version")
+)

cli/azd/pkg/extensions/manager.go

@@ -24,6 +24,9 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Masterminds/semver/v3"
+	"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"
 	"github.com/azure/azure-dev/cli/azd/pkg/config"
 	"github.com/azure/azure-dev/cli/azd/pkg/osutil"
 	"github.com/azure/azure-dev/cli/azd/pkg/rzip"
@@ -280,11 +283,16 @@ func (m *Manager) Install(
 	ctx context.Context,
 	extension *ExtensionMetadata,
 	versionPreference string,
-) (*ExtensionVersion, error) {
+) (extVersion *ExtensionVersion, err error) {
 	if extension == nil {
 		return nil, fmt.Errorf("extension metadata cannot be nil")
 	}
 
+	ctx, span := tracing.Start(ctx, events.ExtensionInstallEvent)
+	defer func() {
+		span.EndWithStatus(err)
+	}()
+
 	installed, err := m.GetInstalled(FilterOptions{Id: extension.Id})
 	if err == nil && installed != nil {
 		return nil, fmt.Errorf("%s %w", extension.Id, ErrExtensionInstalled)
@@ -489,6 +497,10 @@ func (m *Manager) Install(
 		return nil, fmt.Errorf("failed to save user config: %w", err)
 	}
 
+	span.SetAttributes(
+		fields.ExtensionId.String(extension.Id),
+		fields.ExtensionVersion.String(selectedVersion.Version))
+
 	log.Printf(
 		"Extension '%s' (version %s) installed successfully to %s\n",
 		extension.Id,

cli/azd/pkg/extensions/runner.go

@@ -62,6 +62,18 @@ func (r *Runner) Invoke(ctx context.Context, extension *Extension, options *Invo
 	}
 
 	runResult, err := r.commandRunner.Run(ctx, runArgs)
+	if err != nil {
+		return &runResult, &ExtensionRunError{Err: err, ExtensionId: extension.Id}
+	}
+	return &runResult, nil
+}
+
+// ExtensionRunError represents an error that occurred while running an extension.
+type ExtensionRunError struct {
+	ExtensionId string
+	Err         error
+}
 
-	return &runResult, err
+func (e *ExtensionRunError) Error() string {
+	return fmt.Sprintf("extension '%s' run failed: %v", e.ExtensionId, e.Err)
 }