// Package memory provides a worker for memory background tasks. package memory import ( "context" "log/slog" "sync" "time" ) // MemoryWorker is a background worker that handles memory operations. // It can be integrated with the agent's worker pool. type MemoryWorker struct { mu sync.Mutex // integration provides memory operations integration *AgentMemoryIntegration // stopChan signals the worker to stop stopChan chan struct{} // doneChan signals the worker has stopped doneChan chan struct{} // running indicates if the worker is active running bool // logger for worker events logger *slog.Logger // workQueue for async memory operations workQueue chan MemoryWorkItem // workerCount for parallel processing workerCount int } // MemoryWorkItem represents a unit of work for the memory worker type MemoryWorkItem struct { // Type of work Type WorkItemType // Context for the operation Ctx context.Context // Data for the operation Data interface{} // ResultChan for async results ResultChan chan<- WorkResult } // WorkItemType defines the type of memory work type WorkItemType int const ( WorkStoreEntry WorkItemType = iota WorkStoreBatch WorkConsolidate WorkApplyTTL WorkDeleteUserData ) // WorkResult contains the result of a work item type WorkResult struct { Success bool Error error Data interface{} } // NewMemoryWorker creates a new memory worker func NewMemoryWorker(integration *AgentMemoryIntegration, workerCount int) *MemoryWorker { if workerCount <= 0 { workerCount = 4 } return &MemoryWorker{ integration: integration, stopChan: make(chan struct{}), doneChan: make(chan struct{}), logger: slog.Default().With("component", "memory-worker"), workQueue: make(chan MemoryWorkItem, 1000), workerCount: workerCount, } } // Start starts the memory worker func (w *MemoryWorker) Start() { w.mu.Lock() if w.running { w.mu.Unlock() return } w.running = true w.mu.Unlock() // Start the integration (consolidation worker, etc.) w.integration.Start() // Start worker goroutines var wg sync.WaitGroup for i := 0; i < w.workerCount; i++ { wg.Add(1) go func(id int) { defer wg.Done() w.runWorker(id) }(i) } // Wait for all workers to finish go func() { wg.Wait() close(w.doneChan) }() w.logger.Info("memory worker started", "workers", w.workerCount) } // Stop stops the memory worker func (w *MemoryWorker) Stop() error { w.mu.Lock() if !w.running { w.mu.Unlock() return nil } w.running = false w.mu.Unlock() close(w.stopChan) <-w.doneChan // Stop the integration return w.integration.Stop() } // runWorker processes work items func (w *MemoryWorker) runWorker(id int) { w.logger.Debug("worker started", "id", id) for { select { case <-w.stopChan: w.logger.Debug("worker stopping", "id", id) return case item := <-w.workQueue: w.processWorkItem(item) } } } // processWorkItem handles a single work item func (w *MemoryWorker) processWorkItem(item MemoryWorkItem) { var result WorkResult switch item.Type { case WorkStoreEntry: if entry, ok := item.Data.(*MemoryEntry); ok { err := w.integration.orchestrator.Store(item.Ctx, entry) result = WorkResult{Success: err == nil, Error: err} } case WorkStoreBatch: if entries, ok := item.Data.([]*MemoryEntry); ok { err := w.integration.orchestrator.BatchStore(item.Ctx, entries) result = WorkResult{Success: err == nil, Error: err} } case WorkConsolidate: if sessionID, ok := item.Data.(string); ok { err := w.integration.ConsolidateSession(item.Ctx, sessionID) result = WorkResult{Success: err == nil, Error: err} } case WorkApplyTTL: deleted, err := w.integration.orchestrator.backend.ApplyTTL(item.Ctx) result = WorkResult{Success: err == nil, Error: err, Data: deleted} case WorkDeleteUserData: if req, ok := item.Data.(*DeleteUserDataRequest); ok { err := w.integration.DeleteUserData(item.Ctx, req.UserID, req.RequestedBy) result = WorkResult{Success: err == nil, Error: err} } } // Send result if channel provided if item.ResultChan != nil { select { case item.ResultChan <- result: default: // Channel full or closed } } } // DeleteUserDataRequest holds data for user deletion type DeleteUserDataRequest struct { UserID string RequestedBy string } // SubmitWork submits a work item to the queue func (w *MemoryWorker) SubmitWork(item MemoryWorkItem) bool { select { case w.workQueue <- item: return true default: return false // Queue full } } // SubmitWorkWithResult submits work and waits for result func (w *MemoryWorker) SubmitWorkWithResult(ctx context.Context, item MemoryWorkItem) (WorkResult, error) { resultChan := make(chan WorkResult, 1) item.ResultChan = resultChan item.Ctx = ctx if !w.SubmitWork(item) { return WorkResult{}, context.DeadlineExceeded } select { case <-ctx.Done(): return WorkResult{}, ctx.Err() case result := <-resultChan: return result, nil } } // AsyncStoreEntry submits an entry for async storage func (w *MemoryWorker) AsyncStoreEntry(ctx context.Context, entry *MemoryEntry) bool { return w.SubmitWork(MemoryWorkItem{ Type: WorkStoreEntry, Ctx: ctx, Data: entry, }) } // AsyncStoreBatch submits entries for async batch storage func (w *MemoryWorker) AsyncStoreBatch(ctx context.Context, entries []*MemoryEntry) bool { return w.SubmitWork(MemoryWorkItem{ Type: WorkStoreBatch, Ctx: ctx, Data: entries, }) } // AsyncConsolidate submits a session for async consolidation func (w *MemoryWorker) AsyncConsolidate(ctx context.Context, sessionID string) bool { return w.SubmitWork(MemoryWorkItem{ Type: WorkConsolidate, Ctx: ctx, Data: sessionID, }) } // AsyncDeleteUserData submits a user data deletion request func (w *MemoryWorker) AsyncDeleteUserData(ctx context.Context, userID, requestedBy string) bool { return w.SubmitWork(MemoryWorkItem{ Type: WorkDeleteUserData, Ctx: ctx, Data: &DeleteUserDataRequest{UserID: userID, RequestedBy: requestedBy}, }) } // QueueLength returns the current queue length func (w *MemoryWorker) QueueLength() int { return len(w.workQueue) } // IsRunning returns whether the worker is running func (w *MemoryWorker) IsRunning() bool { w.mu.Lock() defer w.mu.Unlock() return w.running } // WorkerPoolIntegration provides an interface compatible with common worker pools type WorkerPoolIntegration struct { worker *MemoryWorker } // NewWorkerPoolIntegration creates a new worker pool integration func NewWorkerPoolIntegration(worker *MemoryWorker) *WorkerPoolIntegration { return &WorkerPoolIntegration{worker: worker} } // Name returns the worker name func (w *WorkerPoolIntegration) Name() string { return "memory-worker" } // Start starts the worker func (w *WorkerPoolIntegration) Start(ctx context.Context) error { w.worker.Start() return nil } // Stop stops the worker func (w *WorkerPoolIntegration) Stop(ctx context.Context) error { return w.worker.Stop() } // Health returns the health status func (w *WorkerPoolIntegration) Health() error { if !w.worker.IsRunning() { return context.Canceled } return nil } // Metrics returns worker metrics func (w *WorkerPoolIntegration) Metrics() map[string]interface{} { metrics := w.worker.integration.GetMetrics() return map[string]interface{}{ "total_queries": metrics.TotalQueries, "total_stores": metrics.TotalStores, "avg_retrieval_ms": metrics.AvgRetrievalMs, "consolidation_runs": metrics.ConsolidationRuns, "queue_length": w.worker.QueueLength(), "active_working_memories": metrics.ActiveWorkingMemories, } } // HealthCheckTask performs periodic health checks type HealthCheckTask struct { integration *AgentMemoryIntegration interval time.Duration stopChan chan struct{} } // NewHealthCheckTask creates a new health check task func NewHealthCheckTask(integration *AgentMemoryIntegration, interval time.Duration) *HealthCheckTask { if interval <= 0 { interval = 30 * time.Second } return &HealthCheckTask{ integration: integration, interval: interval, stopChan: make(chan struct{}), } } // Start starts the health check task func (h *HealthCheckTask) Start() { go func() { ticker := time.NewTicker(h.interval) defer ticker.Stop() for { select { case <-h.stopChan: return case <-ticker.C: metrics := h.integration.GetMetrics() slog.Info("memory health check", "active_memories", metrics.ActiveWorkingMemories, "total_queries", metrics.TotalQueries, "avg_retrieval_ms", metrics.AvgRetrievalMs) } } }() } // Stop stops the health check task func (h *HealthCheckTask) Stop() { close(h.stopChan) }