From 5b52ce27a82a91a3d558241e7d76b74c8d8f4371 Mon Sep 17 00:00:00 2001 From: viettranx Date: Sun, 12 Apr 2026 13:38:44 +0700 Subject: [PATCH] fix(vault): aggregate enrichment progress across per-agent batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After rescan agent_key fix, docs split into per-agent batch queues. Each batch was calling Start()/Finish() independently, causing the progress bar to flash 0/1 per agent instead of showing global progress. Now: handler Start(total) sets global count, worker batches call TrackBatch()/MarkBatchDone() for lifecycle, AddDone(n) auto-completes when done >= total. Progress bar shows smooth 0→N across all agents. --- internal/vault/enrich_progress.go | 59 ++++++++++++++++++++++++------- internal/vault/enrich_worker.go | 13 ++----- 2 files changed, 49 insertions(+), 23 deletions(-) diff --git a/internal/vault/enrich_progress.go b/internal/vault/enrich_progress.go index a8b4ec7b..00eab55b 100644 --- a/internal/vault/enrich_progress.go +++ b/internal/vault/enrich_progress.go @@ -9,13 +9,17 @@ import ( ) // EnrichProgress tracks enrichment pipeline progress and broadcasts via WS events. +// Lifecycle: handler calls Start(total) once with the global count, +// worker batches call AddDone(n) as docs complete, and MarkBatchDone() +// when a per-agent batch finishes. Auto-completes when done >= total. type EnrichProgress struct { - mu sync.Mutex - msgBus bus.EventPublisher - tenantID uuid.UUID - total int - done int - running bool + mu sync.Mutex + msgBus bus.EventPublisher + tenantID uuid.UUID + total int + done int + running bool + activeBatches int // number of worker batches in flight } // NewEnrichProgress creates a progress tracker that broadcasts to WS clients. @@ -49,33 +53,62 @@ func (p *EnrichProgress) broadcast(e EnrichEvent) { bus.BroadcastForTenant(p.msgBus, protocol.EventVaultEnrichProgress, p.tenantID, e) } -// Start signals enrichment with updated total and tenant scope. -// Only resets done counter when starting a new batch (was idle). -// Subsequent calls within the same batch just update the total. +// Start signals enrichment with the global total. Called by the HTTP rescan +// handler ONCE with the full count. Resets counters for a fresh run. func (p *EnrichProgress) Start(total int, tenantID uuid.UUID) { p.mu.Lock() defer p.mu.Unlock() - if !p.running { - p.done = 0 - } + p.done = 0 p.total = total p.tenantID = tenantID p.running = true - p.broadcast(EnrichEvent{Phase: "enriching", Done: p.done, Total: p.total, Running: true}) + p.activeBatches = 0 + p.broadcast(EnrichEvent{Phase: "enriching", Done: 0, Total: total, Running: true}) +} + +// TrackBatch increments the active batch counter. Called by worker when +// a per-agent batch starts processing. +func (p *EnrichProgress) TrackBatch() { + p.mu.Lock() + defer p.mu.Unlock() + p.activeBatches++ } // AddDone increments completed count by n and broadcasts progress. +// Auto-completes when done >= total. func (p *EnrichProgress) AddDone(n int) { p.mu.Lock() defer p.mu.Unlock() p.done += n + if p.done >= p.total && p.total > 0 { + p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false}) + p.running = false + return + } p.broadcast(EnrichEvent{Phase: "enriching", Done: p.done, Total: p.total, Running: true}) } +// MarkBatchDone decrements the active batch counter. If all batches +// finished and done < total (some docs failed), force-complete. +func (p *EnrichProgress) MarkBatchDone() { + p.mu.Lock() + defer p.mu.Unlock() + p.activeBatches-- + if p.activeBatches <= 0 && p.running { + p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false}) + p.running = false + p.activeBatches = 0 + } +} + // Finish signals the enrichment pipeline has completed and resets counters. +// Kept for backward compat but prefer AddDone auto-complete or MarkBatchDone. func (p *EnrichProgress) Finish() { p.mu.Lock() defer p.mu.Unlock() + if !p.running { + return // already completed via auto-complete + } p.broadcast(EnrichEvent{Phase: "complete", Done: p.done, Total: p.total, Running: false}) p.done = 0 p.total = 0 diff --git a/internal/vault/enrich_worker.go b/internal/vault/enrich_worker.go index 4acaade7..af62a0cc 100644 --- a/internal/vault/enrich_worker.go +++ b/internal/vault/enrich_worker.go @@ -12,7 +12,6 @@ import ( "sync" "time" - "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/bus" "github.com/nextlevelbuilder/goclaw/internal/eventbus" "github.com/nextlevelbuilder/goclaw/internal/providers" @@ -133,26 +132,20 @@ type enriched struct { // Items are chunked into enrichBatchSize groups so bulk rescan doesn't // overwhelm the LLM provider with hundreds of concurrent requests. func (w *enrichWorker) processBatch(ctx context.Context, key string) { + w.progress.TrackBatch() var totalQueued int for { items := w.queue.Drain(key) if len(items) == 0 { if w.queue.TryFinish(key) { - if totalQueued > 0 { - w.progress.Finish() - } + w.progress.MarkBatchDone() return } continue } totalQueued += len(items) - tenantID := uuid.Nil - if len(items) > 0 { - tenantID, _ = uuid.Parse(items[0].TenantID) - } - w.progress.Start(totalQueued, tenantID) // Process in chunks of enrichBatchSize, up to enrichMaxConcurrent in parallel. var wg sync.WaitGroup @@ -172,7 +165,7 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) { wg.Wait() if w.queue.TryFinish(key) { - w.progress.Finish() + w.progress.MarkBatchDone() return } }