vecna/pkg/metrics/middleware.go

package metrics

import (
	"fmt"
	"net/http"

	"github.com/uptrace/bunrouter"
)

// Middleware returns a bunrouter middleware that records per-request Prometheus metrics.
// It reads timing from the RequestTrace stored in the context (set by server/trace.go).
// The trace target/url labels are optional; pass empty strings if not applicable.
func (r *Registry) Middleware(getTrace func(req bunrouter.Request) TraceSnapshot) bunrouter.MiddlewareFunc {
	return func(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
		return func(w http.ResponseWriter, req bunrouter.Request) error {
			rw := &statusWriter{ResponseWriter: w, status: http.StatusOK}
			err := next(rw, req)

			snap := getTrace(req)
			endpoint := req.URL.Path
			status := fmt.Sprintf("%d", rw.status)

			r.RequestsTotal.WithLabelValues(endpoint, status).Inc()
			r.RequestDuration.WithLabelValues(endpoint).Observe(snap.TotalSeconds)
			if snap.ForwardTarget != "" {
				r.ForwardDuration.WithLabelValues(snap.ForwardTarget, snap.ForwardURL).Observe(snap.ForwardSeconds)
			}
			if snap.AdapterType != "" {
				r.TranslateDuration.WithLabelValues(snap.AdapterType).Observe(snap.TranslateSeconds)
			}
			if snap.PromptTokens > 0 || snap.TotalTokens > 0 {
				r.AddTokens(snap.ForwardTarget, snap.ForwardModel, snap.PromptTokens, snap.TotalTokens)
			}

			return err
		}
	}
}

// TraceSnapshot carries the timing and usage values the metrics middleware needs.
type TraceSnapshot struct {
	TotalSeconds     float64
	ForwardSeconds   float64
	TranslateSeconds float64
	ForwardTarget    string
	ForwardURL       string
	ForwardModel     string
	AdapterType      string
	PromptTokens     int
	TotalTokens      int
}

// statusWriter wraps http.ResponseWriter to capture the written status code.
type statusWriter struct {
	http.ResponseWriter
	status int
}

func (sw *statusWriter) WriteHeader(code int) {
	sw.status = code
	sw.ResponseWriter.WriteHeader(code)
}

func (sw *statusWriter) Flush() {
	if f, ok := sw.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}