fix(rag): correctly attribute and track query embedding usage for active sessions

pkg/rag/embed/embed.go

@@ -14,9 +14,8 @@ import (
 // Embedder generates vector embeddings for text
 type Embedder struct {
 	provider       provider.Provider
-	usageHandler   func(ctx context.Context, tokens int64, cost float64) // Callback to emit usage events
-	batchSize      int                                                   // Batch size for API calls
-	maxConcurrency int                                                   // Maximum concurrent embedding batch requests
+	batchSize      int // Batch size for API calls
+	maxConcurrency int // Maximum concurrent embedding batch requests
 }
 
 // Option is a functional option for configuring the Embedder
@@ -51,44 +50,33 @@ func New(p provider.Provider, opts ...Option) *Embedder {
 	return e
 }
 
-// SetUsageHandler sets a callback to be called after each embedding with usage info
-func (e *Embedder) SetUsageHandler(handler func(ctx context.Context, tokens int64, cost float64)) {
-	e.usageHandler = handler
-}
-
 // Embed generates an embedding for a single text
-// Emits usage event immediately via handler if set
-func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, error) {
+func (e *Embedder) Embed(ctx context.Context, text string) ([]float64, int64, error) {
 	// Try to use the provider's embedding API if it implements EmbeddingProvider.
 	if embeddingProvider, ok := e.provider.(provider.EmbeddingProvider); ok {
 		result, err := embeddingProvider.CreateEmbedding(ctx, text)
 		if err != nil {
-			return nil, err
-		}
-
-		// Emit usage event immediately
-		if e.usageHandler != nil {
-			e.usageHandler(ctx, result.TotalTokens, result.Cost)
+			return nil, 0, err
 		}
 
 		slog.DebugContext(ctx, "Embedding generated",
 			"provider", e.provider.ID(),
 			"tokens", result.TotalTokens,
 			"cost", result.Cost)
 
-		return result.Embedding, nil
+		return result.Embedding, result.TotalTokens, nil
 	}
 
 	// Provider does not support embeddings via the standard interface; fail fast.
-	return nil, fmt.Errorf("provider %s does not support embeddings", e.provider.ID())
+	return nil, 0, fmt.Errorf("provider %s does not support embeddings", e.provider.ID())
 }
 
 // EmbedBatch generates embeddings for multiple texts using intelligent batching
 // If the provider supports batch embeddings, it will use parallel batch API calls
 // Otherwise, it falls back to sequential processing
-func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, error) {
+func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64, int64, error) {
 	if len(texts) == 0 {
-		return [][]float64{}, nil
+		return [][]float64{}, 0, nil
 	}
 
 	// Check if provider supports batch embeddings.
@@ -102,19 +90,21 @@ func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64,
 		"text_count", len(texts))
 
 	embeddings := make([][]float64, len(texts))
+	var totalTokens int64
 	for i, text := range texts {
-		embedding, err := e.Embed(ctx, text)
+		embedding, tokens, err := e.Embed(ctx, text)
+		totalTokens += tokens
 		if err != nil {
-			return nil, fmt.Errorf("failed to embed text %d: %w", i, err)
+			return nil, totalTokens, fmt.Errorf("failed to embed text %d: %w", i, err)
 		}
 		embeddings[i] = embedding
 	}
 
-	return embeddings, nil
+	return embeddings, totalTokens, nil
 }
 
 // embedBatchOptimized processes texts in optimized batches with parallel API calls
