Add distributed tracing and structured error handling for extensions

cli/azd/.vscode/cspell.yaml

@@ -243,6 +243,9 @@ overrides:
       - userosscache
       - docstates
       - dylib
+  - filename: cli/azd/pkg/azdext/trace_context.go
+    words:
+      - traceparent
 ignorePaths:
   - "**/*_test.go"
   - "**/mock*.go"

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/docs/extension-framework.md

@@ -334,6 +334,42 @@ 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. To ensure your extension's operations are correctly correlated with the parent `azd` process, you must hydrate the context in your extension's entry point.
+
+**Update `main.go`:**
+
+```go
+import (
+    "context"
+    "os"
+    "github.com/azure/azure-dev/cli/azd/pkg/azdext"
+)
+
+func main() {
+    ctx := context.Background()
+
+    // Hydrate context with traceparent from environment if present
+    // This ensures the extension process participates in the active trace
+    if traceparent := os.Getenv("TRACEPARENT"); traceparent != "" {
+        ctx = azdext.ContextFromTraceParent(ctx, traceparent)
+    }
+
+    rootCmd := cmd.NewRootCommand()
+
+    if err := rootCmd.ExecuteContext(ctx); err != nil {
+        // Handle error
+    }
+}
+```
+
+The `ExtensionHost` automatically handles trace propagation for all gRPC communication:
+1. **Outgoing**: Traces are injected into messages sent back to `azd` core.
+2. **Incoming**: Traces from `azd` core are extracted and injected into the `context.Context` passed to your handlers.
+
+This ensures that calls to Azure SDKs (which support `TRACEPARENT`) within your handlers will automatically be correlated.
+
 ### Developer Extension
 
 The easiest way to get started building extensions is to install the `azd` Developer extension.

cli/azd/grpc/proto/event.proto

@@ -25,6 +25,10 @@ message EventMessage {
     InvokeServiceHandler invoke_service_handler = 5;
     ServiceHandlerStatus service_handler_status = 6;
   }
+
+  // W3C traceparent format for distributed tracing propagation.
+  // Format: "00-{traceId}-{spanId}-{flags}"
+  string trace_parent = 98;
 }
 
 // Client subscribes to project-related events

cli/azd/grpc/proto/framework_service.proto

@@ -17,6 +17,10 @@ service FrameworkService {
 // Envelope for all possible framework service messages (requests and responses)
 message FrameworkServiceMessage {
   string request_id = 1;
+  // W3C Trace Context traceparent header value for distributed tracing.
+  // Format: "00-{trace-id}-{span-id}-{trace-flags}"
+  // See: https://www.w3.org/TR/trace-context/
+  string trace_parent = 98;
   FrameworkServiceErrorMessage error = 99;
   oneof message_type {
     RegisterFrameworkServiceRequest register_framework_service_request = 2;
@@ -41,6 +45,10 @@ message FrameworkServiceMessage {
 message FrameworkServiceErrorMessage {
   string message = 1;
   string details = 2;
+  // Structured error information for telemetry
+  string error_code = 3;     // Error code from the service (e.g., "Conflict", "NotFound")
+  int32 status_code = 4;     // HTTP status code (e.g., 409, 404, 500)
+  string service_name = 5;   // Service name for telemetry (e.g., "ai.azure.com", "cognitiveservices")
 }
 
 // Request to register a framework service provider

cli/azd/grpc/proto/service_target.proto

@@ -17,6 +17,10 @@ service ServiceTargetService {
 // Envelope for all possible service target messages (requests and responses)
 message ServiceTargetMessage {
   string request_id = 1;
+  // W3C Trace Context traceparent header value for distributed tracing.
+  // Format: "00-{trace-id}-{span-id}-{trace-flags}"
+  // See: https://www.w3.org/TR/trace-context/
+  string trace_parent = 98;
   ServiceTargetErrorMessage error = 99;
   oneof message_type {
     RegisterServiceTargetRequest register_service_target_request = 2;
@@ -83,6 +87,10 @@ message RegisterServiceTargetResponse {
 message ServiceTargetErrorMessage {
   string message = 2;
   string details = 3;
+  // Structured error information for telemetry
+  string error_code = 4;     // Error code from the service (e.g., "Conflict", "NotFound")
+  int32 status_code = 5;     // HTTP status code (e.g., 409, 404, 500)
+  string service_name = 6;   // Service name for telemetry (e.g., "ai.azure.com", "cognitiveservices")
 }
 
 // GetTargetResource request and response

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 extRespErr *azdext.ExtensionResponseError
 	if errors.As(err, &respErr) {
 		serviceName := "other"
 		statusCode := -1
@@ -84,6 +86,24 @@ 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, &extRespErr) {
+		// Handle structured errors from extensions (e.g., service target providers)
+		if extRespErr.HasServiceInfo() {
+			// Extension provided service information - use it for telemetry
+			serviceName, hostDomain := mapService(extRespErr.ServiceName)
+			errDetails = append(errDetails,
+				fields.ServiceName.String(serviceName),
+				fields.ServiceHost.String(hostDomain),
+				fields.ServiceStatusCode.Int(extRespErr.StatusCode),
+			)
+			if extRespErr.ErrorCode != "" {
+				errDetails = append(errDetails, fields.ServiceErrorCode.String(extRespErr.ErrorCode))
+			}
+			errCode = fmt.Sprintf("ext.service.%s.%d", serviceName, extRespErr.StatusCode)
+		} else {
+			// Extension error without service info
+			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: "WithExtensionResponseError",
+			err: &azdext.ExtensionResponseError{
+				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/event_message_envelope.go

@@ -96,6 +96,16 @@ func (ops *EventMessageEnvelope) SetError(msg *EventMessage, err error) {
 	// No-op: EventMessage uses status strings, not Error field
 }
 
+// GetTraceParent returns the W3C traceparent value from the message for distributed tracing.
+func (ops *EventMessageEnvelope) GetTraceParent(msg *EventMessage) string {
+	return msg.GetTraceParent()
+}
+
+// SetTraceParent sets the W3C traceparent value on the message for distributed tracing.
+func (ops *EventMessageEnvelope) SetTraceParent(msg *EventMessage, traceParent string) {
+	msg.TraceParent = traceParent
+}
+
 // GetInnerMessage returns the inner message from the oneof field
 func (ops *EventMessageEnvelope) GetInnerMessage(msg *EventMessage) any {
 	// The MessageType field is a oneof wrapper. We need to extract the actual inner message.