naco-api/internal/server/metricsMiddleware.go

package server

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	log "github.com/nacorid/logger"
	"github.com/nacorid/naco-api/internal/utils"
)

// meanStat is a cumulative mean over a mutex; used for per-tool latency (ms).
type meanStat struct {
	mu  sync.Mutex
	sum float64
	n   int64
}

func (m *meanStat) Record(v float64) {
	m.mu.Lock()
	m.sum += v
	m.n++
	m.mu.Unlock()
}

func (m *meanStat) GetMean() float64 {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.n == 0 {
		return 0
	}
	return m.sum / float64(m.n)
}

type toolStat struct {
	calls   utils.Counter
	errors  utils.Counter
	latency meanStat
}

type MetricsMiddleware struct {
	mu             sync.Mutex
	requestsBucket utils.RollingStats
	totalRequests  utils.Counter
	totalErrors    utils.Counter
	enrichEvents   utils.Counter
	toolStats      map[string]*toolStat
}

func NewMetricsHook() *MetricsMiddleware {
	return &MetricsMiddleware{
		requestsBucket: *utils.NewRollingStats(60, 1*time.Second),
		totalRequests:  *utils.NewCounter(),
		totalErrors:    *utils.NewCounter(),
		enrichEvents:   *utils.NewCounter(),
		toolStats:      make(map[string]*toolStat),
	}
}

func (mh *MetricsMiddleware) statsFor(tool string) *toolStat {
	mh.mu.Lock()
	defer mh.mu.Unlock()
	s, ok := mh.toolStats[tool]
	if !ok {
		s = &toolStat{}
		mh.toolStats[tool] = s
	}
	return s
}

func (mh *MetricsMiddleware) RecordTool(tool string, err error, duration time.Duration) {
	mh.totalRequests.Increment()
	if err != nil {
		mh.totalErrors.Increment()
	}
	s := mh.statsFor(tool)
	s.calls.Increment()
	if err != nil {
		s.errors.Increment()
	}
	s.latency.Record(float64(duration) / float64(time.Millisecond))
}

func (mh *MetricsMiddleware) RecordEnrichment() {
	mh.enrichEvents.Increment()
}

func (mh *MetricsMiddleware) OnCall(next mcp.MethodHandler) mcp.MethodHandler {
	return func(ctx context.Context, method string, request mcp.Request) (mcp.Result, error) {
		if method == "tools/call" {
			toolName := "unknown"
			if callReq, ok := request.(*mcp.CallToolRequest); ok && callReq.Params != nil {
				toolName = callReq.Params.Name
			}
			start := time.Now()
			mh.requestsBucket.Increment()
			result, err := next(ctx, method, request)
			mh.RecordTool(toolName, err, time.Since(start))
			return result, err
		}
		return next(ctx, method, request)
	}
}

func (mh *MetricsMiddleware) OnCallHTTP(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		mh.requestsBucket.Increment()
		next.ServeHTTP(w, r)
		duration := fmt.Sprintf("%.2fms", float64(time.Since(start))/float64(time.Millisecond))
		log.DebugWithContext(r.Context(), "MetricsMiddleware: HTTP Handler", "duration", duration)
	})
}

func (mh *MetricsMiddleware) GetRPS() float64 {
	return mh.requestsBucket.GetRPS()
}

func (mh *MetricsMiddleware) GetTotalRequests() int64 {
	return mh.totalRequests.Get()
}

// MetricsText renders Prometheus-style exposition text for /metrics.
func (mh *MetricsMiddleware) MetricsText() string {
	mh.mu.Lock()
	tools := make([]string, 0, len(mh.toolStats))
	for t := range mh.toolStats {
		tools = append(tools, t)
	}
	sort.Strings(tools)
	mh.mu.Unlock()

	var b strings.Builder
	fmt.Fprintf(&b, "# HELP naco_requests_total Total MCP tool calls.\n")
	fmt.Fprintf(&b, "# TYPE naco_requests_total counter\n")
	fmt.Fprintf(&b, "naco_requests_total %d\n", mh.totalRequests.Get())
	fmt.Fprintf(&b, "# HELP naco_errors_total Total MCP tool calls that returned an error.\n")
	fmt.Fprintf(&b, "# TYPE naco_errors_total counter\n")
	fmt.Fprintf(&b, "naco_errors_total %d\n", mh.totalErrors.Get())
	fmt.Fprintf(&b, "# HELP naco_enrichment_total Total background enrichment runs.\n")
	fmt.Fprintf(&b, "# TYPE naco_enrichment_total counter\n")
	fmt.Fprintf(&b, "naco_enrichment_total %d\n", mh.enrichEvents.Get())
	fmt.Fprintf(&b, "# HELP naco_requests_rps Rolling requests per second.\n")
	fmt.Fprintf(&b, "# TYPE naco_requests_rps gauge\n")
	fmt.Fprintf(&b, "naco_requests_rps %.3f\n", mh.GetRPS())

	for _, t := range tools {
		s := mh.statsFor(t)
		fmt.Fprintf(&b, "# HELP naco_tool_calls_total Tool calls for %s.\n", t)
		fmt.Fprintf(&b, "# TYPE naco_tool_calls_total counter\n")
		fmt.Fprintf(&b, "naco_tool_calls_total{tool=%q} %d\n", t, s.calls.Get())
		fmt.Fprintf(&b, "# HELP naco_tool_errors_total Tool errors for %s.\n", t)
		fmt.Fprintf(&b, "# TYPE naco_tool_errors_total counter\n")
		fmt.Fprintf(&b, "naco_tool_errors_total{tool=%q} %d\n", t, s.errors.Get())
		fmt.Fprintf(&b, "# HELP naco_tool_latency_ms Rolling mean latency (ms) for %s.\n", t)
		fmt.Fprintf(&b, "# TYPE naco_tool_latency_ms gauge\n")
		fmt.Fprintf(&b, "naco_tool_latency_ms{tool=%q} %.3f\n", t, s.latency.GetMean())
	}
	return b.String()
}