Add distributed tracing and structured error handling for extensions

cli/azd/cmd/extensions.go

@@ -176,6 +176,11 @@ func (a *extensionAction) Run(ctx context.Context) (*actions.ActionResult, error
 		fmt.Sprintf("AZD_ACCESS_TOKEN=%s", jwtToken),
 	)
 
+	// Propagate trace context to the extension process
+	if traceEnv := tracing.Environ(ctx); len(traceEnv) > 0 {
+		allEnv = append(allEnv, traceEnv...)
+	}
+
 	options := &extensions.InvokeOptions{
 		Args:   a.args,
 		Env:    allEnv,

cli/azd/cmd/middleware/extensions.go

@@ -15,6 +15,7 @@ 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/pkg/extensions"
 	"github.com/azure/azure-dev/cli/azd/pkg/input"
 	"github.com/azure/azure-dev/cli/azd/pkg/ioc"
@@ -132,6 +133,11 @@ func (m *ExtensionsMiddleware) Run(ctx context.Context, next NextFn) (*actions.A
 					allEnv = append(allEnv, "FORCE_COLOR=1")
 				}
 
+				// Propagate trace context to the extension process
+				if traceEnv := tracing.Environ(ctx); len(traceEnv) > 0 {
+					allEnv = append(allEnv, traceEnv...)
+				}
+
 				args := []string{"listen"}
 				if debugEnabled, _ := m.options.Flags.GetBool("debug"); debugEnabled {
 					args = append(args, "--debug")

cli/azd/docs/extension-framework.md

@@ -334,6 +334,34 @@ The build process automatically creates binaries for multiple platforms and arch
 > [!NOTE]
 > Build times may vary depending on your hardware and extension complexity.
 
+### Distributed Tracing
+
+`azd` uses OpenTelemetry and W3C Trace Context for distributed tracing. `azd` sets `TRACEPARENT` in the environment when it launches the extension process.
+
+Use `azdext.NewContext()` to hydrate the root context with trace context:
+
+```go
+func main() {
+  ctx := azdext.NewContext()
+  rootCmd := cmd.NewRootCommand()
+  if err := rootCmd.ExecuteContext(ctx); err != nil {
+    // Handle error
+  }
+}
+```
+
+To correlate Azure SDK calls with the parent trace, add the correlation policy to your client options:
+
+```go
+import "github.com/azure/azure-dev/cli/azd/pkg/azsdk"
+
+clientOptions := &policy.ClientOptions{
+    PerCallPolicies: []policy.Policy{
+        azsdk.NewMsCorrelationPolicy(),
+    },
+}
+```
+
 ### Developer Extension
 
 The easiest way to get started building extensions is to install the `azd` Developer extension.
@@ -1895,7 +1923,7 @@ func (r *RustFrameworkProvider) Package(ctx context.Context, serviceConfig *azde
 
 // Register the framework provider
 func main() {
-    ctx := azdext.WithAccessToken(context.Background())
+    ctx := azdext.WithAccessToken(azdext.NewContext())
     azdClient, err := azdext.NewAzdClient()
     if err != nil {
         log.Fatal(err)
@@ -2011,7 +2039,7 @@ func (v *VMServiceTargetProvider) Endpoints(ctx context.Context, serviceConfig *
 
 // Register the service target provider
 func main() {
-    ctx := azdext.WithAccessToken(context.Background())
+    ctx := azdext.WithAccessToken(azdext.NewContext())
     azdClient, err := azdext.NewAzdClient()
     if err != nil {
         log.Fatal(err)

cli/azd/grpc/proto/errors.proto

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+syntax = "proto3";
+
+package azdext;
+
+option go_package = "github.com/azure/azure-dev/cli/azd/pkg/azdext";
+
+// ErrorOrigin indicates where an error originated from.
+// This helps with telemetry categorization and error handling strategies.
+enum ErrorOrigin {
+  ERROR_ORIGIN_UNSPECIFIED = 0;  // Default/unknown origin
+  ERROR_ORIGIN_LOCAL       = 1;  // Local errors: config, filesystem, auth state, validation, etc.
+  ERROR_ORIGIN_SERVICE     = 2;  // HTTP/gRPC upstream service errors
+  ERROR_ORIGIN_TOOL        = 3;  // Subprocess, plugin, or external tool errors
+}
+
+// ServiceErrorDetail contains structured error information from an HTTP/gRPC service.
+// Used when ErrorOrigin is ERROR_ORIGIN_SERVICE.
+message ServiceErrorDetail {
+  string error_code   = 1;  // Error code from the service (e.g., "Conflict", "NotFound", "RateLimitExceeded")
+  int32 status_code   = 2;  // HTTP status code (e.g., 409, 404, 500)
+  string service_name = 3;  // Service host/name for telemetry (e.g., "ai.azure.com", "management.azure.com")
+}
+
+// ExtensionError is a unified error message that can represent errors from different sources.
+// It provides structured error information for telemetry and error handling.
+message ExtensionError {
+  string message = 2;       // Human-readable error message
+  string details = 3;       // Additional error details
+  ErrorOrigin origin = 4;   // Where the error originated from
+
+  // Source-specific structured details. Only one should be set based on origin.
+  oneof source {
+    ServiceErrorDetail service_error = 10;
+    // ToolErrorDetail tool_error    = 11;  // Reserved for future use
+    // LocalErrorDetail local_error  = 12;  // Reserved for future use
+  }
+}

cli/azd/grpc/proto/framework_service.proto

@@ -8,6 +8,7 @@ option go_package = "github.com/azure/azure-dev/cli/azd/pkg/azdext";
 
 import "models.proto";
 import "service_target.proto";
+import "errors.proto";
 
 service FrameworkService {
   // Bidirectional stream for framework service requests and responses
@@ -17,7 +18,7 @@ service FrameworkService {
 // Envelope for all possible framework service messages (requests and responses)
 message FrameworkServiceMessage {
   string request_id = 1;
-  FrameworkServiceErrorMessage error = 99;
+  ExtensionError error = 99;
   oneof message_type {
     RegisterFrameworkServiceRequest register_framework_service_request = 2;
     RegisterFrameworkServiceResponse register_framework_service_response = 3;
@@ -37,12 +38,6 @@ message FrameworkServiceMessage {
   }
 }
 
-// Error message for framework service operations
-message FrameworkServiceErrorMessage {
-  string message = 1;
-  string details = 2;
-}
-
 // Request to register a framework service provider
 message RegisterFrameworkServiceRequest {
   string language = 1; // unique identifier for the language/framework (e.g., "rust", "go", "php")

cli/azd/grpc/proto/service_target.proto

@@ -8,6 +8,7 @@ option go_package = "github.com/azure/azure-dev/cli/azd/pkg/azdext";
 
 import "include/google/protobuf/struct.proto";
 import "models.proto";
+import "errors.proto";
 
 service ServiceTargetService {
   // Bidirectional stream for service target requests and responses
@@ -17,7 +18,7 @@ service ServiceTargetService {
 // Envelope for all possible service target messages (requests and responses)
 message ServiceTargetMessage {
   string request_id = 1;
-  ServiceTargetErrorMessage error = 99;
+  ExtensionError error = 99;
   oneof message_type {
     RegisterServiceTargetRequest register_service_target_request = 2;
     RegisterServiceTargetResponse register_service_target_response = 3;
@@ -80,11 +81,6 @@ message RegisterServiceTargetResponse {
   // Add fields as needed (empty for now)
 }
 
-message ServiceTargetErrorMessage {
-  string message = 2;
-  string details = 3;
-}
-
 // GetTargetResource request and response
 message GetTargetResourceRequest {
   string subscription_id = 1;

cli/azd/internal/cmd/errors.go

@@ -17,6 +17,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
 	"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/azdext"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
 	"github.com/azure/azure-dev/cli/azd/pkg/extensions"
 	"go.opentelemetry.io/otel/attribute"
@@ -33,6 +34,7 @@ func MapError(err error, span tracing.Span) {
 	var toolExecErr *exec.ExitError
 	var authFailedErr *auth.AuthFailedError
 	var extensionRunErr *extensions.ExtensionRunError
+	var extServiceErr *azdext.ServiceError
 	if errors.As(err, &respErr) {
 		serviceName := "other"
 		statusCode := -1
@@ -84,6 +86,22 @@ 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, &extServiceErr) {
+		// Handle structured service errors from extensions
+		if extServiceErr.StatusCode > 0 && extServiceErr.ServiceName != "" {
+			serviceName, hostDomain := mapService(extServiceErr.ServiceName)
+			errDetails = append(errDetails,
+				fields.ServiceName.String(serviceName),
+				fields.ServiceHost.String(hostDomain),
+				fields.ServiceStatusCode.Int(extServiceErr.StatusCode),
+			)
+			if extServiceErr.ErrorCode != "" {
+				errDetails = append(errDetails, fields.ServiceErrorCode.String(extServiceErr.ErrorCode))
+			}
+			errCode = fmt.Sprintf("ext.service.%s.%d", serviceName, extServiceErr.StatusCode)
+		} else {
+			errCode = "ext.service.failed"
+		}
 	} else if errors.As(err, &toolExecErr) {
 		toolName := "other"
 		cmdName := cmdAsName(toolExecErr.Cmd)

cli/azd/internal/cmd/errors_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/internal/tracing/fields"
 	"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/azdext"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
 	"github.com/azure/azure-dev/cli/azd/test/mocks/mocktracing"
 	"github.com/stretchr/testify/require"
@@ -148,6 +149,23 @@ func Test_MapError(t *testing.T) {
 				fields.ErrorKey(fields.ServiceCorrelationId.Key).String("12345"),
 			},
 		},
+		{
+			name: "WithExtServiceError",
+			err: &azdext.ServiceError{
+				Message:     "Rate limit exceeded",
+				Details:     "Too many requests",
+				ErrorCode:   "RateLimitExceeded",
+				StatusCode:  429,
+				ServiceName: "openai.azure.com",
+			},
+			wantErrReason: "ext.service.openai.429",
+			wantErrDetails: []attribute.KeyValue{
+				fields.ErrorKey(fields.ServiceName.Key).String("openai"),
+				fields.ErrorKey(fields.ServiceHost.Key).String("openai.azure.com"),
+				fields.ErrorKey(fields.ServiceStatusCode.Key).Int(429),
+				fields.ErrorKey(fields.ServiceErrorCode.Key).String("RateLimitExceeded"),
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

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

@@ -29,6 +29,9 @@ var Domains = []Domain{
 	{"azurefd.net", "frontdoor"},
 	{"scm.azurewebsites.net", "kudu"},
 	{"azurewebsites.net", "websites"},
+	{"services.ai.azure.com", "ai"},
+	{"cognitiveservices.azure.com", "cognitiveservices"},
+	{"openai.azure.com", "openai"},
 	{"blob.core.windows.net", "blob"},
 	{"cloudapp.azure.com", "vm"},
 	{"cloudapp.net", "vm"},

cli/azd/pkg/azdext/azd_client.go

@@ -123,7 +123,7 @@ func (c *AzdClient) Deployment() DeploymentServiceClient {
 	return c.deploymentClient
 }
 
-// Deployment returns the deployment service client.
+// Events returns the event service client.
 func (c *AzdClient) Events() EventServiceClient {
 	if c.eventsClient == nil {
 		c.eventsClient = NewEventServiceClient(c.connection)

cli/azd/pkg/azdext/context.go

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azdext
+
+import (
+	"context"
+	"os"
+
+	"go.opentelemetry.io/otel/propagation"
+)
+
+const (
+	TraceparentKey = "traceparent"
+	TracestateKey  = "tracestate"
+
+	// https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/env-carriers.md
+
+	TraceparentEnv = "TRACEPARENT"
+	TracestateEnv  = "TRACESTATE"
+)
+
+// NewContext initializes a new context with tracing information extracted from environment variables.
+func NewContext() context.Context {
+	ctx := context.Background()
+	parent := os.Getenv(TraceparentEnv)
+	state := os.Getenv(TracestateEnv)
+
+	if parent != "" {
+		tc := propagation.TraceContext{}
+		return tc.Extract(ctx, propagation.MapCarrier{
+			TraceparentKey: parent,
+			TracestateKey:  state,
+		})
+	}
+
+	return ctx
+}

cli/azd/pkg/azdext/context_test.go

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azdext
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func TestNewContext_FromEnvironment(t *testing.T) {
+	traceparent := "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+	tracestate := "rojo=00f067aa0ba902b7"
+
+	t.Setenv(TraceparentEnv, traceparent)
+	t.Setenv(TracestateEnv, tracestate)
+
+	ctx := NewContext()
+	sc := trace.SpanContextFromContext(ctx)
+
+	require.True(t, sc.IsValid(), "span context should be extracted from environment")
+	traceID, err := trace.TraceIDFromHex("4bf92f3577b34da6a3ce929d0e0e4736")
+	require.NoError(t, err)
+	spanID, err := trace.SpanIDFromHex("00f067aa0ba902b7")
+	require.NoError(t, err)
+	require.Equal(t, traceID, sc.TraceID())
+	require.Equal(t, spanID, sc.SpanID())
+	require.Equal(t, tracestate, sc.TraceState().String())
+}
+
+func TestNewContext_NoEnvironment(t *testing.T) {
+	t.Setenv(TraceparentEnv, "")
+	t.Setenv(TracestateEnv, "")
+
+	ctx := NewContext()
+	sc := trace.SpanContextFromContext(ctx)
+
+	require.False(t, sc.IsValid(), "span context should be absent when no environment is set")
+}