Propagate trace context to the extension process

Author: JeffreyCA

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/pkg/azdext/trace_context.go

@@ -11,10 +11,13 @@ import (
 
 const traceparentKey = "traceparent"
 
-// GetTraceParentFromContext extracts the current span context and formats it as a W3C traceparent string.
-// Returns empty string if no valid span context is found.
-func GetTraceParentFromContext(ctx context.Context) string {
-	carrier := propagation.MapCarrier{}
-	propagation.TraceContext{}.Inject(ctx, carrier)
-	return carrier.Get(traceparentKey)
+// ContextFromTraceParent creates a new context with the span context extracted from the traceparent string.
+func ContextFromTraceParent(ctx context.Context, traceparent string) context.Context {
+	if traceparent == "" {
+		return ctx
+	}
+	carrier := propagation.MapCarrier{
+		traceparentKey: traceparent,
+	}
+	return propagation.TraceContext{}.Extract(ctx, carrier)
 }