feat(memory): add scheduled memory pruning with two-path strategy (#1350)

Implement MemoryBank-style retention scoring R=exp(-t/S) with two
complementary pruning paths:

- Path 1 (event-driven): async cap enforcement on Store() when user
  exceeds max_memories_per_user
- Path 2 (background sweep): periodic time.Ticker goroutine prunes
  decayed memories for inactive users in batches

Includes Prometheus metrics, graceful shutdown, multi-replica support
via prune_sweep_enabled flag, config template, and documentation.

Signed-off-by: Abdallah Samara <abdallahsamabd@gmail.com>

Author: abdallahsamabd

e2e/testing/09-memory-features-test.py

@@ -605,10 +605,10 @@ def test_02_pronoun_resolution(self):
             self.milvus.flush()
             time.sleep(3)
 
-        # Ask using pronoun in NEW SESSION (no previous_response_id)
-        # Memory context from Milvus should help resolve the pronoun
+        # Ask about the person by name in NEW SESSION (no previous_response_id)
+        # Memory context from Milvus should provide the answer
         result = self.send_memory_request(
-            message="Where does she live?",
+            message="Where does Sarah live?",
             auto_store=False,
             # NO previous_response_id - this is a new session!
         )

src/semantic-router/pkg/config/config.go

@@ -1080,6 +1080,21 @@ type MemoryQualityScoringConfig struct {
 
 	// MaxMemoriesPerUser caps memories per user; if over, lowest-R memories are deleted first (0 = no cap).
 	MaxMemoriesPerUser int `yaml:"max_memories_per_user,omitempty"`
+
+	// PruneInterval is the interval between background sweep runs (e.g. "6h", "24h").
+	// Empty or "0" disables the background sweep. Requires PruneSweepEnabled to be true.
+	PruneInterval string `yaml:"prune_interval,omitempty"`
+
+	// PruneBatchSize is the number of users to prune per batch during the background sweep (default: 50).
+	PruneBatchSize int `yaml:"prune_batch_size,omitempty"`
+
+	// PruneSweepEnabled enables the background prune sweep on this replica.
+	// In multi-replica deployments, set this to true on only one replica to avoid duplicate work.
+	PruneSweepEnabled bool `yaml:"prune_sweep_enabled,omitempty"`
+
+	// MaxConcurrentPrunes limits how many pruneIfOverCap goroutines can run at once
+	// to avoid overwhelming Milvus during high-throughput Store() bursts (default: 10).
+	MaxConcurrentPrunes int `yaml:"max_concurrent_prunes,omitempty"`
 }
 
 // MemoryMilvusConfig contains Milvus-specific configuration for memory storage.

src/semantic-router/pkg/extproc/router.go

@@ -54,6 +54,9 @@ type OpenAIRouter struct {
 	// RateLimiter enforces per-user/model rate limits from multiple sources
 	// (Envoy RLS → local limiter). Initialized in NewOpenAIRouter.
 	RateLimiter *ratelimit.RateLimitResolver
+
+	// StopPruneSweep stops the background memory prune sweep goroutine (nil if not started).
+	StopPruneSweep func()
 }
 
 // Ensure OpenAIRouter implements the ext_proc calls
@@ -397,6 +400,13 @@ func NewOpenAIRouter(configPath string) (*OpenAIRouter, error) {
 		}
 	}
 
