Add token metrics and observability to optimizer integration (#3347)

* feat: Add token metrics and observability to optimizer integration

Author: Nigel Brown

.gitignore

@@ -42,4 +42,5 @@ cmd/thv-operator/.task/checksum/crdref-gen
 # Test coverage
 coverage*
 
-crd-helm-wrapper
+crd-helm-wrapper
+cmd/vmcp/__debug_bin*

examples/vmcp-config-optimizer.yaml

@@ -95,6 +95,19 @@ optimizer:
   # embeddingService: embedding-service-name
   # (vMCP will resolve the service DNS name)
 
+# =============================================================================
+# TELEMETRY CONFIGURATION (for Jaeger tracing)
+# =============================================================================
+# Configure OpenTelemetry to send traces to Jaeger
+telemetry:
+  endpoint: "localhost:4318"  # OTLP HTTP endpoint (Jaeger collector) - no http:// prefix needed with insecure: true
+  serviceName: "vmcp-optimizer"
+  serviceVersion: "1.0.0"  # Optional: service version
+  tracingEnabled: true
+  metricsEnabled: false  # Set to true if you want metrics too
+  samplingRate: "1.0"  # 100% sampling for development (use lower in production)
+  insecure: true  # Use HTTP instead of HTTPS
+
 # =============================================================================
 # USAGE
 # =============================================================================

pkg/optimizer/ingestion/service.go

@@ -4,10 +4,15 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/mark3labs/mcp-go/mcp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
 
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/optimizer/db"
@@ -47,6 +52,11 @@ type Service struct {
 	tokenCounter     *tokens.Counter
 	backendServerOps *db.BackendServerOps
 	backendToolOps   *db.BackendToolOps
+	tracer           trace.Tracer
+
+	// Embedding time tracking
+	embeddingTimeMu    sync.Mutex
+	totalEmbeddingTime time.Duration
 }
 
 // NewService creates a new ingestion service
@@ -80,27 +90,58 @@ func NewService(config *Config) (*Service, error) {
 	// Initialize token counter
 	tokenCounter := tokens.NewCounter()
 
-	// Create chromem-go embeddingFunc from our embedding manager
-	embeddingFunc := func(_ context.Context, text string) ([]float32, error) {
+	// Initialize tracer
+	tracer := otel.Tracer("github.com/stacklok/toolhive/pkg/optimizer/ingestion")
+
+	svc := &Service{
+		config:             config,
+		database:           database,
+		embeddingManager:   embeddingManager,
+		tokenCounter:       tokenCounter,
+		tracer:             tracer,
+		totalEmbeddingTime: 0,
+	}
+
+	// Create chromem-go embeddingFunc from our embedding manager with tracing
+	embeddingFunc := func(ctx context.Context, text string) ([]float32, error) {
+		// Create a span for embedding calculation
+		_, span := svc.tracer.Start(ctx, "optimizer.ingestion.calculate_embedding",
+			trace.WithAttributes(
+				attribute.String("operation", "embedding_calculation"),
+			))
+		defer span.End()
+
+		start := time.Now()
+
 		// Our manager takes a slice, so wrap the single text
 		embeddingsResult, err := embeddingManager.GenerateEmbedding([]string{text})
 		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
 			return nil, err
 		}
 		if len(embeddingsResult) == 0 {
-			return nil, fmt.Errorf("no embeddings generated")
+			err := fmt.Errorf("no embeddings generated")
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
+			return nil, err
 		}
+
+		// Track embedding time
+		duration := time.Since(start)
+		svc.embeddingTimeMu.Lock()
+		svc.totalEmbeddingTime += duration
+		svc.embeddingTimeMu.Unlock()
+
+		span.SetAttributes(
+			attribute.Int64("embedding.duration_ms", duration.Milliseconds()),
+		)
+
 		return embeddingsResult[0], nil
 	}
 
-	svc := &Service{
-		config:           config,
-		database:         database,
-		embeddingManager: embeddingManager,
-		tokenCounter:     tokenCounter,
-		backendServerOps: db.NewBackendServerOps(database, embeddingFunc),
-		backendToolOps:   db.NewBackendToolOps(database, embeddingFunc),
-	}
+	svc.backendServerOps = db.NewBackendServerOps(database, embeddingFunc)
+	svc.backendToolOps = db.NewBackendToolOps(database, embeddingFunc)
 
 	logger.Info("Ingestion service initialized for event-driven ingestion (chromem-go)")
 	return svc, nil
@@ -129,6 +170,16 @@ func (s *Service) IngestServer(
 	description *string,
 	tools []mcp.Tool,
 ) error {
+	// Create a span for the entire ingestion operation
+	ctx, span := s.tracer.Start(ctx, "optimizer.ingestion.ingest_server",
+		trace.WithAttributes(
+			attribute.String("server.id", serverID),
+			attribute.String("server.name", serverName),
+			attribute.Int("tools.count", len(tools)),
+		))
+	defer span.End()
+
+	start := time.Now()
 	logger.Infof("Ingesting server: %s (%d tools) [serverID=%s]", serverName, len(tools), serverID)
 
 	// Create backend server record (simplified - vMCP manages lifecycle)
@@ -144,25 +195,51 @@ func (s *Service) IngestServer(
 
 	// Create or update server (chromem-go handles embeddings)
 	if err := s.backendServerOps.Update(ctx, backendServer); err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return fmt.Errorf("failed to create/update server %s: %w", serverName, err)
 	}
 	logger.Debugf("Created/updated server: %s", serverName)
 
 	// Sync tools for this server
 	toolCount, err := s.syncBackendTools(ctx, serverID, serverName, tools)
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return fmt.Errorf("failed to sync tools for %s: %w", serverName, err)
 	}
 
-	logger.Infof("Successfully ingested server %s with %d tools", serverName, toolCount)
+	duration := time.Since(start)
+	span.SetAttributes(
+		attribute.Int64("ingestion.duration_ms", duration.Milliseconds()),
+		attribute.Int("tools.ingested", toolCount),
+	)
+
+	logger.Infow("Successfully ingested server",
+		"server_name", serverName,
+		"server_id", serverID,
+		"tools_count", toolCount,
+		"duration_ms", duration.Milliseconds())
 	return nil
 }
 
 // syncBackendTools synchronizes tools for a backend server
 func (s *Service) syncBackendTools(ctx context.Context, serverID string, serverName string, tools []mcp.Tool) (int, error) {
+	// Create a span for tool synchronization
+	ctx, span := s.tracer.Start(ctx, "optimizer.ingestion.sync_backend_tools",
+		trace.WithAttributes(
+			attribute.String("server.id", serverID),
+			attribute.String("server.name", serverName),
+			attribute.Int("tools.count", len(tools)),
+		))
+	defer span.End()
+
 	logger.Debugf("syncBackendTools: server=%s, serverID=%s, tool_count=%d", serverName, serverID, len(tools))
+
 	// Delete existing tools
 	if err := s.backendToolOps.DeleteByServer(ctx, serverID); err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return 0, fmt.Errorf("failed to delete existing tools: %w", err)
 	}
 
@@ -178,6 +255,8 @@ func (s *Service) syncBackendTools(ctx context.Context, serverID string, serverN
 		// Convert InputSchema to JSON
 		schemaJSON, err := json.Marshal(tool.InputSchema)
 		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
 			return 0, fmt.Errorf("failed to marshal input schema for tool %s: %w", tool.Name, err)
 		}
 
@@ -193,6 +272,8 @@ func (s *Service) syncBackendTools(ctx context.Context, serverID string, serverN
 		}
 
 		if err := s.backendToolOps.Create(ctx, backendTool, serverName); err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, err.Error())
 			return 0, fmt.Errorf("failed to create tool %s: %w", tool.Name, err)
 		}
 	}
