Add distributed tracing and structured error handling for extensions

cli/azd/.vscode/cspell.yaml

@@ -243,6 +243,14 @@ overrides:
       - userosscache
       - docstates
       - dylib
+  - filename: cli/azd/pkg/azdext/trace_context.go
+    words:
+      - traceparent
+      - tracestate
+  - filename: internal/grpcserver/server.go
+    words:
+      - traceparents
+      - tracestates
 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/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. Trace context is propagated through gRPC metadata, and `azd` sets `TRACEPARENT` in the environment when it launches an extension.
+
+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/framework_service.proto

@@ -41,6 +41,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

@@ -83,6 +83,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/grpcserver/server.go

@@ -11,6 +11,7 @@ import (
 	"net"
 
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
+	"go.opentelemetry.io/otel/propagation"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -81,7 +82,13 @@ func (s *Server) Start() (*ServerInfo, error) {
 	var serverInfo ServerInfo
 
 	s.grpcServer = grpc.NewServer(
-		grpc.UnaryInterceptor(s.tokenAuthInterceptor(&serverInfo)),
+		grpc.ChainUnaryInterceptor(
+			s.tokenAuthInterceptor(&serverInfo),
+			traceContextUnaryInterceptor(),
+		),
+		grpc.ChainStreamInterceptor(
+			traceContextStreamInterceptor(),
+		),
 	)
 
 	// Use ":0" to let the system assign an available random port
@@ -174,3 +181,67 @@ func generateSigningKey() ([]byte, error) {
 	}
 	return bytes, nil
 }
+
+// traceContextUnaryInterceptor extracts W3C traceparent from incoming gRPC metadata
+// and injects it into the context for distributed tracing correlation.
+func traceContextUnaryInterceptor() grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		ctx = extractTraceContextFromMetadata(ctx)
+		return handler(ctx, req)
+	}
+}
+
+// traceContextStreamInterceptor extracts W3C traceparent from incoming gRPC metadata
+// and injects it into the context for distributed tracing correlation.
+func traceContextStreamInterceptor() grpc.StreamServerInterceptor {
+	return func(
+		srv interface{},
+		ss grpc.ServerStream,
+		info *grpc.StreamServerInfo,
+		handler grpc.StreamHandler,
+	) error {
+		ctx := extractTraceContextFromMetadata(ss.Context())
+		if ctx != ss.Context() {
+			// Wrap the stream with the new context so downstream handlers see the trace span.
+			ss = &wrappedServerStream{ServerStream: ss, ctx: ctx}
+		}
+
+		return handler(srv, ss)
+	}
+}
+
+func extractTraceContextFromMetadata(ctx context.Context) context.Context {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return ctx
+	}
+
+	traceparents := md.Get(azdext.TraceparentKey)
+	if len(traceparents) == 0 {
+		return ctx
+	}
+
+	carrier := propagation.MapCarrier{azdext.TraceparentKey: traceparents[0]}
+
+	tracestates := md.Get(azdext.TracestateKey)
+	if len(tracestates) > 0 && tracestates[0] != "" {
+		carrier[azdext.TracestateKey] = tracestates[0]
+	}
+
+	return propagation.TraceContext{}.Extract(ctx, carrier)
+}
+
+// wrappedServerStream wraps a grpc.ServerStream with a custom context.
+type wrappedServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (w *wrappedServerStream) Context() context.Context {
+	return w.ctx
+}

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

@@ -35,7 +35,13 @@ type AzdClient struct {
 // WithAddress sets the address of the `azd` gRPC server.
 func WithAddress(address string) AzdClientOption {
 	return func(c *AzdClient) error {
-		connection, err := grpc.NewClient(address, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		opts := []grpc.DialOption{
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+			grpc.WithChainUnaryInterceptor(tracingUnaryInterceptor),
+			grpc.WithChainStreamInterceptor(tracingStreamInterceptor),
+		}
+
+		connection, err := grpc.NewClient(address, opts...)
 		if err != nil {
 			return err
 		}
@@ -123,7 +129,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)
@@ -189,3 +195,27 @@ func (c *AzdClient) Account() AccountServiceClient {
 
 	return c.accountClient
 }
+
+func tracingUnaryInterceptor(
+	ctx context.Context,
+	method string,
+	req, reply interface{},
+	cc *grpc.ClientConn,
+	invoker grpc.UnaryInvoker,
+	opts ...grpc.CallOption,
+) error {
+	ctx = WithTracing(ctx)
+	return invoker(ctx, method, req, reply, cc, opts...)
+}
+
+func tracingStreamInterceptor(
+	ctx context.Context,
+	desc *grpc.StreamDesc,
+	cc *grpc.ClientConn,
+	method string,
+	streamer grpc.Streamer,
+	opts ...grpc.CallOption,
+) (grpc.ClientStream, error) {
+	ctx = WithTracing(ctx)
+	return streamer(ctx, desc, cc, method, opts...)
+}

cli/azd/pkg/azdext/context.go

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azdext
+
+import (
+	"context"
+	"os"
+
+	"go.opentelemetry.io/otel/propagation"
+)
+
+// 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
+}