homelab-frontend/internal/tracing/middleware.go

package tracing

import (
	"net/http"
	"time"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
	"go.opentelemetry.io/otel/trace"
)

const tracerName = "api-gateway"

// responseWriter wraps http.ResponseWriter to capture status code.
type responseWriter struct {
	http.ResponseWriter
	statusCode int
	written    int64
}

func newResponseWriter(w http.ResponseWriter) *responseWriter {
	return &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
}

func (rw *responseWriter) WriteHeader(code int) {
	rw.statusCode = code
	rw.ResponseWriter.WriteHeader(code)
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	n, err := rw.ResponseWriter.Write(b)
	rw.written += int64(n)
	return n, err
}

// Flush forwards to the wrapped writer so that streaming handlers can still
// detect http.Flusher through this wrapper. Without it a bare
// `w.(http.Flusher)` assertion fails for every request that passes through
// this middleware, which silently disables SSE flushing and keep-alives.
func (rw *responseWriter) Flush() {
	if f, ok := rw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// Unwrap exposes the wrapped writer to http.ResponseController.
func (rw *responseWriter) Unwrap() http.ResponseWriter {
	return rw.ResponseWriter
}

// Middleware returns an HTTP middleware that adds tracing to requests.
func Middleware(next http.Handler) http.Handler {
	tracer := otel.Tracer(tracerName)
	propagator := otel.GetTextMapPropagator()

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		// Extract any existing trace context from incoming request
		ctx := propagator.Extract(r.Context(), propagation.HeaderCarrier(r.Header))

		// Start a new span
		spanName := r.Method + " " + r.URL.Path
		ctx, span := tracer.Start(ctx, spanName,
			trace.WithSpanKind(trace.SpanKindServer),
			trace.WithAttributes(
				semconv.HTTPRequestMethodKey.String(r.Method),
				semconv.URLPath(r.URL.Path),
				semconv.URLScheme(scheme(r)),
				semconv.ServerAddress(r.Host),
				semconv.UserAgentOriginal(r.UserAgent()),
				semconv.NetworkPeerAddress(r.RemoteAddr),
			),
		)
		defer span.End()

		// Add query parameters if present
		if r.URL.RawQuery != "" {
			span.SetAttributes(semconv.URLQuery(r.URL.RawQuery))
		}

		// Add model attribute for LLM requests
		if model := r.Header.Get("X-Model"); model != "" {
			span.SetAttributes(attribute.String("llm.model", model))
		}

		// Wrap response writer to capture status
		rw := newResponseWriter(w)

		// Inject trace context into response headers (for debugging)
		propagator.Inject(ctx, propagation.HeaderCarrier(w.Header()))

		// Call the next handler with traced context
		next.ServeHTTP(rw, r.WithContext(ctx))

		// Record response attributes
		duration := time.Since(start)
		span.SetAttributes(
			semconv.HTTPResponseStatusCode(rw.statusCode),
			attribute.Int64("http.response.body.size", rw.written),
			attribute.Float64("http.request.duration_ms", float64(duration.Milliseconds())),
		)

		// Set span status based on HTTP status code
		if rw.statusCode >= 400 {
			span.SetStatus(codes.Error, http.StatusText(rw.statusCode))
		} else {
			span.SetStatus(codes.Ok, "")
		}
	})
}

func scheme(r *http.Request) string {
	if r.TLS != nil {
		return "https"
	}
	if s := r.Header.Get("X-Forwarded-Proto"); s != "" {
		return s
	}
	return "http"
}