refactor(workers): per-tenant provider resolution for background workers

Replace static provider/model fields with per-tenant resolution at
processing time. Fixes architectural mismatch where singleton workers
used tenant-specific system_configs.

Changes:
- Add shared ResolveBackgroundProvider() in providerresolve package
- Refactor vault enrichWorker, episodic, dreaming workers to resolve
  provider per-event using registry + systemConfigs
- Remove hot-reload machinery (no longer needed)
- Update ConsolidationDeps to use Registry/SystemConfigs

Fallback chain: background.provider → agent.default_provider → first
registered provider.
This commit is contained in:
viettranx
2026-04-12 16:19:55 +07:00
parent e021934e68
commit 3fa3fb5ebf
17 changed files with 259 additions and 191 deletions
+16 -15
View File
@@ -198,8 +198,8 @@ func runGateway() {
KGStore: pgStores.KnowledgeGraph,
SessionStore: pgStores.Sessions,
EventBus: domainBus,
Provider: bgProvider,
Model: bgModel,
SystemConfigs: pgStores.SystemConfigs,
Registry: providerRegistry,
Extractor: kgExtractor,
AgentStore: pgStores.Agents,
})
@@ -211,21 +211,22 @@ func runGateway() {
}
// V3: Wire vault enrichment worker (async summary + embedding + auto-linking).
// Provider is resolved per-tenant at runtime — no static provider needed.
var enrichProgress *vault.EnrichProgress
var updateVaultProvider vault.ProviderUpdater
if pgStores.Vault != nil && bgProvider != nil {
cleanupVaultEnrich, ep, updater := vault.RegisterEnrichWorker(vault.EnrichWorkerDeps{
VaultStore: pgStores.Vault,
Provider: bgProvider,
Model: bgModel,
EventBus: domainBus,
MsgBus: msgBus,
TeamStore: pgStores.Teams, // Phase 04 task-based auto-linking
var enrichWorker *vault.EnrichWorker
if pgStores.Vault != nil && providerRegistry != nil {
cleanupVaultEnrich, ep, ew := vault.RegisterEnrichWorker(vault.EnrichWorkerDeps{
VaultStore: pgStores.Vault,
SystemConfigs: pgStores.SystemConfigs,
Registry: providerRegistry,
EventBus: domainBus,
MsgBus: msgBus,
TeamStore: pgStores.Teams,
})
enrichProgress = ep
updateVaultProvider = updater
enrichWorker = ew
defer cleanupVaultEnrich()
slog.Info("vault enrichment worker registered", "provider", bgProvider.Name(), "model", bgModel)
slog.Info("vault enrichment worker registered (per-tenant provider resolution)")
}
loadBootstrapFiles(pgStores, workspace, agentCfg)
@@ -298,8 +299,8 @@ func runGateway() {
agentRouter: agentRouter,
toolsReg: toolsReg,
skillsLoader: skillsLoader,
enrichProgress: enrichProgress,
updateVaultProvider: updateVaultProvider,
enrichProgress: enrichProgress,
enrichWorker: enrichWorker,
workspace: workspace,
dataDir: dataDir,
domainBus: domainBus,
+2 -2
View File
@@ -28,8 +28,8 @@ type gatewayDeps struct {
toolsReg *tools.Registry
skillsLoader *skills.Loader // optional: enables skill creation in evolution approval
permCache *cache.PermissionCache // nil if no tenant store; closed on shutdown to stop sweep goroutines
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
updateVaultProvider vault.ProviderUpdater // nil if enrichment worker not registered; hot-swaps LLM
enrichProgress *vault.EnrichProgress // nil if enrichment worker not registered
enrichWorker *vault.EnrichWorker // nil if enrichment worker not registered; for stop/enqueue
workspace string
dataDir string
domainBus eventbus.DomainEventBus
+2 -6
View File
@@ -118,12 +118,8 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(
pgMem.UpdateChunkConfig(mem.MaxChunkLen, mem.ChunkOverlap)
}
}
// Hot-swap vault enrichment provider/model if config changed
if d.updateVaultProvider != nil {
if p, m := resolveBackgroundProvider(d.cfg, d.providerRegistry); p != nil {
d.updateVaultProvider(p, m)
}
}
// Note: vault enrichment provider is resolved per-tenant at runtime,
// no hot-reload needed here
slog.Debug("system_configs refreshed to in-memory config", "keys", len(sysConfigs))
}
})
+2 -15
View File
@@ -106,21 +106,8 @@ func (d *gatewayDeps) runLifecycle(
slog.Info("tts config reloaded", "provider", newMgr.PrimaryProvider(), "auto", string(newMgr.AutoMode()))
})
// Hot-swap vault enrichment provider/model on config changes via pub/sub.
if d.updateVaultProvider != nil {
d.msgBus.Subscribe("vault-enrich-config-reload", func(evt bus.Event) {
if evt.Name != bus.TopicConfigChanged {
return
}
updatedCfg, ok := evt.Payload.(*config.Config)
if !ok {
return
}
if p, m := resolveBackgroundProvider(updatedCfg, d.providerRegistry); p != nil {
d.updateVaultProvider(p, m)
}
})
}
// Note: vault enrichment provider is resolved per-tenant at runtime,
// no hot-reload handler needed here
// Log orphaned providers on agent deletion. Auto-delete is unsafe because
// providers can be referenced by heartbeats (FK), OAuth tokens, media chains.
@@ -109,8 +109,7 @@ func TestDreamingWorkerHandleHonoursCustomThreshold(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEpisodic,
memoryStore: mockMemory,
provider: mockProvider,
model: "test",
registry: testRegistry(mockProvider),
threshold: 5, // global default says skip at count=2
debounce: 1 * time.Second,
resolveConfig: func(_ context.Context, _ string) *config.DreamingConfig {
@@ -36,15 +36,14 @@ func TestDreamingWorkerUsesScoredListing(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEp,
memoryStore: mockMem,
provider: mockProv,
model: "test",
registry: testRegistry(mockProv),
threshold: 5,
debounce: 1 * time.Second,
}
err := worker.Handle(context.Background(), eventbus.DomainEvent{
Type: eventbus.EventEpisodicCreated,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: "agent-scored",
UserID: "user-scored",
Payload: &eventbus.EpisodicCreatedPayload{},
@@ -92,15 +91,14 @@ func TestDreamingWorkerFiltersBelowThreshold(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEp,
memoryStore: mockMem,
provider: mockProv,
model: "test",
registry: testRegistry(mockProv),
threshold: 5,
debounce: 1 * time.Second,
}
err := worker.Handle(context.Background(), eventbus.DomainEvent{
Type: eventbus.EventEpisodicCreated,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: "agent-filter",
UserID: "user-filter",
Payload: &eventbus.EpisodicCreatedPayload{},
@@ -142,15 +140,14 @@ func TestDreamingWorkerFilterEmptyStampsDebounce(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEp,
memoryStore: newMockMemoryStore(),
provider: &mockProvider{chatResp: &providers.ChatResponse{Content: "noop"}},
model: "test",
registry: testRegistry(&mockProvider{chatResp: &providers.ChatResponse{Content: "noop"}}),
threshold: 5,
debounce: 10 * time.Minute, // realistic
}
ev := eventbus.DomainEvent{
Type: eventbus.EventEpisodicCreated,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: "agent-loop",
UserID: "user-loop",
Payload: &eventbus.EpisodicCreatedPayload{},
+20 -6
View File
@@ -11,6 +11,7 @@ import (
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
@@ -26,8 +27,8 @@ const (
type dreamingWorker struct {
episodicStore store.EpisodicStore
memoryStore store.MemoryStore
provider providers.Provider
model string // LLM model for synthesis
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
// threshold/debounce are the global defaults. Per-agent overrides come
// from resolveConfig which reads the agent's MemoryConfig.Dreaming JSONB.
@@ -38,6 +39,11 @@ type dreamingWorker struct {
lastRun sync.Map // key: "agentID:userID" → time.Time
}
// resolveProvider delegates to shared background provider resolution.
func (w *dreamingWorker) resolveProvider(ctx context.Context, tenantID uuid.UUID) (providers.Provider, string) {
return providerresolve.ResolveBackgroundProvider(ctx, tenantID, w.registry, w.systemConfigs)
}
// formatEntryForSynthesis renders a single episodic entry with recall
// metadata for the LLM synthesis prompt. Entries with recall signal are
// tagged so the LLM can weight them higher; unrecalled entries pass through
@@ -143,8 +149,16 @@ func (w *dreamingWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
return nil
}
// Resolve provider for this tenant at processing time.
tenantUUID, _ := uuid.Parse(event.TenantID)
provider, model := w.resolveProvider(ctx, tenantUUID)
if provider == nil {
slog.Warn("dreaming: no provider available", "tenant", event.TenantID, "agent", agentID)
return nil
}
// Build LLM prompt and call provider.
synthesis, err := w.synthesize(ctx, entries)
synthesis, err := w.synthesize(ctx, provider, model, entries)
if err != nil {
slog.Warn("dreaming: LLM synthesis failed", "err", err, "agent", agentID)
return nil
@@ -181,19 +195,19 @@ func (w *dreamingWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
// synthesize calls the LLM to extract long-term facts from session summaries.
// Each entry is annotated with its recall metadata so the LLM can weight
// frequently-recalled memories higher during synthesis.
func (w *dreamingWorker) synthesize(ctx context.Context, entries []store.EpisodicSummary) (string, error) {
func (w *dreamingWorker) synthesize(ctx context.Context, provider providers.Provider, model string, entries []store.EpisodicSummary) (string, error) {
summaries := make([]string, len(entries))
for i, e := range entries {
summaries[i] = formatEntryForSynthesis(e)
}
body := strings.Join(summaries, "\n---\n")
resp, err := w.provider.Chat(ctx, providers.ChatRequest{
resp, err := provider.Chat(ctx, providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: dreamingSystemPrompt},
{Role: "user", Content: "Session summaries:\n---\n" + body + "\n---"},
},
Model: w.model,
Model: model,
Options: map[string]any{
providers.OptMaxTokens: dreamingMaxTokens,
},
+25 -15
View File
@@ -10,16 +10,22 @@ import (
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/eventbus"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// episodicWorker handles session.completed events → creates episodic summaries.
type episodicWorker struct {
store store.EpisodicStore
sessions store.SessionCoreStore // for reading session messages during summarization
provider providers.Provider
model string
eventBus eventbus.DomainEventBus
store store.EpisodicStore
sessions store.SessionCoreStore // for reading session messages during summarization
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
eventBus eventbus.DomainEventBus
}
// resolveProvider delegates to shared background provider resolution.
func (w *episodicWorker) resolveProvider(ctx context.Context, tenantID uuid.UUID) (providers.Provider, string) {
return providerresolve.ResolveBackgroundProvider(ctx, tenantID, w.registry, w.systemConfigs)
}
// Handle processes a session.completed event into an episodic summary.
@@ -57,15 +63,19 @@ func (w *episodicWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
// Use compaction summary if available, else call LLM
summary := payload.Summary
if summary == "" && w.provider != nil {
summary, err = w.summarizeSession(ctx, payload)
if err != nil {
return fmt.Errorf("episodic: summarize: %w", err)
if summary == "" {
provider, model := w.resolveProvider(ctx, tenantUUID)
if provider != nil {
summary, err = w.summarizeSession(ctx, provider, model, payload)
if err != nil {
return fmt.Errorf("episodic: summarize: %w", err)
}
}
}
if summary == "" {
provider, _ := w.resolveProvider(ctx, tenantUUID)
slog.Warn("episodic: no summary available, skipping", "session", payload.SessionKey,
"compaction_summary_empty", payload.Summary == "", "provider_nil", w.provider == nil)
"compaction_summary_empty", payload.Summary == "", "provider_nil", provider == nil)
return nil
}
slog.Debug("episodic: creating summary", "session", payload.SessionKey, "summary_len", len(summary))
@@ -113,12 +123,12 @@ func (w *episodicWorker) Handle(ctx context.Context, event eventbus.DomainEvent)
}
// summarizeSession reads actual session messages and calls LLM to summarize.
func (w *episodicWorker) summarizeSession(ctx context.Context, payload *eventbus.SessionCompletedPayload) (string, error) {
func (w *episodicWorker) summarizeSession(ctx context.Context, provider providers.Provider, model string, payload *eventbus.SessionCompletedPayload) (string, error) {
// Try reading session messages for a real summary.
if w.sessions != nil {
messages := w.sessions.GetHistory(ctx, payload.SessionKey)
if len(messages) > 0 {
return w.summarizeFromMessages(ctx, messages)
return w.summarizeFromMessages(ctx, provider, model, messages)
}
// Messages may have been compacted away — try existing session summary.
if summary := w.sessions.GetSummary(ctx, payload.SessionKey); summary != "" {
@@ -129,7 +139,7 @@ func (w *episodicWorker) summarizeSession(ctx context.Context, payload *eventbus
}
// summarizeFromMessages builds a conversation excerpt and calls LLM.
func (w *episodicWorker) summarizeFromMessages(ctx context.Context, messages []providers.Message) (string, error) {
func (w *episodicWorker) summarizeFromMessages(ctx context.Context, provider providers.Provider, model string, messages []providers.Message) (string, error) {
var sb strings.Builder
for _, m := range messages {
if m.Role == "system" {
@@ -153,12 +163,12 @@ func (w *episodicWorker) summarizeFromMessages(ctx context.Context, messages []p
sctx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
resp, err := w.provider.Chat(sctx, providers.ChatRequest{
resp, err := provider.Chat(sctx, providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: summarizationPrompt},
{Role: "user", Content: sb.String()},
},
Model: w.model,
Model: model,
Options: map[string]any{"max_tokens": 1024, "temperature": 0.3},
})
if err != nil {
+11
View File
@@ -3,10 +3,21 @@ package consolidation
import (
"context"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/knowledgegraph"
"github.com/nextlevelbuilder/goclaw/internal/providers"
)
// testRegistry creates a Registry with the given provider registered under MasterTenantID.
// This allows tests to use the new registry-based provider resolution.
func testRegistry(p providers.Provider) *providers.Registry {
r := providers.NewRegistry(func(ctx context.Context) uuid.UUID {
return providers.MasterTenantID
})
r.Register(p)
return r
}
// mockExtractor implements EntityExtractor for testing.
type mockExtractor struct {
result *knowledgegraph.ExtractionResult
+9 -9
View File
@@ -21,8 +21,8 @@ type ConsolidationDeps struct {
KGStore store.KnowledgeGraphStore
SessionStore store.SessionCoreStore // for reading session messages during summarization
EventBus eventbus.DomainEventBus
Provider providers.Provider // for LLM summarization
Model string
SystemConfigs store.SystemConfigStore // per-tenant provider config
Registry *providers.Registry // provider resolution
Extractor EntityExtractor
// AgentStore is optional: when present, the dreaming worker reads
// per-agent overrides from MemoryConfig.Dreaming. If nil, the worker
@@ -34,11 +34,11 @@ type ConsolidationDeps struct {
// Returns a cleanup function that unsubscribes all handlers.
func Register(deps ConsolidationDeps) func() {
episodic := &episodicWorker{
store: deps.EpisodicStore,
sessions: deps.SessionStore,
provider: deps.Provider,
model: deps.Model,
eventBus: deps.EventBus,
store: deps.EpisodicStore,
sessions: deps.SessionStore,
systemConfigs: deps.SystemConfigs,
registry: deps.Registry,
eventBus: deps.EventBus,
}
semantic := &semanticWorker{
kgStore: deps.KGStore,
@@ -52,8 +52,8 @@ func Register(deps ConsolidationDeps) func() {
dreaming := &dreamingWorker{
episodicStore: deps.EpisodicStore,
memoryStore: deps.MemoryStore,
provider: deps.Provider,
model: deps.Model,
systemConfigs: deps.SystemConfigs,
registry: deps.Registry,
threshold: dreamingDefaultThreshold,
debounce: dreamingDefaultDebounce,
resolveConfig: newAgentStoreResolver(deps.AgentStore),
+7 -11
View File
@@ -301,15 +301,14 @@ func TestEpisodicWorkerHandle_WithSummary(t *testing.T) {
worker := &episodicWorker{
store: mockStore,
provider: mockProvider,
model: "test-model",
registry: testRegistry(mockProvider),
eventBus: mockEventBus,
}
ctx := context.Background()
event := eventbus.DomainEvent{
Type: eventbus.EventSessionCompleted,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: uuid.New().String(),
UserID: "test-user",
Payload: &eventbus.SessionCompletedPayload{
@@ -695,8 +694,7 @@ func TestDreamingWorkerHandle_MeetsThreshold(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEpisodic,
memoryStore: mockMemory,
provider: mockProvider,
model: "test-model",
registry: testRegistry(mockProvider),
threshold: 5,
debounce: 1 * time.Second,
}
@@ -704,7 +702,7 @@ func TestDreamingWorkerHandle_MeetsThreshold(t *testing.T) {
ctx := context.Background()
event := eventbus.DomainEvent{
Type: eventbus.EventEpisodicCreated,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: "agent-123",
UserID: "user-123",
Payload: &eventbus.EpisodicCreatedPayload{},
@@ -752,8 +750,7 @@ func TestDreamingWorkerHandle_DebounceSkip(t *testing.T) {
worker := &dreamingWorker{
episodicStore: mockEpisodic,
memoryStore: mockMemory,
provider: mockProvider,
model: "test-model",
registry: testRegistry(mockProvider),
threshold: 5,
debounce: 10 * time.Second,
}
@@ -763,7 +760,7 @@ func TestDreamingWorkerHandle_DebounceSkip(t *testing.T) {
// First run should succeed
event1 := eventbus.DomainEvent{
Type: eventbus.EventEpisodicCreated,
TenantID: uuid.New().String(),
TenantID: providers.MasterTenantID.String(),
AgentID: "agent-123",
UserID: "user-123",
Payload: &eventbus.EpisodicCreatedPayload{
@@ -828,8 +825,7 @@ func TestRegister_WiresAllWorkers(t *testing.T) {
KGStore: mockKG,
SessionStore: mockSession,
EventBus: mockEventBus,
Provider: mockProvider,
Model: "test-model",
Registry: testRegistry(mockProvider),
Extractor: mockExtractor,
}
@@ -0,0 +1,64 @@
package providerresolve
import (
"context"
"github.com/google/uuid"
"github.com/nextlevelbuilder/goclaw/internal/providers"
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// ResolveBackgroundProvider resolves the LLM provider for background workers.
// Fallback chain: background.provider → agent.default_provider → first registered.
// Used by vault enrichment, episodic summarization, dreaming consolidation.
func ResolveBackgroundProvider(
ctx context.Context,
tenantID uuid.UUID,
registry *providers.Registry,
systemConfigs store.SystemConfigStore,
) (providers.Provider, string) {
if registry == nil {
return nil, ""
}
// Load system configs for the tenant
var configs map[string]string
if systemConfigs != nil {
tctx := store.WithTenantID(ctx, tenantID)
configs, _ = systemConfigs.List(tctx)
}
// tryResolve attempts to get a provider by name
tryResolve := func(name, model string) (providers.Provider, string, bool) {
if name == "" {
return nil, "", false
}
p, err := registry.GetForTenant(tenantID, name)
if err != nil || p == nil {
return nil, "", false
}
if model == "" {
model = p.DefaultModel()
}
return p, model, true
}
// 1. Explicit background config
if p, m, ok := tryResolve(configs["background.provider"], configs["background.model"]); ok {
return p, m
}
// 2. Agent default provider
if p, m, ok := tryResolve(configs["agent.default_provider"], configs["agent.default_model"]); ok {
return p, m
}
// 3. First registered provider
names := registry.ListForTenant(tenantID)
if len(names) == 0 {
return nil, ""
}
p, err := registry.GetForTenant(tenantID, names[0])
if err != nil {
return nil, ""
}
return p, p.DefaultModel()
}
+2 -2
View File
@@ -18,7 +18,7 @@ import (
//
// Nil-tolerant: teamStore may be unset (e.g. during tests), in which
// case the phase is a silent no-op.
func (w *enrichWorker) phase25TaskLinking(ctx context.Context, embedded []enriched, docMap map[string]*store.VaultDocument) {
func (w *EnrichWorker) phase25TaskLinking(ctx context.Context, embedded []enriched, docMap map[string]*store.VaultDocument) {
if w.teamStore == nil || len(embedded) == 0 {
return
}
@@ -102,7 +102,7 @@ func (w *enrichWorker) phase25TaskLinking(ctx context.Context, embedded []enrich
//
// Batched query via VaultStore.BatchFindByDelegationIDs; single
// CreateLinks call. No-op when no embedded doc carries metadata.delegation_id.
func (w *enrichWorker) phase26DelegationLinking(ctx context.Context, embedded []enriched, docMap map[string]*store.VaultDocument) {
func (w *EnrichWorker) phase26DelegationLinking(ctx context.Context, embedded []enriched, docMap map[string]*store.VaultDocument) {
if len(embedded) == 0 {
return
}
+6 -8
View File
@@ -36,8 +36,7 @@ type candidatePair struct {
}
// classifyLinks orchestrates LLM-based link classification for enriched docs.
func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID string, results []enriched) {
provider, _ := w.llm()
func (w *EnrichWorker) classifyLinks(ctx context.Context, provider providers.Provider, model, tenantID, agentID string, results []enriched) {
if provider == nil {
return
}
@@ -70,7 +69,7 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri
chunk := allCandidates[chunkStart:chunkEnd]
system, user := buildClassifyPrompt(source, chunk)
raw, err := w.callClassifyWithRetry(ctx, system, user)
raw, err := w.callClassifyWithRetry(ctx, provider, model, system, user)
if err != nil {
slog.Warn("vault.classify: llm_failed", "doc", sourceDocID, "chunk", chunkStart, "err", err)
continue
@@ -80,7 +79,7 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri
if err != nil {
slog.Warn("vault.classify: parse_failed_first", "doc", sourceDocID, "err", err, "raw_len", len(raw), "raw", raw)
hint := fmt.Sprintf("\n\nPrevious response was invalid JSON (error: %s). Output ONLY a valid JSON array.", err.Error())
raw2, err2 := w.callClassifyWithRetry(ctx, system, user+hint)
raw2, err2 := w.callClassifyWithRetry(ctx, provider, model, system, user+hint)
if err2 != nil {
slog.Warn("vault.classify: retry_parse_failed", "doc", sourceDocID, "err", err2)
continue
@@ -122,7 +121,7 @@ func (w *enrichWorker) classifyLinks(ctx context.Context, tenantID, agentID stri
}
}
func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, _ string, results []enriched) map[string][]candidatePair {
func (w *EnrichWorker) gatherCandidates(ctx context.Context, tenantID, _ string, results []enriched) map[string][]candidatePair {
seen := make(map[string]bool)
out := make(map[string][]candidatePair)
@@ -177,9 +176,8 @@ func (w *enrichWorker) gatherCandidates(ctx context.Context, tenantID, _ string,
}
// callClassifyWithRetry calls the LLM with shared retry logic.
func (w *enrichWorker) callClassifyWithRetry(ctx context.Context, system, user string) (string, error) {
_, model := w.llm()
return w.chatWithRetry(ctx, "vault.classify", providers.ChatRequest{
func (w *EnrichWorker) callClassifyWithRetry(ctx context.Context, provider providers.Provider, model, system, user string) (string, error) {
return w.chatWithRetry(ctx, provider, "vault.classify", providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: system},
{Role: "user", Content: user},
@@ -19,13 +19,10 @@ func TestCallClassifyWithRetry_ResponseWhitespaceStripping(t *testing.T) {
errors: []error{nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
resp, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry failed: %v", err)
@@ -55,13 +52,10 @@ func TestCallClassifyWithRetry_SecondAttemptSucceeds(t *testing.T) {
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
resp, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err != nil {
t.Fatalf("Should succeed on second attempt, got error: %v", err)
@@ -83,13 +77,10 @@ func TestCallClassifyWithRetry_EmptyResponse(t *testing.T) {
errors: []error{nil, nil, nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
resp, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
// Empty response is still a successful LLM call, should return empty string
if err != nil {
+8 -20
View File
@@ -52,13 +52,10 @@ func TestCallClassifyWithRetry_Success(t *testing.T) {
errors: []error{nil},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
resp, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry failed: %v", err)
@@ -89,13 +86,10 @@ func TestCallClassifyWithRetry_RetryThenSuccess(t *testing.T) {
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
resp, err := worker.callClassifyWithRetry(ctx, "system", "user")
resp, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err != nil {
t.Fatalf("callClassifyWithRetry should succeed after retries, got error: %v", err)
@@ -122,13 +116,10 @@ func TestCallClassifyWithRetry_AllFail(t *testing.T) {
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx := context.Background()
_, err := worker.callClassifyWithRetry(ctx, "system", "user")
_, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err == nil {
t.Fatalf("callClassifyWithRetry should return error after exhausting retries")
@@ -154,15 +145,12 @@ func TestCallClassifyWithRetry_ContextCancellation(t *testing.T) {
},
}
worker := &enrichWorker{
provider: provider,
model: "test",
}
worker := &EnrichWorker{}
ctx, cancel := context.WithCancel(context.Background())
cancel() // Cancel immediately
_, err := worker.callClassifyWithRetry(ctx, "system", "user")
_, err := worker.callClassifyWithRetry(ctx, provider, "test", "system", "user")
if err == nil {
t.Fatalf("callClassifyWithRetry should return error for cancelled context")
+72 -56
View File
@@ -12,9 +12,11 @@ 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"
"github.com/nextlevelbuilder/goclaw/internal/providerresolve"
"github.com/nextlevelbuilder/goclaw/internal/store"
"golang.org/x/sync/semaphore"
)
@@ -37,69 +39,78 @@ var (
// EnrichWorkerDeps bundles dependencies for the vault enrichment worker.
type EnrichWorkerDeps struct {
VaultStore store.VaultStore
Provider providers.Provider
Model string
EventBus eventbus.DomainEventBus
MsgBus bus.EventPublisher // for WS event broadcast
TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe)
VaultStore store.VaultStore
SystemConfigs store.SystemConfigStore // per-tenant provider config
Registry *providers.Registry // provider resolution
EventBus eventbus.DomainEventBus
MsgBus bus.EventPublisher // for WS event broadcast
TeamStore store.TaskCommentStore // for Phase 2.5 task-based auto-linking (nil-safe)
}
// ProviderUpdater allows hot-swapping the LLM provider/model at runtime.
type ProviderUpdater func(p providers.Provider, model string)
// RegisterEnrichWorker subscribes the enrichment worker to vault doc events.
// Returns (unsubscribe func, progress tracker, provider updater for hot-reload).
func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, ProviderUpdater) {
// Returns (unsubscribe func, progress tracker, EnrichWorker for stop/enqueue).
func RegisterEnrichWorker(deps EnrichWorkerDeps) (func(), *EnrichProgress, *EnrichWorker) {
progress := NewEnrichProgress(deps.MsgBus)
w := &enrichWorker{
vault: deps.VaultStore,
teamStore: deps.TeamStore,
provider: deps.Provider,
model: deps.Model,
dedup: make(map[string]string),
sem: semaphore.NewWeighted(enrichMaxConcurrent),
progress: progress,
w := &EnrichWorker{
vault: deps.VaultStore,
teamStore: deps.TeamStore,
systemConfigs: deps.SystemConfigs,
registry: deps.Registry,
msgBus: deps.MsgBus,
dedup: make(map[string]string),
sem: semaphore.NewWeighted(enrichMaxConcurrent),
progress: progress,
cancelFuncs: &sync.Map{},
}
unsub := deps.EventBus.Subscribe(eventbus.EventVaultDocUpserted, w.Handle)
return unsub, progress, w.UpdateProvider
return unsub, progress, w
}
// enrichWorker processes vault document upsert events to generate summaries,
// EnrichWorker processes vault document upsert events to generate summaries,
// embeddings, and semantic links between related documents.
type enrichWorker struct {
vault store.VaultStore
teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil
provider providers.Provider
model string
llmMu sync.RWMutex // guards provider + model hot-swap
queue enrichBatchQueue
progress *EnrichProgress
// Exported so HTTP handlers can call Stop/EnqueueUnenriched.
type EnrichWorker struct {
vault store.VaultStore
teamStore store.TaskCommentStore // nil-tolerant — Phase 2.5 disabled when nil
systemConfigs store.SystemConfigStore // per-tenant provider config
registry *providers.Registry // provider resolution
msgBus bus.EventPublisher // for error event broadcast
queue enrichBatchQueue
progress *EnrichProgress
// Bounded dedup: docID → content_hash. Prevents re-processing unchanged files.
dedupMu sync.Mutex
dedup map[string]string
sem *semaphore.Weighted // limits concurrent LLM summarize calls
// Per-tenant cancel functions for stop capability
cancelFuncs *sync.Map // key: tenantID string, value: context.CancelFunc
}
// UpdateProvider hot-swaps the LLM provider and model used by the enrichment worker.
// Called when background worker config changes at runtime.
func (w *enrichWorker) UpdateProvider(p providers.Provider, model string) {
w.llmMu.Lock()
defer w.llmMu.Unlock()
if w.provider != nil && w.provider.Name() == p.Name() && w.model == model {
return // no change
// resolveProviderForTenant delegates to shared background provider resolution.
func (w *EnrichWorker) resolveProviderForTenant(ctx context.Context, tenantID string) (providers.Provider, string) {
tid, err := uuid.Parse(tenantID)
if err != nil {
tid = providers.MasterTenantID
}
w.provider = p
w.model = model
slog.Info("vault.enrich: provider updated", "provider", p.Name(), "model", model)
return providerresolve.ResolveBackgroundProvider(ctx, tid, w.registry, w.systemConfigs)
}
// llm returns the current provider and model, safe for concurrent reads.
func (w *enrichWorker) llm() (providers.Provider, string) {
w.llmMu.RLock()
defer w.llmMu.RUnlock()
return w.provider, w.model
// Stop cancels in-flight enrichment for the given tenant.
// Safe to call even if no enrichment is running.
func (w *EnrichWorker) Stop(tenantID string) {
if cancel, ok := w.cancelFuncs.Load(tenantID); ok {
cancel.(context.CancelFunc)()
w.cancelFuncs.Delete(tenantID)
w.progress.Finish()
slog.Info("vault.enrich: stopped by user", "tenant", tenantID)
}
}
// IsRunning returns true if enrichment is in progress for the tenant.
func (w *EnrichWorker) IsRunning(tenantID string) bool {
_, ok := w.cancelFuncs.Load(tenantID)
return ok
}
// enrichTaskSiblingCap bounds the number of auto-linked siblings per
@@ -115,7 +126,7 @@ var enrichTaskSiblingCap = func() int {
}()
// Handle is the EventBus handler for vault.doc_upserted events.
func (w *enrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) error {
func (w *EnrichWorker) Handle(ctx context.Context, event eventbus.DomainEvent) error {
payload, ok := event.Payload.(eventbus.VaultDocUpsertedPayload)
if !ok {
return nil
@@ -151,7 +162,7 @@ type enriched struct {
// processBatch drains and processes queued vault doc events in a loop.
// 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) {
func (w *EnrichWorker) processBatch(ctx context.Context, key string) {
for {
items := w.queue.Drain(key)
if len(items) == 0 {
@@ -188,7 +199,7 @@ func (w *enrichWorker) processBatch(ctx context.Context, key string) {
// processChunk runs the 4-phase enrichment pipeline for a single chunk of docs.
// Phase 1 batches all files into a single LLM call for summarization.
func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultDocUpsertedPayload) {
func (w *EnrichWorker) processChunk(ctx context.Context, items []eventbus.VaultDocUpsertedPayload) {
// Phase 0 — Prepare: dedup check, batch-fetch existing docs, read file content.
type prepared struct {
payload eventbus.VaultDocUpsertedPayload
@@ -214,6 +225,13 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD
// Batch-fetch all existing docs in a single query.
tenantID := pending[0].TenantID
// Resolve provider once per chunk (all items share tenantID)
provider, model := w.resolveProviderForTenant(ctx, tenantID)
if provider == nil {
slog.Warn("vault.enrich: no provider available", "tenant", tenantID)
return
}
docIDs := make([]string, len(pending))
for i, item := range pending {
docIDs[i] = item.DocID
@@ -286,7 +304,7 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD
paths[i] = all[idx].payload.Path
contents[i] = all[idx].content
}
summaries := w.batchSummarize(ctx, paths, contents)
summaries := w.batchSummarize(ctx, provider, model, paths, contents)
for i, idx := range needLLM {
if i < len(summaries) && summaries[i] != "" {
all[idx].summary = summaries[i]
@@ -340,7 +358,7 @@ func (w *enrichWorker) processChunk(ctx context.Context, items []eventbus.VaultD
// Phase 3 — Classify links for this chunk.
if len(embedded) > 0 {
first := embedded[0].payload
w.classifyLinks(ctx, first.TenantID, first.AgentID, embedded)
w.classifyLinks(ctx, provider, model, first.TenantID, first.AgentID, embedded)
}
// Phase 4 — Record dedup + wikilinks.
@@ -355,14 +373,13 @@ Output a JSON array: [{"idx":1,"summary":"..."},{"idx":2,"summary":"..."}]
idx is 1-based matching the document number. Output ONLY valid JSON, no preamble.`
// batchSummarize sends multiple files in a single LLM call and parses JSON summaries.
func (w *enrichWorker) batchSummarize(ctx context.Context, paths, contents []string) []string {
func (w *EnrichWorker) batchSummarize(ctx context.Context, provider providers.Provider, model string, paths, contents []string) []string {
var b strings.Builder
for i := range paths {
fmt.Fprintf(&b, "[%d] File: %s\n%s\n\n", i+1, paths[i], contents[i])
}
_, model := w.llm()
raw, err := w.chatWithRetry(ctx, "vault.batch_summarize", providers.ChatRequest{
raw, err := w.chatWithRetry(ctx, provider, "vault.batch_summarize", providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: batchSummarizePrompt},
{Role: "user", Content: b.String()},
@@ -412,7 +429,7 @@ func parseBatchSummaries(raw string, expected int) []string {
// chatWithRetry is the shared retry loop for all enrichment LLM calls.
// Escalating timeouts and backoffs prevent transient provider failures
// (e.g. 529 overloaded) from permanently skipping documents.
func (w *enrichWorker) chatWithRetry(ctx context.Context, logPrefix string, req providers.ChatRequest) (string, error) {
func (w *EnrichWorker) chatWithRetry(ctx context.Context, provider providers.Provider, logPrefix string, req providers.ChatRequest) (string, error) {
var lastErr error
for attempt := range enrichMaxRetries {
if attempt > 0 {
@@ -422,7 +439,6 @@ func (w *enrichWorker) chatWithRetry(ctx context.Context, logPrefix string, req
case <-time.After(enrichRetryBackoffs[attempt]):
}
}
provider, _ := w.llm()
cctx, cancel := context.WithTimeout(ctx, enrichRetryTimeouts[attempt])
resp, err := provider.Chat(cctx, req)
cancel()
@@ -443,7 +459,7 @@ func (w *enrichWorker) chatWithRetry(ctx context.Context, logPrefix string, req
// Skips binary/media/document files to avoid parsing garbage data as wikilinks
// (PDFs and office docs are binary and would produce garbage [[...]] matches
// while wasting a 4MB read buffer).
func (w *enrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUpsertedPayload) {
func (w *EnrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUpsertedPayload) {
doc, err := w.vault.GetDocumentByID(ctx, p.TenantID, p.DocID)
if err != nil || doc == nil {
return
@@ -473,7 +489,7 @@ func (w *enrichWorker) syncWikilinks(ctx context.Context, p eventbus.VaultDocUps
// recordDedup stores a processed hash and evicts ~25% entries if over capacity.
func (w *enrichWorker) recordDedup(docID, hash string) {
func (w *EnrichWorker) recordDedup(docID, hash string) {
w.dedupMu.Lock()
defer w.dedupMu.Unlock()
w.dedup[docID] = hash