pkg/metrics/metrics.go

package metrics

import (
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"sync"
	"sync/atomic"
)

// Counter is a monotonically increasing counter
type Counter struct {
	val  uint64
	desc string
}

func NewCounter(desc string) *Counter {
	return &Counter{desc: desc}
}

func (c *Counter) Inc() {
	atomic.AddUint64(&c.val, 1)
}

func (c *Counter) Add(delta uint64) {
	atomic.AddUint64(&c.val, delta)
}

func (c *Counter) Write(buf *bytes.Buffer, name string, labels string) {
	val := atomic.LoadUint64(&c.val)
	// Write HELP and TYPE only once per metric family, but here we do it per metric which is redundant but acceptable for simple text format
	// Actually, prom format prefers HELP/TYPE once.
	// We'll handle that in Registry.Handler
	if labels != "" {
		buf.WriteString(fmt.Sprintf("%s{%s} %d\n", name, labels, val))
	} else {
		buf.WriteString(fmt.Sprintf("%s %d\n", name, val))
	}
}

// CounterVec manages a set of counters with labels
type CounterVec struct {
	desc       string
	labelNames []string
	metrics    sync.Map // map[string]*Counter (key is formatted label string)
}

func NewCounterVec(desc string, labelNames []string) *CounterVec {
	return &CounterVec{desc: desc, labelNames: labelNames}
}

func (cv *CounterVec) WithLabelValues(values ...string) *Counter {
	if len(values) != len(cv.labelNames) {
		return &Counter{} // Return dummy to avoid panic
	}

	var keyBuilder bytes.Buffer
	for i, name := range cv.labelNames {
		if i > 0 {
			keyBuilder.WriteString(",")
		}
		keyBuilder.WriteString(fmt.Sprintf("%s=\"%s\"", name, values[i]))
	}
	key := keyBuilder.String()

	if v, ok := cv.metrics.Load(key); ok {
		return v.(*Counter)
	}

	c := NewCounter(cv.desc)
	actual, loaded := cv.metrics.LoadOrStore(key, c)
	if loaded {
		return actual.(*Counter)
	}
	return c
}

// Registry holds all metrics
type Registry struct {
	metrics sync.Map
}

var DefaultRegistry = &Registry{}

func (r *Registry) Register(name string, m interface{}) {
	r.metrics.Store(name, m)
}

// Handler returns an HTTP handler for metrics
func (r *Registry) Handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		buf := new(bytes.Buffer)

		var names []string
		r.metrics.Range(func(key, value interface{}) bool {
			names = append(names, key.(string))
			return true
		})
		sort.Strings(names)

		for _, name := range names {
			val, _ := r.metrics.Load(name)
			
			switch m := val.(type) {
			case *Counter:
				buf.WriteString(fmt.Sprintf("# HELP %s %s\n", name, m.desc))
				buf.WriteString(fmt.Sprintf("# TYPE %s counter\n", name))
				m.Write(buf, name, "")
				
			case *CounterVec:
				buf.WriteString(fmt.Sprintf("# HELP %s %s\n", name, m.desc))
				buf.WriteString(fmt.Sprintf("# TYPE %s counter\n", name))
				
				// Collect keys to sort them for consistent output
				var keys []string
				m.metrics.Range(func(k, v interface{}) bool {
					keys = append(keys, k.(string))
					return true
				})
				sort.Strings(keys)
				
				for _, key := range keys {
					if c, ok := m.metrics.Load(key); ok {
						counter := c.(*Counter)
						counter.Write(buf, name, key)
					}
				}
			}
		}

		w.Header().Set("Content-Type", "text/plain; version=0.0.4")
		w.Write(buf.Bytes())
	})
}

// Global metrics
var (
	AgentIterations = NewCounter("Total number of agent iterations")
	ToolExecutions  = NewCounterVec("Total number of tool executions", []string{"tool", "status"})
	ToolCacheHits   = NewCounterVec("Total number of tool cache hits", []string{"tool"})
	ToolCacheMisses = NewCounterVec("Total number of tool cache misses", []string{"tool"})
	SandboxExecs    = NewCounterVec("Total number of sandboxed executions", []string{"status"})
	LLMRequests     = NewCounterVec("Total number of LLM requests", []string{"model", "status"})
)

func init() {
	DefaultRegistry.Register("agent_iterations_total", AgentIterations)
	DefaultRegistry.Register("tool_executions_total", ToolExecutions)
	DefaultRegistry.Register("tool_cache_hits_total", ToolCacheHits)
	DefaultRegistry.Register("tool_cache_misses_total", ToolCacheMisses)
	DefaultRegistry.Register("sandbox_executions_total", SandboxExecs)
	DefaultRegistry.Register("llm_requests_total", LLMRequests)
}