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,34 @@ 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, float64, 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, 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, result.Cost, 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, 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, float64, error) {
 	if len(texts) == 0 {
-		return [][]float64{}, nil
+		return [][]float64{}, 0, 0, nil
 	}
 
 	// Check if provider supports batch embeddings.
@@ -102,19 +91,23 @@ func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float64,
 		"text_count", len(texts))
 
 	embeddings := make([][]float64, len(texts))
+	var totalTokens int64
+	var totalCost float64
 	for i, text := range texts {
-		embedding, err := e.Embed(ctx, text)
+		embedding, tokens, cost, err := e.Embed(ctx, text)
 		if err != nil {
-			return nil, fmt.Errorf("failed to embed text %d: %w", i, err)
+			return nil, 0, 0, fmt.Errorf("failed to embed text %d: %w", i, err)
 		}
 		embeddings[i] = embedding
+		totalTokens += tokens
+		totalCost += cost
 	}
 
-	return embeddings, nil
+	return embeddings, totalTokens, totalCost, 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, float64, error) {
 	totalTexts := len(texts)
 	slog.DebugContext(ctx, "Starting optimized batch embedding",
 		"provider", e.provider.ID(),
@@ -125,6 +118,8 @@ func (e *Embedder) embedBatchOptimized(ctx context.Context, batchProvider provid
 	// Pre-allocate results
 	embeddings := make([][]float64, totalTexts)
 	var mu sync.Mutex
+	var totalTokens int64
+	var totalCost float64
 
 	// Create errgroup with concurrency limit
 	g, ctx := errgroup.WithContext(ctx)
@@ -154,13 +149,10 @@ 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
+			totalCost += result.Cost
 			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 +165,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, 0, 0, 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, totalCost, nil
 }

pkg/rag/manager.go

@@ -236,7 +236,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
@@ -270,13 +270,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",
@@ -309,7 +309,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 					"num_results", len(results))
 			}
 
-			return results, nil
+			return results, usage, nil
 		}
 	}
 
@@ -321,6 +321,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
 	}
 
@@ -337,10 +338,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)
@@ -356,8 +358,10 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 				"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,
@@ -375,15 +379,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",
@@ -421,7 +425,7 @@ func (m *Manager) Query(ctx context.Context, query string) (results []database.S
 	// 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())

pkg/rag/strategy/bm25.go

@@ -271,21 +271,21 @@ func (s *BM25Strategy) Initialize(ctx context.Context, docPaths []string, chunki
 }
 
 // Query searches for relevant documents using BM25 scoring
-func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error) {
+func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) {
 	// Tokenize query
 	queryTerms := s.tokenize(query)
 	if len(queryTerms) == 0 {
-		return nil, errors.New("query contains no valid terms")
+		return nil, types.Usage{}, errors.New("query contains no valid terms")
 	}
 
 	// Get all documents
 	allDocs, err := s.getAllDocuments(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to retrieve documents: %w", err)
+		return nil, types.Usage{}, fmt.Errorf("failed to retrieve documents: %w", err)
 	}
 
 	if len(allDocs) == 0 {
-		return []database.SearchResult{}, nil
+		return []database.SearchResult{}, types.Usage{}, nil
 	}
 
 	// Pre-tokenize all documents once: build term frequency maps and lengths.
@@ -333,7 +333,7 @@ func (s *BM25Strategy) Query(ctx context.Context, query string, numResults int,
 		scores = scores[:numResults]
 	}
 
-	return scores, nil
+	return scores, types.Usage{}, nil
 }
 
 // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed

pkg/rag/strategy/types.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/docker/docker-agent/pkg/rag/database"
+	"github.com/docker/docker-agent/pkg/rag/types"
 )
 
 // Strategy defines the interface for different retrieval strategies.
@@ -14,7 +15,7 @@ type Strategy interface {
 
 	// Query searches for relevant documents using the strategy's retrieval method.
 	// numResults is the maximum number of candidates to retrieve (before fusion).
-	Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error)
+	Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error)
 
 	// CheckAndReindexChangedFiles checks for file changes and re-indexes if needed.
 	CheckAndReindexChangedFiles(ctx context.Context, docPaths []string, chunking ChunkingConfig) error

pkg/rag/strategy/vector_store.go

@@ -147,14 +147,6 @@ func NewVectorStore(cfg VectorStoreConfig) *VectorStore {
 		embeddingConcurrency:  cfg.EmbeddingConcurrency,
 		fileIndexConcurrency:  cfg.FileIndexConcurrency,
 	}
-
-	// Set usage handler to calculate cost from models.dev and emit events with CUMULATIVE totals
-	// This matches how chat completions calculate cost in runtime.go
-	cfg.Embedder.SetUsageHandler(func(ctx context.Context, tokens int64, _ float64) {
-		cost := s.calculateCost(ctx, tokens)
-		s.recordUsage(tokens, cost)
-	})
-
 	return s
 }
 
@@ -385,15 +377,18 @@ func (s *VectorStore) Initialize(ctx context.Context, docPaths []string, chunkin
 }
 
 // Query searches for relevant documents using vector similarity
-func (s *VectorStore) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, error) {
-	queryEmbedding, err := s.embedder.Embed(ctx, query)
+func (s *VectorStore) Query(ctx context.Context, query string, numResults int, threshold float64) ([]database.SearchResult, types.Usage, error) {
+	queryEmbedding, tokens, _, err := s.embedder.Embed(ctx, query)
 	if err != nil {
-		return nil, fmt.Errorf("failed to generate query embedding: %w", err)
+		return nil, types.Usage{}, fmt.Errorf("failed to generate query embedding: %w", err)
 	}
 
+	cost := s.calculateCost(ctx, tokens)
+	usage := types.Usage{TotalTokens: tokens, Cost: cost}
+
 	results, err := s.db.SearchSimilarVectors(ctx, queryEmbedding, numResults)
 	if err != nil {
-		return nil, fmt.Errorf("failed to search: %w", err)
+		return nil, types.Usage{}, fmt.Errorf("failed to search: %w", err)
 	}
 
 	// Convert internal result type to public SearchResult type
@@ -407,7 +402,7 @@ func (s *VectorStore) Query(ctx context.Context, query string, numResults int, t
 		}
 	}
 
-	return filtered, nil
+	return filtered, usage, nil
 }
 
 // CheckAndReindexChangedFiles checks for file changes and re-indexes if needed
@@ -607,11 +602,14 @@ func (s *VectorStore) indexFile(ctx context.Context, filePath string) error {
 		"path", filePath,
 		"chunk_count", len(validChunks))
 
-	embeddings, err := s.embedder.EmbedBatch(ctx, chunkContents)
+	embeddings, batchTokens, _, err := s.embedder.EmbedBatch(ctx, chunkContents)
 	if err != nil {
 		return fmt.Errorf("failed to generate embeddings: %w", classifyModelCallError(err))
 	}
 
+	batchCost := s.calculateCost(ctx, batchTokens)
+	s.recordUsage(batchTokens, batchCost)
+
 	if len(embeddings) != len(validChunks) {
 		return fmt.Errorf("embedding count mismatch: got %d embeddings for %d chunks", len(embeddings), len(validChunks))
 	}