-func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, error) {
+func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provider.BatchEmbeddingProvider, texts []string) ([][]float64, int64, error) {
 	totalTexts := len(texts)
 	slog.DebugContext(ctx, "Starting optimized batch embedding",
 		"provider", e.provider.ID(),
@@ -125,6 +115,7 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid
 	// Pre-allocate results
 	embeddings := make([][]float64, totalTexts)
 	var mu sync.Mutex
+	var totalTokens int64
 
 	// Create errgroup with concurrency limit
 	g, ctx := errgroup.WithContext(ctx)
@@ -154,13 +145,9 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid
 			// Store results (mutex protects slice writes)
 			mu.Lock()
 			copy(embeddings[start:end], result.Embeddings)
+			totalTokens += result.TotalTokens
 			mu.Unlock()
 
-			// Emit usage event (handler should be thread-safe)
-			if e.usageHandler != nil {
-				e.usageHandler(ctx, result.TotalTokens, result.Cost)
-			}
-
 			slog.DebugContext(ctx, "Batch completed",
 				"batch", batchNum,
 				"embeddings", len(result.Embeddings),
@@ -173,13 +160,13 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid
 
 	// Wait for all batches and return first error if any
 	if err := g.Wait(); err != nil {
-		return nil, err
+		return nil, totalTokens, err
 	}
 
 	slog.DebugContext(ctx, "Batch embedding completed",
 		"provider", e.provider.ID(),
 		"total_embeddings", len(embeddings),
 		"batches_processed", (totalTexts+e.batchSize-1)/e.batchSize)
 
-	return embeddings, nil
+	return embeddings, totalTokens, nil
 }

pkg/rag/embed/embed_test.go

@@ -0,0 +1,63 @@
+package embed
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/docker/docker-agent/pkg/model/provider"
+	"github.com/docker/docker-agent/pkg/modelsdev"
+	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// mockEmbeddingProvider fakes a provider that fails on the 3rd embedding
+type mockEmbeddingProvider struct {
+	provider.Provider
+	callCount int
+}
+
+func (m *mockEmbeddingProvider) ID() modelsdev.ID {
+	return modelsdev.ID{}
+}
+
+func (m *mockEmbeddingProvider) BaseConfig() base.Config {
+	return base.Config{}
+}
+
+
+func (m *mockEmbeddingProvider) CreateEmbedding(ctx context.Context, text string) (*base.EmbeddingResult, error) {
+	m.callCount++
+	if m.callCount == 3 {
+		return nil, errors.New("simulated failure")
+	}
+	return &base.EmbeddingResult{
+		Embedding:   []float64{0.1, 0.2, 0.3},
+		TotalTokens: 10,
+		Cost:        0.001,
+	}, nil
+}
+
+// verify we implement EmbeddingProvider
+var _ provider.EmbeddingProvider = (*mockEmbeddingProvider)(nil)
+var _ provider.Provider = (*mockEmbeddingProvider)(nil)
+
+func TestEmbedBatch_PartialUsageOnError(t *testing.T) {
+	mockProv := &mockEmbeddingProvider{}
+	embedder := New(mockProv, WithBatchSize(10), WithMaxConcurrency(1))
+
+	// Provide 4 texts. The 3rd will fail.
+	texts := []string{"text1", "text2", "text3", "text4"}
+
+	embeddings, tokens, err := embedder.EmbedBatch(context.Background(), texts)
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "simulated failure")
+
+	// EmbedBatch should return nil for embeddings since it failed
+	assert.Nil(t, embeddings)
+
+	// Tokens should be 20 (10 from text1 + 10 from text2)
+	assert.Equal(t, int64(20), tokens)
+}

pkg/rag/manager.go

@@ -241,7 +241,7 @@ func (m *Manager) Initialize(ctx context.Context) (err error) {
 
 // Query searches for relevant documents using all configured strategies
 // If multiple strategies are configured, results are combined using the fusion strategy
-func (m *Manager) Query(ctx context.Context, query string) (results []database.SearchResult, err error) {
+func (m *Manager) Query(ctx context.Context, query string) (results []database.SearchResult, usage types.Usage, err error) {
 	// Start a `retrieval {rag_name}` span per the OTel GenAI semconv.
 	// The query text itself is sensitive so we never capture it on the
 	// span here — content capture is gated by a separate environment
@@ -275,13 +275,13 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 			// `:=`) so the deferred span closure sees the live values
 			// even if a future change replaces the explicit
 			// `return X, Y` form below with a bare `return`.
-			results, err = strategyImpl.Query(ctx, query, strategyCfg.Limit, strategyCfg.Threshold)
+			results, usage, err = strategyImpl.Query(ctx, query, strategyCfg.Limit, strategyCfg.Threshold)
 			if err != nil {
 				slog.ErrorContext(ctx, "[RAG Manager] Strategy query failed",
 					"rag_name", m.name,
 					"strategy", strategyName,
 					"error", err)
-				return nil, err
+				return nil, types.Usage{}, err
 			}
 
 			slog.DebugContext(ctx, "[RAG Manager] Single strategy results",
@@ -314,7 +314,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 					"num_results", len(results))
 			}
 
-			return results, nil
+			return results, usage, nil
 		}
 	}
 
@@ -326,6 +326,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 	type strategyResult struct {
 		name    string
 		results []database.SearchResult
+		usage   types.Usage
 		err     error
 	}
 
@@ -342,10 +343,11 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 			"strategy_threshold", strategyCfg.Threshold)
 
 		go func(name string, strategyImpl strategy.Strategy, cfg strategy.Config) {
-			results, err := strategyImpl.Query(ctx, query, cfg.Limit, cfg.Threshold)
+			results, usage, err := strategyImpl.Query(ctx, query, cfg.Limit, cfg.Threshold)
 			resultsChan <- strategyResult{
 				name:    name,
 				results: results,
+				usage:   usage,
 				err:     err,
 			}
 		}(strategyName, strategyImpl, strategyCfg)
