Use gRPC metadata to pass trace context

Author: JeffreyCA

cli/azd/.vscode/cspell.yaml

@@ -246,6 +246,9 @@ overrides:
   - filename: cli/azd/pkg/azdext/trace_context.go
     words:
       - traceparent
+  - filename: internal/grpcserver/server.go
+    words:
+      - traceparents
 ignorePaths:
   - "**/*_test.go"
   - "**/mock*.go"

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

@@ -336,40 +336,32 @@ The build process automatically creates binaries for multiple platforms and arch
 
 ### 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.
+`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.
 
-**Update `main.go`:**
+Use `azdext.NewContext()` to hydrate the root context with trace context:
 
 ```go
-import (
-    "context"
-    "os"
-    "github.com/azure/azure-dev/cli/azd/pkg/azdext"
-)
-
 func main() {
-    ctx := context.Background()
+  ctx := azdext.NewContext()
+  rootCmd := cmd.NewRootCommand()
+  if err := rootCmd.ExecuteContext(ctx); err != nil {
+    // Handle error
+  }
+}
+```
 
-    // 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)
-    }
+To correlate Azure SDK calls with the parent trace, add the correlation policy to your client options:
 
-    rootCmd := cmd.NewRootCommand()
+```go
+import "github.com/azure/azure-dev/cli/azd/pkg/azsdk"
 
-    if err := rootCmd.ExecuteContext(ctx); err != nil {
-        // Handle error
-    }
+clientOptions := &policy.ClientOptions{
+    PerCallPolicies: []policy.Policy{
+        azsdk.NewMsCorrelationPolicy(),
+    },
 }
 ```
 
-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.
@@ -1931,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)
@@ -2047,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/event.proto

@@ -25,10 +25,6 @@ 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,10 +17,6 @@ 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;

cli/azd/grpc/proto/service_target.proto

@@ -17,10 +17,6 @@ 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;

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/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,13 @@ 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
+}