spatial_memory.go

// Package yaad provides a graph-native persistent memory engine for AI coding
// agents. It stores, indexes, and retrieves contextual memories using a DAG with
// typed nodes and edges, a 4-signal hybrid search (BM25 + vector + graph +
// temporal) fused via reciprocal-rank fusion, and tiered spatial memory
// (hot/warm/cold) with promotion thresholds and token budgets.
//
// Spatial memory management (this file) divides entries into three access
// tiers, promotes entries based on recent hit counts, demotes idle entries,
// and enforces per-tier token budgets by evicting the least-recently accessed
// entries first. The system is safe for concurrent use.
package yaad

import (
	"sort"
	"sync"
	"time"
)

// TierType represents the access tier of a memory entry.
type TierType int

const (
	// TierHot is the fast path for frequently accessed memories.
	TierHot TierType = iota
	// TierWarm holds recently accessed memories.
	TierWarm
	// TierCold holds rarely accessed memories, ready for archival.
	TierCold
)

// Default tier promotion and demotion thresholds.
const (
	defaultHotAccessThreshold = 5                // accesses needed in window to promote Warm->Hot
	defaultHotAccessWindow    = 1 * time.Hour    // window for hot access count
	defaultHotDemotionAfter   = 30 * time.Minute // demote Hot->Warm after this idle period
	defaultWarmDemotionAfter  = 24 * time.Hour   // demote Warm->Cold after this idle period
)

// MemoryEntry is a single memory unit tracked by the spatial memory system.
type MemoryEntry struct {
	ID           string
	Content      string
	Tier         TierType
	AccessCount  int
	LastAccessed time.Time
	CreatedAt    time.Time
	TokenCount   int
}

// SpatialMemory manages a tiered set of memory entries with token budgets per tier.
// It is safe for concurrent use.
type SpatialMemory struct {
	mu              sync.RWMutex
	entries         map[string]*MemoryEntry
	hotBudget       int
	warmBudget      int
	coldBudget      int
	hotThreshold    int
	hotWindow       time.Duration
	hotDemoteAfter  time.Duration
	warmDemoteAfter time.Duration
}

// NewSpatialMemory creates a SpatialMemory with the given token budgets per tier.
func NewSpatialMemory(hotBudget, warmBudget, coldBudget int) *SpatialMemory {
	return &SpatialMemory{
		entries:         make(map[string]*MemoryEntry),
		hotBudget:       hotBudget,
		warmBudget:      warmBudget,
		coldBudget:      coldBudget,
		hotThreshold:    defaultHotAccessThreshold,
		hotWindow:       defaultHotAccessWindow,
		hotDemoteAfter:  defaultHotDemotionAfter,
		warmDemoteAfter: defaultWarmDemotionAfter,
	}
}

// Add inserts a new memory entry. If the entry already exists it is overwritten.
// New entries always land in TierCold; LastAccessed defaults to now and
// CreatedAt fills in if zero.
func (sm *SpatialMemory) Add(entry MemoryEntry) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	now := time.Now()
	entry.LastAccessed = now
	if entry.CreatedAt.IsZero() {
		entry.CreatedAt = now
	}
	if entry.CreatedAt.IsZero() {
		entry.CreatedAt = now
	}
	entry.Tier = TierCold
	entry.AccessCount = 0

	sm.entries[entry.ID] = &entry

	sm.enforceBudget(TierCold)
}

// Access records an access to the memory entry with the given id,
// updating its tier according to the promotion rules.
// Returns a copy of the entry, or false if the id was not found.
func (sm *SpatialMemory) Access(id string) (MemoryEntry, bool) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	e, ok := sm.entries[id]
	if !ok {
		return MemoryEntry{}, false
	}

	e.AccessCount++
	e.LastAccessed = time.Now()

	sm.promote(e)
	sm.enforceBudget(e.Tier)

	return *e, true
}

// Remove deletes the memory entry with the given id.
func (sm *SpatialMemory) Remove(id string) bool {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	_, ok := sm.entries[id]
	if ok {
		delete(sm.entries, id)
	}
	return ok
}