@@ -359,17 +361,19 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 		select {
 		case result = <-resultsChan:
 		case <-ctx.Done():
-			return nil, ctx.Err()
+			return nil, types.Usage{}, ctx.Err()
 		}
 
 		if result.err != nil {
 			slog.ErrorContext(ctx, "[RAG Manager] Strategy query failed",
 				"rag_name", m.name,
 				"strategy", result.name,
 				"error", result.err)
-			return nil, fmt.Errorf("strategy %s failed: %w", result.name, result.err)
+			return nil, types.Usage{}, fmt.Errorf("strategy %s failed: %w", result.name, result.err)
 		}
 
+		usage.Add(result.usage)
+
 		slog.DebugContext(ctx, "[RAG Manager] Strategy returned results",
 			"rag_name", m.name,
 			"strategy", result.name,
@@ -386,15 +390,15 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 
 	// Safety check: fusion should never be nil with multiple strategies
 	if m.fusion == nil {
-		return nil, errors.New("fusion strategy is nil but multiple strategies are configured (this is a bug)")
+		return nil, types.Usage{}, errors.New("fusion strategy is nil but multiple strategies are configured (this is a bug)")
 	}
 
 	fusedResults, err := m.fusion.Fuse(strategyResults)
 	if err != nil {
 		slog.ErrorContext(ctx, "[RAG Manager] Fusion failed",
 			"rag_name", m.name,
 			"error", err)
-		return nil, fmt.Errorf("failed to fuse results: %w", err)
+		return nil, types.Usage{}, fmt.Errorf("failed to fuse results: %w", err)
 	}
 
 	slog.DebugContext(ctx, "[RAG Manager] Fusion complete",
@@ -428,11 +432,8 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 			"num_results", len(fusedResults))
 	}
 
-	// TODO: Track and emit query embedding usage
-	// For queries during agent execution, usage should be added to agent's session
-	// This requires passing session context through the RAG tool
 
-	return fusedResults, nil
+	return fusedResults, usage, nil
 }
 
 // Helper to get strategy names for logging

pkg/rag/manager_rerank_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/docker/docker-agent/pkg/modelerrors"
 	"github.com/docker/docker-agent/pkg/rag/database"
 	"github.com/docker/docker-agent/pkg/rag/strategy"
+	"github.com/docker/docker-agent/pkg/rag/types"
 )
 
 // failingReranker counts calls and always fails with a fixed error.
@@ -35,8 +36,8 @@ func (s *staticStrategy) Initialize(context.Context, []string, strategy.Chunking
 	return nil
 }
 
-func (s *staticStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, error) {
-	return s.results, nil
+func (s *staticStrategy) Query(context.Context, string, int, float64) ([]database.SearchResult, types.Usage, error) {
+	return s.results, types.Usage{}, nil
 }
 
 func (s *staticStrategy) CheckAndReindexChangedFiles(context.Context, []string, strategy.ChunkingConfig) error {
@@ -83,7 +84,7 @@ func TestQueryDisablesRerankerAfterNonRetryableError(t *testing.T) {
 	m, reranker := newRerankTestManager(t, rerankErr)
 
 	for range 3 {
-		results, err := m.Query(t.Context(), "some query")
+		results, _, err := m.Query(t.Context(), "some query")
 		require.NoError(t, err, "rerank failures must not fail the query")
 		assert.Len(t, results, 2, "original results are returned as fallback")
 	}
@@ -101,7 +102,7 @@ func TestQueryKeepsRerankerOnTransientError(t *testing.T) {
 	m, reranker := newRerankTestManager(t, rerankErr)
 
 	for range 3 {
-		results, err := m.Query(t.Context(), "some query")
+		results, _, err := m.Query(t.Context(), "some query")
 		require.NoError(t, err)
 		assert.Len(t, results, 2)
 	}
@@ -114,7 +115,7 @@ func TestQueryKeepsRerankerOnContextCancellation(t *testing.T) {
 	t.Parallel()
 	m, reranker := newRerankTestManager(t, context.Canceled)
 
-	results, err := m.Query(t.Context(), "some query")
+	results, _, err := m.Query(t.Context(), "some query")
 	require.NoError(t, err)
 	assert.Len(t, results, 2)
 	assert.Equal(t, int64(1), reranker.calls.Load())