@@ -228,6 +309,20 @@ func (s *Service) GetTotalToolTokens(ctx context.Context) int {
 	return 0
 }
 
+// GetTotalEmbeddingTime returns the total time spent calculating embeddings
+func (s *Service) GetTotalEmbeddingTime() time.Duration {
+	s.embeddingTimeMu.Lock()
+	defer s.embeddingTimeMu.Unlock()
+	return s.totalEmbeddingTime
+}
+
+// ResetEmbeddingTime resets the total embedding time counter
+func (s *Service) ResetEmbeddingTime() {
+	s.embeddingTimeMu.Lock()
+	defer s.embeddingTimeMu.Unlock()
+	s.totalEmbeddingTime = 0
+}
+
 // Close releases resources
 func (s *Service) Close() error {
 	var errs []error

pkg/optimizer/ingestion/service_test.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/stretchr/testify/require"
@@ -108,6 +109,78 @@ func TestServiceCreationAndIngestion(t *testing.T) {
 	require.True(t, toolNamesFound["search_web"], "search_web should be in results")
 }
 
+// TestService_EmbeddingTimeTracking tests that embedding time is tracked correctly
+func TestService_EmbeddingTimeTracking(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	tmpDir := t.TempDir()
+
+	// Try to use Ollama if available, otherwise skip test
+	embeddingConfig := &embeddings.Config{
+		BackendType: "ollama",
+		BaseURL:     "http://localhost:11434",
+		Model:       "all-minilm",
+		Dimension:   384,
+	}
+
+	embeddingManager, err := embeddings.NewManager(embeddingConfig)
+	if err != nil {
+		t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err)
+		return
+	}
+	_ = embeddingManager.Close()
+
+	// Initialize service
+	config := &Config{
+		DBConfig: &db.Config{
+			PersistPath: filepath.Join(tmpDir, "test-db"),
+		},
+		EmbeddingConfig: &embeddings.Config{
+			BackendType: "ollama",
+			BaseURL:     "http://localhost:11434",
+			Model:       "all-minilm",
+			Dimension:   384,
+		},
+	}
+
+	svc, err := NewService(config)
+	require.NoError(t, err)
+	defer func() { _ = svc.Close() }()
+
+	// Initially, embedding time should be 0
+	initialTime := svc.GetTotalEmbeddingTime()
+	require.Equal(t, time.Duration(0), initialTime, "Initial embedding time should be 0")
+
+	// Create test tools
+	tools := []mcp.Tool{
+		{
+			Name:        "test_tool_1",
+			Description: "First test tool for embedding",
+		},
+		{
+			Name:        "test_tool_2",
+			Description: "Second test tool for embedding",
+		},
+	}
+
+	// Reset embedding time before ingestion
+	svc.ResetEmbeddingTime()
+
+	// Ingest server with tools (this will generate embeddings)
+	err = svc.IngestServer(ctx, "test-server-id", "TestServer", nil, tools)
+	require.NoError(t, err)
+
+	// After ingestion, embedding time should be greater than 0
+	totalEmbeddingTime := svc.GetTotalEmbeddingTime()
+	require.Greater(t, totalEmbeddingTime, time.Duration(0),
+		"Total embedding time should be greater than 0 after ingestion")
+
+	// Reset and verify it's back to 0
+	svc.ResetEmbeddingTime()
+	resetTime := svc.GetTotalEmbeddingTime()
+	require.Equal(t, time.Duration(0), resetTime, "Embedding time should be 0 after reset")
+}
+
 // TestServiceWithOllama demonstrates using real embeddings (requires Ollama running)
 // This test can be enabled locally to verify Ollama integration
 func TestServiceWithOllama(t *testing.T) {