git.stella-ops.org/src/Cli/StellaOps.Cli/Telemetry/TraceparentHttpMessageHandler.cs

using System;
using System.Diagnostics;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;

namespace StellaOps.Cli.Telemetry;

/// <summary>
/// HTTP message handler that propagates W3C Trace Context (traceparent) headers.
/// Per CLI-OBS-50-001, ensures CLI HTTP client propagates traceparent headers for all commands,
/// prints correlation IDs on failure, and records trace IDs in verbose logs.
/// </summary>
public sealed class TraceparentHttpMessageHandler : DelegatingHandler
{
    private const string TraceparentHeader = "traceparent";
    private const string TracestateHeader = "tracestate";
    private const string RequestIdHeader = "x-request-id";
    private const string CorrelationIdHeader = "x-correlation-id";

    private readonly ILogger<TraceparentHttpMessageHandler> _logger;
    private readonly bool _verbose;

    public TraceparentHttpMessageHandler(
        ILogger<TraceparentHttpMessageHandler> logger,
        bool verbose = false)
    {
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _verbose = verbose;
    }

    protected override async Task<HttpResponseMessage> SendAsync(
        HttpRequestMessage request,
        CancellationToken cancellationToken)
    {
        var activity = Activity.Current;
        string? traceId = null;
        string? spanId = null;

        // Generate or use existing trace context
        if (activity is not null)
        {
            traceId = activity.TraceId.ToString();
            spanId = activity.SpanId.ToString();

            // Add W3C traceparent header if not already present
            if (!request.Headers.Contains(TraceparentHeader))
            {
                var traceparent = $"00-{traceId}-{spanId}-{(activity.Recorded ? "01" : "00")}";
                request.Headers.TryAddWithoutValidation(TraceparentHeader, traceparent);

                if (_verbose)
                {
                    _logger.LogDebug("Added traceparent header: {Traceparent}", traceparent);
                }
            }

            // Add tracestate if present
            if (!string.IsNullOrWhiteSpace(activity.TraceStateString) &&
                !request.Headers.Contains(TracestateHeader))
            {
                request.Headers.TryAddWithoutValidation(TracestateHeader, activity.TraceStateString);
            }
        }
        else
        {
            // Generate a new trace ID if no activity exists
            traceId = Guid.NewGuid().ToString("N");
            spanId = Guid.NewGuid().ToString("N")[..16];

            if (!request.Headers.Contains(TraceparentHeader))
            {
                var traceparent = $"00-{traceId}-{spanId}-00";
                request.Headers.TryAddWithoutValidation(TraceparentHeader, traceparent);

                if (_verbose)
                {
                    _logger.LogDebug("Generated new traceparent header: {Traceparent}", traceparent);
                }
            }
        }

        // Also add x-request-id for legacy compatibility
        if (!request.Headers.Contains(RequestIdHeader))
        {
            request.Headers.TryAddWithoutValidation(RequestIdHeader, traceId);
        }

        if (_verbose)
        {
            _logger.LogDebug(
                "Sending {Method} {Uri} with trace_id={TraceId}",
                request.Method,
                ScrubUrl(request.RequestUri),
                traceId);
        }

        HttpResponseMessage response;
        try
        {
            response = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
        }
        catch (Exception ex) when (ex is not OperationCanceledException)
        {
            _logger.LogError(
                "Request failed: {Method} {Uri} trace_id={TraceId} error={Error}",
                request.Method,
                ScrubUrl(request.RequestUri),
                traceId,
                ex.Message);
            throw;
        }

        // Extract correlation ID from response if present
        var responseTraceId = GetResponseTraceId(response);

        if (!response.IsSuccessStatusCode)
        {
            _logger.LogWarning(
                "Request returned {StatusCode}: {Method} {Uri} trace_id={TraceId} response_trace_id={ResponseTraceId}",
                (int)response.StatusCode,
                request.Method,
                ScrubUrl(request.RequestUri),
                traceId,
                responseTraceId ?? "(not provided)");
        }
        else if (_verbose)
        {
            _logger.LogDebug(
                "Request completed {StatusCode}: {Method} {Uri} trace_id={TraceId}",
                (int)response.StatusCode,
                request.Method,
                ScrubUrl(request.RequestUri),
                traceId);
        }

        return response;
    }

    private static string? GetResponseTraceId(HttpResponseMessage response)
    {
        if (response.Headers.TryGetValues(CorrelationIdHeader, out var correlationValues))
        {
            return string.Join(",", correlationValues);
        }

        if (response.Headers.TryGetValues(RequestIdHeader, out var requestIdValues))
        {
            return string.Join(",", requestIdValues);
        }

        if (response.Headers.TryGetValues("x-trace-id", out var traceIdValues))
        {
            return string.Join(",", traceIdValues);
        }

        return null;
    }

    private static string ScrubUrl(Uri? uri)
    {
        if (uri is null)
            return "(null)";

        // Remove query string to avoid logging sensitive parameters
        return $"{uri.Scheme}://{uri.Authority}{uri.AbsolutePath}";
    }
}

/// <summary>
/// Extension to add traceparent propagation to HTTP client.
/// </summary>
public static class TraceparentHttpClientBuilderExtensions
{
    /// <summary>
    /// Adds W3C Trace Context (traceparent) header propagation to the HTTP client.
    /// Per CLI-OBS-50-001.
    /// </summary>
    public static IHttpClientBuilder AddTraceparentPropagation(
        this IHttpClientBuilder builder,
        bool verbose = false)
    {
        return builder.AddHttpMessageHandler(sp =>
        {
            var loggerFactory = sp.GetRequiredService<ILoggerFactory>();
            return new TraceparentHttpMessageHandler(
                loggerFactory.CreateLogger<TraceparentHttpMessageHandler>(),
                verbose);
        });
    }
}