+	// Start background prune sweep if memory store is available and configured
+	var stopPruneSweep func()
+	if memoryStore != nil && cfg.Memory.QualityScoring.PruneSweepEnabled && cfg.Memory.QualityScoring.PruneInterval != "" {
+		stopPruneSweep = memory.StartPruneSweep(cfg.Memory.QualityScoring, memoryStore)
+		logging.Infof("Memory prune sweep enabled (interval=%s)", cfg.Memory.QualityScoring.PruneInterval)
+	}
+
 	// Create memory extractor if memory_extraction external model is configured
 	var memoryExtractor *memory.MemoryExtractor
 	if memoryEnabled && cfg.FindExternalModelByRole(config.ModelRoleMemoryExtraction) != nil {
@@ -435,6 +445,7 @@ func NewOpenAIRouter(configPath string) (*OpenAIRouter, error) {
 		MemoryExtractor:      memoryExtractor,
 		CredentialResolver:   credResolver,
 		RateLimiter:          rateLimiter,
+		StopPruneSweep:       stopPruneSweep,
 	}
 
 	return router, nil

src/semantic-router/pkg/extproc/server.go

@@ -134,8 +134,13 @@ func (s *Server) Start() error {
 	return nil
 }
 
-// Stop stops the gRPC server
+// Stop stops the gRPC server and background goroutines
 func (s *Server) Stop() {
+	if s.service != nil {
+		if r := s.service.current.Load(); r != nil && r.StopPruneSweep != nil {
+			r.StopPruneSweep()
+		}
+	}
 	if s.server != nil {
 		s.server.GracefulStop()
 		logging.Infof("Server stopped")

src/semantic-router/pkg/memory/milvus_store.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/client"
@@ -22,6 +23,9 @@ const (
 	DefaultRetryBaseDelay = 100
 )
 
+// DefaultMaxConcurrentPrunes is the default limit for concurrent pruneIfOverCap goroutines.
+const DefaultMaxConcurrentPrunes = 10
+
 // MilvusStore provides memory retrieval from Milvus with similarity threshold filtering
 type MilvusStore struct {
 	client          client.Client
@@ -31,6 +35,8 @@ type MilvusStore struct {
 	maxRetries      int
 	retryBaseDelay  time.Duration
 	embeddingConfig EmbeddingConfig // Unified embedding configuration
+	pruneSem        chan struct{}   // bounds concurrent prune goroutines
+	pruneInFlight   sync.Map        // tracks userIDs with an active prune goroutine (dedup)
 }
 
 // MilvusStoreOptions contains configuration for creating a MilvusStore
@@ -79,6 +85,11 @@ func NewMilvusStore(options MilvusStoreOptions) (*MilvusStore, error) {
 		embeddingCfg = EmbeddingConfig{Model: EmbeddingModelBERT}
 	}
 
+	maxPrunes := cfg.QualityScoring.MaxConcurrentPrunes
+	if maxPrunes <= 0 {
+		maxPrunes = DefaultMaxConcurrentPrunes
+	}
+
 	store := &MilvusStore{
 		client:          options.Client,
 		collectionName:  options.CollectionName,
@@ -87,6 +98,7 @@ func NewMilvusStore(options MilvusStoreOptions) (*MilvusStore, error) {
 		maxRetries:      DefaultMaxRetries,
 		retryBaseDelay:  DefaultRetryBaseDelay * time.Millisecond,
 		embeddingConfig: embeddingCfg,
+		pruneSem:        make(chan struct{}, maxPrunes),
 	}
 
 	// Auto-create collection if it doesn't exist
@@ -684,9 +696,93 @@ func (m *MilvusStore) Store(ctx context.Context, memory *Memory) error {
 	}
 
 	logging.Debugf("MilvusStore.Store: successfully stored memory id=%s", memory.ID)
+
+	// Path 1: event-driven cap enforcement — async prune if user exceeds max_memories_per_user.
+	// Uses context.Background() intentionally: the goroutine must outlive the request ctx.
+	// Two layers of protection against Milvus pressure:
+	//   1. pruneInFlight (sync.Map): dedup — at most one goroutine per user at any time
+	//   2. pruneSem (channel): semaphore — at most maxConcurrentPrunes goroutines globally
+	if m.config.QualityScoring.MaxMemoriesPerUser > 0 {
+		if _, alreadyRunning := m.pruneInFlight.LoadOrStore(memory.UserID, struct{}{}); !alreadyRunning {
+			select {
+			case m.pruneSem <- struct{}{}:
+				go func(userID string) {
+					defer func() {
+						<-m.pruneSem
+						m.pruneInFlight.Delete(userID)
+					}()
+					m.pruneIfOverCap(context.Background(), userID)
+				}(memory.UserID)
+			default:
+				m.pruneInFlight.Delete(memory.UserID)
+				logging.Debugf("MilvusStore.Store: prune semaphore full, skipping cap check for user_id=%s", memory.UserID)
+			}
+		}
+	}
+
 	return nil
 }
 
+// pruneIfOverCap counts the user's memories and calls PruneUser if over MaxMemoriesPerUser.
+// Designed to run in a goroutine triggered by Store().
+func (m *MilvusStore) pruneIfOverCap(ctx context.Context, userID string) {
+	cap := m.config.QualityScoring.MaxMemoriesPerUser
+	if cap <= 0 {
+		return
+	}
+
+	count, err := m.countUserMemories(ctx, userID)
+	if err != nil {
+		logging.Warnf("MilvusStore.pruneIfOverCap: count failed for user_id=%s: %v", userID, err)
+		return
+	}
+
+	if count <= cap {
+		return
+	}
+
+	PruneCapTriggeredTotal.Inc()
+	logging.Infof("MilvusStore.pruneIfOverCap: user_id=%s has %d memories (cap=%d), pruning", userID, count, cap)
+
+	deleted, err := m.PruneUser(ctx, userID)
+	if err != nil {
+		logging.Warnf("MilvusStore.pruneIfOverCap: PruneUser failed for user_id=%s: %v", userID, err)
+		return
+	}
+	if deleted > 0 {
+		PruneDeletedTotal.WithLabelValues("cap").Add(float64(deleted))
+		logging.Infof("MilvusStore.pruneIfOverCap: user_id=%s pruned %d memories", userID, deleted)
+	}
+}
+
+// countUserMemories returns the number of memories stored for a given user.
+func (m *MilvusStore) countUserMemories(ctx context.Context, userID string) (int, error) {
+	filterExpr := fmt.Sprintf("user_id == \"%s\"", userID)
+
+	var queryResult []entity.Column
+	err := m.retryWithBackoff(ctx, func() error {
+		var retryErr error
+		queryResult, retryErr = m.client.Query(
+			ctx,
+			m.collectionName,
+			[]string{},
+			filterExpr,
+			[]string{"id"},
+		)
+		return retryErr
+	})
+	if err != nil {
+		return 0, fmt.Errorf("milvus query failed: %w", err)
+	}
+
+	for _, col := range queryResult {
+		if col.Name() == "id" {
+			return col.Len(), nil
+		}
+	}
+	return 0, nil
+}
+
 // upsert atomically replaces a row in Milvus by primary key.
 // The memory must be fully populated (including Embedding, timestamps, etc.).
 // Used by Update to avoid the delete+insert data-loss window.
@@ -1485,6 +1581,57 @@ func (m *MilvusStore) PruneUser(ctx context.Context, userID string) (deleted int
 	return deleted, nil
 }
 
+// ListStaleUserIDs queries Milvus for memories with created_at older than cutoffUnix
+// and returns the deduplicated set of user_id values. This targets users whose oldest
+// memories may have decayed below the prune threshold, without iterating all users.
+func (m *MilvusStore) ListStaleUserIDs(ctx context.Context, cutoffUnix int64) ([]string, error) {
+	if !m.enabled {
+		return nil, fmt.Errorf("milvus store is not enabled")
+	}
+
+	filterExpr := fmt.Sprintf("created_at < %d", cutoffUnix)
+	outputFields := []string{"user_id"}
+
+	var queryResult []entity.Column
+	err := m.retryWithBackoff(ctx, func() error {
+		var retryErr error
+		queryResult, retryErr = m.client.Query(
+			ctx,
+			m.collectionName,
+			[]string{},
+			filterExpr,
+			outputFields,
+		)
+		return retryErr
+	})
+	if err != nil {
+		return nil, fmt.Errorf("milvus query for stale users failed: %w", err)
+	}
+
+	seen := make(map[string]struct{})
+	for _, col := range queryResult {
+		if col.Name() == "user_id" {
+			vc, ok := col.(*entity.ColumnVarChar)
+			if !ok {
+				continue
+			}
+			for i := 0; i < vc.Len(); i++ {
+				uid, _ := vc.ValueByIdx(i)
+				if uid != "" {
+					seen[uid] = struct{}{}
+				}
+			}
+		}
+	}
+
+	userIDs := make([]string, 0, len(seen))
+	for uid := range seen {
+		userIDs = append(userIDs, uid)
+	}
+	logging.Debugf("MilvusStore.ListStaleUserIDs: found %d users with memories older than %d", len(userIDs), cutoffUnix)
+	return userIDs, nil
+}
+
 // isTransientError checks if an error is transient and should be retried
 func isTransientError(err error) bool {
 	if err == nil {

src/semantic-router/pkg/memory/prune_metrics.go

@@ -0,0 +1,58 @@
+package memory
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+var (
+	// PruneDeletedTotal tracks the total number of memories pruned, labeled by trigger source.
+	PruneDeletedTotal = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "memory_prune_deleted_total",
+			Help: "Total number of memories deleted by pruning",
+		},
+		[]string{"trigger"}, // "cap" or "sweep"
+	)
+
+	// PruneCapTriggeredTotal tracks how many times a Store() call triggered cap enforcement.
+	PruneCapTriggeredTotal = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Name: "memory_prune_cap_triggered_total",
+			Help: "Total number of times cap enforcement was triggered on Store()",
+		},
+	)
+
+	// PruneSweepRunsTotal tracks how many background sweep cycles have completed.
+	PruneSweepRunsTotal = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Name: "memory_prune_sweep_runs_total",
+			Help: "Total number of background prune sweep cycles completed",
+		},
+	)
+
+	// PruneSweepDuration tracks how long each background sweep cycle takes.
+	PruneSweepDuration = promauto.NewHistogram(
+		prometheus.HistogramOpts{
+			Name:    "memory_prune_sweep_duration_seconds",
+			Help:    "Duration of background prune sweep cycles in seconds",
+			Buckets: prometheus.DefBuckets,
+		},
+	)
+
+	// PruneSweepUsersProcessedTotal tracks how many users were evaluated during sweeps.
+	PruneSweepUsersProcessedTotal = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Name: "memory_prune_sweep_users_processed_total",
+			Help: "Total number of users processed during background prune sweeps",
+		},
+	)
+
+	// PruneSweepErrorsTotal tracks errors encountered during background sweeps.
+	PruneSweepErrorsTotal = promauto.NewCounter(
+		prometheus.CounterOpts{
+			Name: "memory_prune_sweep_errors_total",
+			Help: "Total number of errors encountered during background prune sweeps",
+		},
+	)
+)