Add a traceparent flag for Azure request correlation

cli/src/internal/cmd/root.go

@@ -27,6 +27,7 @@ var (
 	noAuth          bool
 	apiVersion      string
 	clientRequestID string
+	traceparent     string
 	urlParams       []string
 	headers         []string
 	headerFile      string
@@ -192,6 +193,8 @@ Examples:
 	rootCmd.PersistentFlags().StringVar(&clientRequestID, "client-request-id", "", "Set the x-ms-client-request-id header for Azure request correlation. Pass the flag without a value to generate a random ID.")
 	// Passing --client-request-id without a value generates a fresh ID for this invocation.
 	rootCmd.PersistentFlags().Lookup("client-request-id").NoOptDefVal = uuid.NewString()
+	rootCmd.PersistentFlags().StringVar(&traceparent, "traceparent", "", "Set the W3C traceparent header. Pass the flag without a value to generate one.")
+	rootCmd.PersistentFlags().Lookup("traceparent").NoOptDefVal = service.TraceparentAutoValue
 	rootCmd.PersistentFlags().StringArrayVar(&urlParams, "url-param", []string{}, "Set or append a URL query parameter (repeatable, format: key=value)")
 	rootCmd.PersistentFlags().StringArrayVarP(&headers, "header", "H", []string{}, "Custom headers (repeatable, format: Key:Value)")
 	rootCmd.PersistentFlags().StringVar(&headerFile, "header-file", "", "Read headers from a file (one Key: Value per line; blank lines and # comments ignored). -H overrides on conflict.")
@@ -267,6 +270,7 @@ func snapshotConfig() config.Config {
 		NoAuth:          noAuth,
 		APIVersion:      apiVersion,
 		ClientRequestID: clientRequestID,
+		Traceparent:     traceparent,
 		URLParams:       urlParams,
 		Headers:         headers,
 		HeaderFile:      headerFile,

cli/src/internal/cmd/root_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/jongio/azd-rest/src/internal/config"
+	"github.com/jongio/azd-rest/src/internal/service"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -31,6 +32,7 @@ func resetGlobalFlags() {
 	noAuth = false
 	apiVersion = ""
 	clientRequestID = ""
+	traceparent = ""
 	urlParams = []string{}
 	headers = []string{}
 	headerFile = ""
@@ -94,13 +96,31 @@ func TestNewRootCmd_SilentFlag(t *testing.T) {
 	assert.Empty(t, flag.Shorthand, "--silent should have no short alias")
 }
 
+func TestNewRootCmd_TraceparentFlag(t *testing.T) {
+	resetGlobalFlags()
+	cmd := NewRootCmd()
+
+	flag := cmd.PersistentFlags().Lookup("traceparent")
+	require.NotNil(t, flag, "--traceparent persistent flag should be registered")
+	assert.Empty(t, flag.DefValue, "--traceparent should default to empty")
+	assert.Equal(t, service.TraceparentAutoValue, flag.NoOptDefVal)
+	assert.Empty(t, flag.Shorthand, "--traceparent should have no short alias")
+}
+
 func TestSnapshotConfig_Silent(t *testing.T) {
 	resetGlobalFlags()
 	silent = true
 	cfg := snapshotConfig()
 	assert.True(t, cfg.Silent, "snapshotConfig should carry the silent flag")
 }
 
+func TestSnapshotConfig_Traceparent(t *testing.T) {
+	resetGlobalFlags()
+	traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+	cfg := snapshotConfig()
+	assert.Equal(t, traceparent, cfg.Traceparent, "snapshotConfig should carry the traceparent flag")
+}
+
 func TestBuildRequestOptions_Headers(t *testing.T) {
 	resetGlobalFlags()
 	headers = []string{"X-Custom: value1", "Authorization: Bearer token"}

cli/src/internal/config/config.go

@@ -13,6 +13,7 @@ type Config struct {
 	NoAuth          bool
 	APIVersion      string
 	ClientRequestID string
+	Traceparent     string
 	URLParams       []string
 	Headers         []string
 	HeaderFile      string

cli/src/internal/service/service.go

@@ -7,6 +7,8 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/rand"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -23,8 +25,14 @@ import (
 	"github.com/jongio/azd-rest/src/internal/config"
 )
 
-// clientRequestIDHeader is the Azure correlation header set by --client-request-id.
-const clientRequestIDHeader = "x-ms-client-request-id"
+const (
+	// TraceparentAutoValue tells BuildRequestOptions to generate a fresh W3C traceparent header.
+	TraceparentAutoValue = "generate"
+
+	// clientRequestIDHeader is the Azure correlation header set by --client-request-id.
+	clientRequestIDHeader = "x-ms-client-request-id"
+	traceparentHeader     = "traceparent"
+)
 
 // TokenProviderFactory creates a TokenProvider. Abstracting this allows tests
 // to inject mocks without touching real Azure credentials.
@@ -286,6 +294,15 @@ func (s *RequestService) BuildRequestOptions(cfg config.Config, method, url stri
 		opts.Headers[clientRequestIDHeader] = cfg.ClientRequestID
 	}
 
+	// The --traceparent flag is authoritative and overrides a matching -H header.
+	if cfg.Traceparent != "" {
+		value, traceErr := prepareTraceparentHeader(cfg.Traceparent)
+		if traceErr != nil {
+			return opts, nil, traceErr
+		}
+		opts.Headers[traceparentHeader] = value
+	}
+
 	// Form fields (#202): build an application/x-www-form-urlencoded body from
 	// repeatable --form-field flags. This is mutually exclusive with a raw body.
 	if len(cfg.FormFields) > 0 {
@@ -469,6 +486,92 @@ func (s *RequestService) Execute(ctx context.Context, cfg config.Config, method,
 	return nil
 }
 
+func prepareTraceparentHeader(value string) (string, error) {
+	if value == TraceparentAutoValue {
+		return generateTraceparent()
+	}
+
+	normalized := strings.ToLower(strings.TrimSpace(value))
+	if err := validateTraceparent(normalized); err != nil {
+		return "", err
+	}
+	return normalized, nil
+}
+
+func generateTraceparent() (string, error) {
+	traceID, err := randomNonZeroHex(16)
+	if err != nil {
+		return "", err
+	}
+	parentID, err := randomNonZeroHex(8)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("00-%s-%s-01", traceID, parentID), nil
+}
+
+func randomNonZeroHex(length int) (string, error) {
+	buf := make([]byte, length)
+	for {
+		if _, err := rand.Read(buf); err != nil {
+			return "", fmt.Errorf("failed to generate traceparent: %w", err)
+		}
+		if !allZeroBytes(buf) {
+			return hex.EncodeToString(buf), nil
+		}
+	}
+}
+
+func validateTraceparent(value string) error {
+	parts := strings.Split(value, "-")
+	if len(parts) != 4 {
+		return fmt.Errorf("invalid traceparent %q: expected version, trace ID, parent ID, and flags", value)
+	}
+	if parts[0] != "00" {
+		return fmt.Errorf("invalid traceparent %q: only version 00 is supported", value)
+	}
+	if !isLowerHex(parts[1], 32) || allZeroHex(parts[1]) {
+		return fmt.Errorf("invalid traceparent %q: trace ID must be 32 non-zero lowercase hex characters", value)
+	}
+	if !isLowerHex(parts[2], 16) || allZeroHex(parts[2]) {
+		return fmt.Errorf("invalid traceparent %q: parent ID must be 16 non-zero lowercase hex characters", value)
+	}
+	if !isLowerHex(parts[3], 2) {
+		return fmt.Errorf("invalid traceparent %q: trace flags must be 2 lowercase hex characters", value)
+	}
+	return nil
+}
+
+func isLowerHex(value string, length int) bool {
+	if len(value) != length {
+		return false
+	}
+	for _, ch := range value {
+		if (ch < '0' || ch > '9') && (ch < 'a' || ch > 'f') {
+			return false
+		}
+	}
+	return true
+}
+
+func allZeroHex(value string) bool {
+	for _, ch := range value {
+		if ch != '0' {
+			return false
+		}
+	}
+	return true
+}
+
+func allZeroBytes(value []byte) bool {
+	for _, b := range value {
+		if b != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 // writeResponseOutput renders the response body to stdout or --output-file,
 // choosing the raw path for binary content and the formatter path otherwise.
 func (s *RequestService) writeResponseOutput(cfg config.Config, resp *client.Response) error {

cli/src/internal/service/service_test.go

@@ -68,6 +68,58 @@ func baseTestConfig(t *testing.T) config.Config {
 	return cfg
 }
 
+func TestBuildRequestOptions_TraceparentSetsHeader(t *testing.T) {
+	svc := newTestService()
+	cfg := baseTestConfig(t)
+	cfg.Headers = []string{"traceparent: 00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-00"}
+	cfg.Traceparent = "00-4BF92F3577B34DA6A3CE929D0E0E4736-00F067AA0BA902B7-01"
+
+	opts, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://example.com")
+	if cleanup != nil {
+		cleanup()
+	}
+	require.NoError(t, err)
+	assert.Equal(t, "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", opts.Headers[traceparentHeader])
+}
+
+func TestBuildRequestOptions_TraceparentGeneratesHeader(t *testing.T) {
+	svc := newTestService()
+	cfg := baseTestConfig(t)
+	cfg.Traceparent = TraceparentAutoValue
+
+	opts, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://example.com")
+	if cleanup != nil {
+		cleanup()
+	}
+	require.NoError(t, err)
+
+	got := opts.Headers[traceparentHeader]
+	require.NotEmpty(t, got)
+	require.NoError(t, validateTraceparent(got))
+}
+
+func TestBuildRequestOptions_TraceparentRejectsInvalidBeforeTokenProvider(t *testing.T) {
+	tokenProviderCalls := 0
+	svc := NewRequestService(
+		func() (client.TokenProvider, error) {
+			tokenProviderCalls++
+			return nil, nil
+		},
+		DefaultHTTPClientFactory,
+	)
+
+	cfg := config.Defaults()
+	cfg.Traceparent = "not-a-traceparent"
+
+	_, cleanup, err := svc.BuildRequestOptions(cfg, "GET", "https://management.azure.com/subscriptions?api-version=2021-04-01")
+	if cleanup != nil {
+		cleanup()
+	}
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "invalid traceparent")
+	assert.Zero(t, tokenProviderCalls, "traceparent validation should run before token creation")
+}
+
 func TestBuildResponseHeaderBlock_StatusAndSortedHeaders(t *testing.T) {
 	resp := &client.Response{
 		Status: "200 OK",