tracing: support environment variable propagation

cli/azd/.vscode/cspell-azd-dictionary.txt

@@ -255,7 +255,9 @@ teamcity
 testdata
 tmpl
 toplevel
+traceparent
 tracesdk
+tracestate
 tracetest
 trafficmanager
 Truef

cli/azd/cmd/middleware/telemetry.go

@@ -42,12 +42,13 @@ func (m *TelemetryMiddleware) Run(ctx context.Context, next NextFn) (*actions.Ac
 	// It does not contain user input, and is safe for telemetry emission.
 	cmdPath := events.GetCommandEventName(m.options.CommandPath)
 
+	eventName := cmdPath
 	extensionId := m.options.Annotations["extension.id"]
 	if extensionId != "" {
-		cmdPath = events.ExtensionRunEvent
+		eventName = events.ExtensionRunEvent
 	}
 
-	spanCtx, span := tracing.Start(ctx, cmdPath)
+	spanCtx, span := tracing.Start(ctx, eventName)
 
 	log.Printf("TraceID: %s", span.SpanContext().TraceID())
 

cli/azd/cmd/root.go

@@ -379,12 +379,12 @@ func NewRootCmd(
 	root.
 		UseMiddleware("debug", middleware.NewDebugMiddleware).
 		UseMiddleware("ux", middleware.NewUxMiddleware).
-		UseMiddlewareWhen("error", middleware.NewErrorMiddleware, func(descriptor *actions.ActionDescriptor) bool {
-			return !descriptor.Options.DisableTroubleshooting
-		}).
 		UseMiddlewareWhen("telemetry", middleware.NewTelemetryMiddleware, func(descriptor *actions.ActionDescriptor) bool {
 			return !descriptor.Options.DisableTelemetry
 		}).
+		UseMiddlewareWhen("error", middleware.NewErrorMiddleware, func(descriptor *actions.ActionDescriptor) bool {
+			return !descriptor.Options.DisableTroubleshooting
+		}).
 		UseMiddlewareWhen("loginGuard", middleware.NewLoginGuardMiddleware, func(descriptor *actions.ActionDescriptor) bool {
 			// Check if the command or any of its parents require login
 			current := descriptor

cli/azd/internal/mcp/host.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"log"
 
+	"github.com/azure/azure-dev/cli/azd/internal/tracing"
 	"github.com/mark3labs/mcp-go/client"
 	"github.com/mark3labs/mcp-go/client/transport"
 	"github.com/mark3labs/mcp-go/mcp"
@@ -75,7 +76,10 @@ func (h *McpHost) Start(ctx context.Context) error {
 
 		switch serverConfig.Type {
 		case "stdio":
-			serverTransport = transport.NewStdio(serverConfig.Command, serverConfig.Env, serverConfig.Args...)
+			env := serverConfig.Env
+			env = append(env, tracing.Environ(ctx)...)
+
+			serverTransport = transport.NewStdio(serverConfig.Command, env, serverConfig.Args...)
 		case "http", "":
 			httpTransport, err := transport.NewStreamableHTTP(serverConfig.Url)
 			if err != nil {

cli/azd/internal/tracing/propagation.go

@@ -0,0 +1,59 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package tracing
+
+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"
+)
+
+// ContextFromEnv initializes the tracing context from environment variables.
+func ContextFromEnv(ctx context.Context) context.Context {
+	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
+}
+
+// Environ returns environment variables for propagating tracing context.
+//
+// This can be used to set environment variables for child processes to
+// continue the trace.
+func Environ(ctx context.Context) []string {
+	tm := propagation.MapCarrier{}
+	tc := propagation.TraceContext{}
+	tc.Inject(ctx, &tm)
+
+	if parent := tm.Get(traceparentKey); parent != "" {
+		environ := []string{
+			traceparentEnv + "=" + parent,
+		}
+
+		if state := tm.Get(tracestateKey); state != "" {
+			environ = append(environ, tracestateEnv+"="+state)
+		}
+		return environ
+	}
+
+	return nil
+}

cli/azd/main.go

@@ -26,6 +26,7 @@ import (
 	"github.com/azure/azure-dev/cli/azd/cmd"
 	"github.com/azure/azure-dev/cli/azd/internal"
 	"github.com/azure/azure-dev/cli/azd/internal/telemetry"
+	"github.com/azure/azure-dev/cli/azd/internal/tracing"
 	"github.com/azure/azure-dev/cli/azd/pkg/config"
 	"github.com/azure/azure-dev/cli/azd/pkg/installer"
 	"github.com/azure/azure-dev/cli/azd/pkg/ioc"
@@ -56,6 +57,9 @@ func main() {
 	log.Printf("azd version: %s", internal.Version)
 
 	ts := telemetry.GetTelemetrySystem()
+	if ts != nil {
+		ctx = tracing.ContextFromEnv(ctx)
+	}
 
 	latest := make(chan semver.Version)
 	go fetchLatestVersion(latest)

cli/azd/test/functional/telemetry_test.go

@@ -240,15 +240,19 @@ func Test_CLI_Telemetry_NestedCommands(t *testing.T) {
 
 	// set environment modifier
 	cli.Env = append(cli.Env, "AZURE_DEV_USER_AGENT=azure_app_space_portal:v1.0.0")
+
+	// set a fixed traceparent to verify trace ID propagation for commands and nested commands
+	traceId := "246cfb7e1b1c5978f4b0fc2e41e98db6"
+	cli.Env = append(cli.Env, "TRACEPARENT="+fmt.Sprintf("00-%s-a16efae62fc848c9-01", traceId))
 	cli.WorkingDirectory = dir
 
 	envName := randomEnvName()
 
 	_, err := cli.RunCommandWithStdIn(
 		ctx,
-		// Choose the default minimal template
-		"Select a template\n\n"+stdinForInit(envName),
-		"init")
+		// Choose default name
+		"\n",
+		"init", "--minimal")
 	require.NoError(t, err)
 
 	// Remove infra folder to avoid lengthy Azure operations while asserting the intended telemetry behavior.
@@ -263,7 +267,9 @@ func Test_CLI_Telemetry_NestedCommands(t *testing.T) {
 	require.NoError(t, err)
 	defer file.Close()
 
-	_, err = cli.RunCommandWithStdIn(ctx, stdinForProvision(), "up", "--trace-log-file", traceFilePath)
+	_, err = cli.RunCommandWithStdIn(ctx,
+		stdinForInit(envName)+stdinForProvision(),
+		"up", "--trace-log-file", traceFilePath)
 	require.Error(t, err)
 
 	traceContent, err := os.ReadFile(traceFilePath)
@@ -274,7 +280,6 @@ func Test_CLI_Telemetry_NestedCommands(t *testing.T) {
 	packageCmdFound := false
 	provisionCmdFound := false
 	upCmdFound := false
-	traceId := ""
 	for scanner.Scan() {
 		if scanner.Text() == "" {
 			continue
@@ -292,9 +297,7 @@ func Test_CLI_Telemetry_NestedCommands(t *testing.T) {
 		if !packageCmdFound {
 			require.Equal(t, "cmd.package", span.Name)
 			packageCmdFound = true
-			require.NoError(t, span.SpanContext.Validate(), "invalid span context")
-			// set the traceID
-			traceId = span.SpanContext.TraceID
+			require.Equal(t, traceId, span.SpanContext.TraceID, "commands do not share a traceID")
 
 			m := attributesMap(span.Attributes)
 			require.Contains(t, m, fields.SubscriptionIdKey)