fix(vault): aggregate enrichment progress across per-agent batches

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.
This commit is contained in:
viettranx committed 2026-04-12 13:38:44 +07:00
1 parent d58e0366fd
commit 5b52ce27a8
2 files changed
+49 -23

No files matched your search

+46 -13
View File
@@ -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
+3 -10
View File
@@ -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
}
}