// TierStats returns the count of entries and total token usage for each tier.
// It snapshots entry pointers under the lock and computes stats outside it,
// reducing lock hold time for large entry sets.
func (sm *SpatialMemory) TierStats() map[TierType]struct{ Count, Tokens int } {
	sm.mu.RLock()
	snapshot := make([]*MemoryEntry, 0, len(sm.entries))
	for _, e := range sm.entries {
		snapshot = append(snapshot, e)
	}
	sm.mu.RUnlock()

	stats := make(map[TierType]struct{ Count, Tokens int })
	for _, e := range snapshot {
		s := stats[e.Tier]
		s.Count++
		s.Tokens += e.TokenCount
		stats[e.Tier] = s
	}
	return stats
}

// Compact demotes entries that violate idle thresholds and enforces
// token budgets by demoting the least-recently-accessed entries first.
func (sm *SpatialMemory) Compact() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	now := time.Now()

	// Demote idle entries.
	for _, e := range sm.entries {
		switch e.Tier {
		case TierHot:
			if now.Sub(e.LastAccessed) > sm.hotDemoteAfter {
				e.Tier = TierWarm
			}
		case TierWarm:
			if now.Sub(e.LastAccessed) > sm.warmDemoteAfter {
				e.Tier = TierCold
			}
		}
	}

	// Enforce budgets in order: hot, warm, cold.
	sm.enforceBudget(TierHot)
	sm.enforceBudget(TierWarm)
	sm.enforceBudget(TierCold)
}

// promote moves an entry to a higher tier based on recent access patterns.
// The caller must hold sm.mu.
func (sm *SpatialMemory) promote(e *MemoryEntry) {
	now := time.Now()

	switch e.Tier {
	case TierCold:
		e.Tier = TierWarm
	case TierWarm:
		// Count accesses within the hot window.
		recentAccesses := sm.recentAccessCount(e, now)
		if recentAccesses >= sm.hotThreshold {
			e.Tier = TierHot
		}
	case TierHot:
		// Already at the highest tier; nothing to do.
	}
}

// recentAccessCount estimates how many accesses occurred in the hot window.
// Because we only store a total AccessCount and a single LastAccessed time,
// we use a heuristic: if the entry was accessed within the window, we assume
// all accumulated accesses happened inside it (conservative for promotion).
func (sm *SpatialMemory) recentAccessCount(e *MemoryEntry, now time.Time) int {
	if now.Sub(e.LastAccessed) > sm.hotWindow {
		return 0
	}
	return e.AccessCount
}

// enforceBudget ensures the given tier does not exceed its token budget by
// demoting the least-recently-accessed entries to the next lower tier.
// The caller must hold sm.mu.
func (sm *SpatialMemory) enforceBudget(tier TierType) {
	budget := sm.budgetForTier(tier)
	if budget <= 0 {
		return
	}

	entries := sm.entriesInTier(tier)
	total := 0
	for _, e := range entries {
		total += e.TokenCount
	}

	if total <= budget {
		return
	}

	// Sort by LastAccessed ascending (oldest first) so we demote the least recent.
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].LastAccessed.Before(entries[j].LastAccessed)
	})

	for _, e := range entries {
		if total <= budget {
			break
		}
		if tier < TierCold {
			e.Tier = tier + 1
			total -= e.TokenCount
		}
	}
}

// budgetForTier returns the configured token budget for the given tier.
func (sm *SpatialMemory) budgetForTier(tier TierType) int {
	switch tier {
	case TierHot:
		return sm.hotBudget
	case TierWarm:
		return sm.warmBudget
	case TierCold:
		return sm.coldBudget
	default:
		return 0
	}
}

// entriesInTier returns a slice of all entries currently in the given tier.
// The caller must hold sm.mu.
func (sm *SpatialMemory) entriesInTier(tier TierType) []*MemoryEntry {
	var result []*MemoryEntry
	for _, e := range sm.entries {
		if e.Tier == tier {
			result = append(result, e)
		}
	}
	return result
}