LocalAI/pkg/model/watchdog_test.go

package model_test

import (
	"errors"
	"sync"
	"time"

	"github.com/mudler/LocalAI/pkg/model"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

// mockProcessManager implements ProcessManager for testing
type mockProcessManager struct {
	mu             sync.Mutex
	shutdownCalls  []string
	forceCalls     []string
	gracefulCalls  []string
	shutdownErrors map[string]error
}

func newMockProcessManager() *mockProcessManager {
	return &mockProcessManager{
		shutdownCalls:  []string{},
		shutdownErrors: make(map[string]error),
	}
}

func (m *mockProcessManager) ShutdownModel(modelName string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.shutdownCalls = append(m.shutdownCalls, modelName)
	m.gracefulCalls = append(m.gracefulCalls, modelName)
	if err, ok := m.shutdownErrors[modelName]; ok {
		return err
	}
	return nil
}

func (m *mockProcessManager) ShutdownModelForce(modelName string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.shutdownCalls = append(m.shutdownCalls, modelName)
	m.forceCalls = append(m.forceCalls, modelName)
	if err, ok := m.shutdownErrors[modelName]; ok {
		return err
	}
	return nil
}

// setShutdownError and clearShutdownError are safe to call while a running
// watchdog is shutting models down.
func (m *mockProcessManager) setShutdownError(modelName string, err error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.shutdownErrors[modelName] = err
}

func (m *mockProcessManager) clearShutdownError(modelName string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.shutdownErrors, modelName)
}

// residencyProcessManager also answers whether the loader still holds a model,
// like ModelLoader.HasModelAtAddress. It lets a test model a shutdown that
// returns an error after the loader already dropped the model.
type residencyProcessManager struct {
	*mockProcessManager
	resMu    sync.Mutex
	resident map[string]bool // keyed by model+"@"+address
}

func newResidencyProcessManager() *residencyProcessManager {
	return &residencyProcessManager{
		mockProcessManager: newMockProcessManager(),
		resident:           map[string]bool{},
	}
}

func (r *residencyProcessManager) setResident(modelName, address string, resident bool) {
	r.resMu.Lock()
	defer r.resMu.Unlock()
	r.resident[modelName+"@"+address] = resident
}

func (r *residencyProcessManager) HasModelAtAddress(modelName, address string) bool {
	r.resMu.Lock()
	defer r.resMu.Unlock()
	return r.resident[modelName+"@"+address]
}

func (m *mockProcessManager) getShutdownCalls() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	result := make([]string, len(m.shutdownCalls))
	copy(result, m.shutdownCalls)
	return result
}

func (m *mockProcessManager) getForceShutdownCalls() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	result := make([]string, len(m.forceCalls))
	copy(result, m.forceCalls)
	return result
}

func (m *mockProcessManager) getGracefulShutdownCalls() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	result := make([]string, len(m.gracefulCalls))
	copy(result, m.gracefulCalls)
	return result
}

var _ = Describe("WatchDog", func() {
	var (
		wd *model.WatchDog
		pm *mockProcessManager
	)

	BeforeEach(func() {
		pm = newMockProcessManager()
	})

	Context("LRU Limit", func() {
		It("should create watchdog with LRU limit", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(5*time.Minute),
				model.WithIdleTimeout(15*time.Minute),
				model.WithLRULimit(2),
			)
			Expect(wd.GetLRULimit()).To(Equal(2))
		})

		It("should allow updating LRU limit dynamically", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
			)
			wd.SetLRULimit(5)
			Expect(wd.GetLRULimit()).To(Equal(5))
		})

		It("should return 0 for disabled LRU", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(0),
			)
			Expect(wd.GetLRULimit()).To(Equal(0))
		})
	})

	Context("Memory Reclaimer Options", func() {
		It("should create watchdog with memory reclaimer settings", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithMemoryReclaimer(true, 0.85),
			)
			enabled, threshold := wd.GetMemoryReclaimerSettings()
			Expect(enabled).To(BeTrue())
			Expect(threshold).To(Equal(0.85))
		})

		It("should allow setting memory reclaimer via separate options", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithMemoryReclaimerEnabled(true),
				model.WithMemoryReclaimerThreshold(0.90),
			)
			enabled, threshold := wd.GetMemoryReclaimerSettings()
			Expect(enabled).To(BeTrue())
			Expect(threshold).To(Equal(0.90))
		})

		It("should use default threshold when not specified", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			_, threshold := wd.GetMemoryReclaimerSettings()
			Expect(threshold).To(Equal(model.DefaultMemoryReclaimerThreshold))
		})

		It("should allow updating memory reclaimer settings dynamically", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			wd.SetMemoryReclaimer(true, 0.80)
			enabled, threshold := wd.GetMemoryReclaimerSettings()
			Expect(enabled).To(BeTrue())
			Expect(threshold).To(Equal(0.80))
		})
	})

	Context("Model Tracking", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(5*time.Minute),
				model.WithIdleTimeout(15*time.Minute),
				model.WithLRULimit(3),
			)
		})

		It("should track loaded models count", func() {
			Expect(wd.GetLoadedModelCount()).To(Equal(0))

			wd.AddAddressModelMap("addr1", "model1")
			Expect(wd.GetLoadedModelCount()).To(Equal(1))

			wd.AddAddressModelMap("addr2", "model2")
			Expect(wd.GetLoadedModelCount()).To(Equal(2))
		})

		It("should update lastUsed time on Mark", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			// The model should now have a lastUsed time set
			// We can verify this indirectly through LRU eviction behavior
		})

		It("should update lastUsed time on UnMark", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)
			wd.UnMark("addr1")
			// The model should now have an updated lastUsed time
		})

		It("should update lastUsed time via UpdateLastUsed", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.UpdateLastUsed("addr1")
			// Verify the time was updated
		})

		It("remains busy until every parallel request completes", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			busySince := wd.GetState().BusyTime["addr1"]
			wd.Mark("addr1")
			wd.UnMark("addr1")

			state := wd.GetState()
			Expect(state.InFlight).To(HaveKeyWithValue("addr1", 1))
			Expect(state.BusyTime).To(HaveKeyWithValue("addr1", busySince))
			Expect(state.IdleTime).NotTo(HaveKey("addr1"))

			wd.UnMark("addr1")
			state = wd.GetState()
			Expect(state.InFlight).NotTo(HaveKey("addr1"))
			Expect(state.BusyTime).NotTo(HaveKey("addr1"))
			Expect(state.IdleTime).To(HaveKey("addr1"))
		})

		It("tracks the oldest active request instead of continuous parallel traffic", func() {
			wd.AddAddressModelMap("addr1", "model1")
			finishFirst := wd.TrackRequest("addr1")
			firstStarted := wd.GetState().BusyTime["addr1"]
			time.Sleep(10 * time.Millisecond)
			finishSecond := wd.TrackRequest("addr1")

			finishFirst()
			state := wd.GetState()
			Expect(state.InFlight).To(HaveKeyWithValue("addr1", 1))
			Expect(state.BusyTime["addr1"]).To(BeTemporally(">", firstStarted))

			finishSecond()
			state = wd.GetState()
			Expect(state.InFlight).NotTo(HaveKey("addr1"))
			Expect(state.BusyTime).NotTo(HaveKey("addr1"))
			Expect(state.IdleTime).To(HaveKey("addr1"))
		})
	})

	Context("EnforceLRULimit", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(5*time.Minute),
				model.WithIdleTimeout(15*time.Minute),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(true), // Enable force eviction for these tests to match old behavior
			)
		})

		It("should not evict when under limit", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle (not busy)

			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(0))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(BeEmpty())
		})

		It("should evict oldest model when at limit", func() {
			// Add two models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Unmark to make it idle

			// Enforce LRU with limit of 2 (need to make room for 1 new model)
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(ContainElement("model1")) // oldest should be evicted
		})

		It("should evict multiple models when needed", func() {
			// Add three models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Unmark to make it idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr3", "model3")
			wd.Mark("addr3")
			wd.UnMark("addr3") // Unmark to make it idle

			// Set limit to 1, should evict 2 oldest + 1 for new = 3 evictions
			wd.SetLRULimit(1)
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(3))
			Expect(result.NeedMore).To(BeFalse())
			shutdowns := pm.getShutdownCalls()
			Expect(shutdowns).To(ContainElement("model1"))
			Expect(shutdowns).To(ContainElement("model2"))
			Expect(shutdowns).To(ContainElement("model3"))
		})

		It("should account for pending loads", func() {
			// Add two models (at limit)
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Unmark to make it idle

			// With 1 pending load, we need to evict 2 (current=2, pending=1, new=1, limit=2)
			// total after = 2 + 1 + 1 = 4, need to evict 4 - 2 = 2
			result := wd.EnforceLRULimit(1)
			Expect(result.EvictedCount).To(Equal(2))
			Expect(result.NeedMore).To(BeFalse())
		})

		It("should not evict when LRU is disabled", func() {
			wd.SetLRULimit(0)

			wd.AddAddressModelMap("addr1", "model1")
			wd.AddAddressModelMap("addr2", "model2")
			wd.AddAddressModelMap("addr3", "model3")

			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(0))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(BeEmpty())
		})

		It("should evict least recently used first", func() {
			wd.SetLRULimit(2)

			// Add models with different lastUsed times
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle
			time.Sleep(20 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Unmark to make it idle
			time.Sleep(20 * time.Millisecond)

			// Touch model1 again to make it more recent
			wd.UpdateLastUsed("addr1")
			time.Sleep(20 * time.Millisecond)

			wd.AddAddressModelMap("addr3", "model3")
			wd.Mark("addr3")
			wd.UnMark("addr3") // Unmark to make it idle

			// Now model2 is the oldest, should be evicted first
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(BeNumerically(">=", 1))
			Expect(result.NeedMore).To(BeFalse())

			shutdowns := pm.getShutdownCalls()
			// model2 should be evicted first (it's the oldest)
			if len(shutdowns) >= 1 {
				Expect(shutdowns[0]).To(Equal("model2"))
			}
		})
	})

	Context("Single Backend Mode (LRU=1)", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(5*time.Minute),
				model.WithIdleTimeout(15*time.Minute),
				model.WithLRULimit(1),
				model.WithForceEvictionWhenBusy(true), // Enable force eviction for these tests
			)
		})

		It("should evict existing model when loading new one", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Unmark to make it idle

			// With limit=1, loading a new model should evict the existing one
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(ContainElement("model1"))
		})

		It("should handle rapid model switches", func() {
			for range 5 {
				wd.AddAddressModelMap("addr", "model")
				wd.Mark("addr")
				wd.UnMark("addr") // Unmark to make it idle
				wd.EnforceLRULimit(0)
			}
			// All previous models should have been evicted
			Expect(len(pm.getShutdownCalls())).To(Equal(5))
		})
	})

	Context("Force Eviction When Busy", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(false), // Default: skip eviction when busy
			)
		})

		It("should skip eviction for busy models when forceEvictionWhenBusy is false", func() {
			// Add two models (at limit of 2, need to evict 1 for new model)
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make model2 idle

			// Keep model1 as busy (simulating active API call)
			// model1 is already marked as busy from the first Mark call

			// Try to enforce LRU - should skip busy model1, evict model2
			result := wd.EnforceLRULimit(0)
			// Should evict model2 (not busy) but skip model1 (busy)
			// Since we evicted 1 (which is what we needed), NeedMore should be false
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeFalse()) // We evicted enough, even though we skipped model1
			Expect(pm.getShutdownCalls()).To(ContainElement("model2"))
			Expect(pm.getShutdownCalls()).ToNot(ContainElement("model1"))
		})

		It("should evict busy models when forceEvictionWhenBusy is true", func() {
			wd.SetForceEvictionWhenBusy(true)

			// Add two models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")

			// Keep model1 as busy (already marked from first Mark call)

			// Try to enforce LRU - should evict model1 even though busy
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(ContainElement("model1"))
		})

		It("should set NeedMore when all models are busy and forceEvictionWhenBusy is false", func() {
			// Add two models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")

			// Mark both as busy
			wd.Mark("addr1")
			wd.Mark("addr2")

			// Try to enforce LRU - should skip both busy models
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(0))
			Expect(result.NeedMore).To(BeTrue())
			Expect(pm.getShutdownCalls()).To(BeEmpty())
		})

		It("should allow updating forceEvictionWhenBusy dynamically", func() {
			// Start with false
			Expect(wd).ToNot(BeNil())

			// Add models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make model2 idle
			// Keep model1 busy (already marked)

			// With forceEvictionWhenBusy=false, should skip busy model1, evict model2
			result := wd.EnforceLRULimit(0)
			Expect(result.NeedMore).To(BeFalse())    // We evicted enough (1 model)
			Expect(result.EvictedCount).To(Equal(1)) // Should evict model2 (not busy)

			// Now enable force eviction
			wd.SetForceEvictionWhenBusy(true)

			// Add models again
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			// Keep model1 busy (already marked)

			// With forceEvictionWhenBusy=true, should evict busy model1
			result = wd.EnforceLRULimit(0)
			Expect(result.NeedMore).To(BeFalse())
			Expect(result.EvictedCount).To(Equal(1))
		})

		It("should continue to next LRU model when busy model is skipped", func() {
			// Add three models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make model2 idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr3", "model3")
			wd.Mark("addr3")
			wd.UnMark("addr3") // Make model3 idle

			// Keep model1 as busy (oldest, already marked)

			// Need to evict 2 models (limit=2, current=3, need room for 1 new)
			// Should skip model1 (busy), evict model2 and model3 (not busy)
			result := wd.EnforceLRULimit(0)
			// Should evict model2 and model3 (2 models, which is what we needed)
			Expect(result.EvictedCount).To(Equal(2))
			Expect(result.NeedMore).To(BeFalse()) // We evicted enough (2 models)
			Expect(pm.getShutdownCalls()).To(ContainElement("model2"))
			Expect(pm.getShutdownCalls()).To(ContainElement("model3"))
		})
	})

	Context("EnforceLRULimitResult", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(false),
			)
		})

		It("should return NeedMore=false when eviction is successful", func() {
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1") // Make idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make idle

			result := wd.EnforceLRULimit(0)
			Expect(result.NeedMore).To(BeFalse())
			Expect(result.EvictedCount).To(Equal(1))
		})

		It("should return NeedMore=true when not enough models can be evicted", func() {
			// Add two models (at limit of 2, need to evict 1 for new model)
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")

			// Mark both as busy (keep them busy)
			// Both are already marked as busy from the Mark calls above

			// Need to evict 1, but both are busy
			result := wd.EnforceLRULimit(0)
			Expect(result.NeedMore).To(BeTrue())
			Expect(result.EvictedCount).To(Equal(0))
		})

		It("should return NeedMore=true when need to evict multiple but some are busy", func() {
			// Set limit to 1, add 3 models (need to evict 2 for new model)
			wd.SetLRULimit(1)
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make model2 idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr3", "model3")
			wd.Mark("addr3")
			// Keep model1 and model3 busy

			// Need to evict 2 models, but model1 and model3 are busy, only model2 is idle
			// Should evict model2 (1 model), but NeedMore=true because we needed 2
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeTrue())
		})

		It("should return correct EvictedCount when some models are evicted", func() {
			// Add three models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2") // Make model2 idle
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr3", "model3")
			wd.Mark("addr3")
			wd.UnMark("addr3") // Make model3 idle

			// Keep model1 as busy (already marked)

			// Need to evict 2 models, but model1 is busy
			// Should evict model2 and model3 (2 models, which is what we needed)
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(2))
			Expect(result.NeedMore).To(BeFalse()) // We evicted enough (2 models)
		})
	})

	Context("Pinned Models", func() {
		It("should set and get pinned models", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			Expect(wd.IsModelPinned("model1")).To(BeFalse())

			wd.SetPinnedModels([]string{"model1", "model2"})
			Expect(wd.IsModelPinned("model1")).To(BeTrue())
			Expect(wd.IsModelPinned("model2")).To(BeTrue())
			Expect(wd.IsModelPinned("model3")).To(BeFalse())
		})

		It("should replace pinned models on subsequent calls", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			wd.SetPinnedModels([]string{"model1"})
			Expect(wd.IsModelPinned("model1")).To(BeTrue())

			wd.SetPinnedModels([]string{"model2"})
			Expect(wd.IsModelPinned("model1")).To(BeFalse())
			Expect(wd.IsModelPinned("model2")).To(BeTrue())
		})

		It("should skip pinned models during LRU eviction", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(true),
			)

			// Add two models, pin the older one
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			wd.SetPinnedModels([]string{"model1"})

			// Enforce LRU - model1 is oldest but pinned, model2 should be evicted
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(pm.getShutdownCalls()).To(ContainElement("model2"))
			Expect(pm.getShutdownCalls()).ToNot(ContainElement("model1"))
		})

		It("should not evict any model when all are pinned and LRU limit reached", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(1),
				model.WithForceEvictionWhenBusy(true),
			)

			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			wd.SetPinnedModels([]string{"model1", "model2"})

			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(0))
			Expect(pm.getShutdownCalls()).To(BeEmpty())
		})

		It("should skip pinned models during idle check", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithIdleTimeout(10*time.Millisecond),
				model.WithIdleCheck(true),
				model.WithWatchdogInterval(50*time.Millisecond),
			)

			// Add two models and make them idle
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			// Pin model1
			wd.SetPinnedModels([]string{"model1"})

			// Start watchdog and wait for idle check
			go wd.Run()
			defer wd.Shutdown()

			// Wait for the idle timeout + watchdog interval to pass
			time.Sleep(200 * time.Millisecond)

			// Only model2 should be shut down
			shutdowns := pm.getShutdownCalls()
			Expect(shutdowns).To(ContainElement("model2"))
			Expect(shutdowns).ToNot(ContainElement("model1"))
		})
	})

	Context("Busy Killer", func() {
		It("force-shuts down a model that is stuck busy past the busy timeout", func() {
			// Regression: a backend stuck on an in-flight gRPC call must be
			// killed via the force path (stop the process first), not the
			// graceful one, which intentionally waits for active calls until
			// its bounded deadline.
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(10*time.Millisecond),
				model.WithBusyCheck(true),
				model.WithWatchdogInterval(20*time.Millisecond),
			)

			wd.AddAddressModelMap("addr1", "stuckModel")
			wd.Mark("addr1") // busy — simulates an in-flight gRPC call

			go wd.Run()
			defer wd.Shutdown()

			Eventually(func() []string {
				return pm.getForceShutdownCalls()
			}, "300ms", "10ms").Should(ContainElement("stuckModel"))
			Expect(pm.getShutdownCalls()).To(ContainElement("stuckModel"))
			Expect(pm.getGracefulShutdownCalls()).To(BeEmpty())
		})

		It("keeps tracking a stuck model whose force shutdown failed", func() {
			pm.setShutdownError("stuckModel", errors.New("waiting to shut down model: context deadline exceeded"))
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(10*time.Millisecond),
				model.WithBusyCheck(true),
				model.WithWatchdogInterval(20*time.Millisecond),
			)

			wd.AddAddressModelMap("addr1", "stuckModel")
			wd.Mark("addr1")

			go wd.Run()
			defer wd.Shutdown()

			Eventually(func() []string {
				return pm.getForceShutdownCalls()
			}, "300ms", "10ms").Should(ContainElement("stuckModel"))
			// The backend still runs: the watchdog must not forget it.
			Consistently(func() map[string]string {
				return wd.GetState().AddressModelMap
			}, "100ms", "10ms").Should(HaveKeyWithValue("addr1", "stuckModel"))
		})
	})

	Context("Idle shutdown failure", func() {
		It("keeps tracking an idle model that could not be stopped, then stops it once it can", func() {
			pm.setShutdownError("idleModel", errors.New("waiting to shut down model: context deadline exceeded"))
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithIdleTimeout(10*time.Millisecond),
				model.WithIdleCheck(true),
				model.WithWatchdogInterval(20*time.Millisecond),
			)

			wd.AddAddressModelMap("addr1", "idleModel")
			wd.Mark("addr1")
			wd.UnMark("addr1")

			go wd.Run()
			defer wd.Shutdown()

			Eventually(func() []string {
				return pm.getShutdownCalls()
			}, "300ms", "10ms").Should(ContainElement("idleModel"))
			Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addr1", "idleModel"),
				"the backend still runs, the watchdog must still know it")

			pm.clearShutdownError("idleModel") // the lifecycle lock is released
			Eventually(func() map[string]string {
				return wd.GetState().AddressModelMap
			}, "500ms", "10ms").ShouldNot(HaveKey("addr1"))
		})
	})

	Context("Concurrency Groups", func() {
		Describe("ReplaceModelGroups / GetModelGroups", func() {
			It("returns nil for unknown models", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				Expect(wd.GetModelGroups("nope")).To(BeNil())
			})

			It("stores and retrieves groups", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy", "vision"},
					"b": {"heavy"},
				})
				Expect(wd.GetModelGroups("a")).To(Equal([]string{"heavy", "vision"}))
				Expect(wd.GetModelGroups("b")).To(Equal([]string{"heavy"}))
				Expect(wd.GetModelGroups("c")).To(BeNil())
			})

			It("replaces previous state on subsequent calls", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{"a": {"heavy"}})
				wd.ReplaceModelGroups(map[string][]string{"b": {"vision"}})
				Expect(wd.GetModelGroups("a")).To(BeNil())
				Expect(wd.GetModelGroups("b")).To(Equal([]string{"vision"}))
			})

			It("clears state when called with an empty map", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{"a": {"heavy"}})
				wd.ReplaceModelGroups(nil)
				Expect(wd.GetModelGroups("a")).To(BeNil())
			})

			It("returns a defensive copy", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{"a": {"heavy"}})
				got := wd.GetModelGroups("a")
				got[0] = "tampered"
				Expect(wd.GetModelGroups("a")).To(Equal([]string{"heavy"}))
			})
		})

		Describe("EnforceGroupExclusivity", func() {
			It("is a no-op when the requested model has no groups", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.AddAddressModelMap("addr1", "model1")
				wd.AddAddressModelMap("addr2", "model2")

				result := wd.EnforceGroupExclusivity("requested")
				Expect(result.EvictedCount).To(Equal(0))
				Expect(result.NeedMore).To(BeFalse())
				Expect(pm.getShutdownCalls()).To(BeEmpty())
			})

			It("is a no-op when no loaded model shares a group", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{
					"loaded":    {"vision"},
					"requested": {"heavy"},
				})
				wd.AddAddressModelMap("addr1", "loaded")

				result := wd.EnforceGroupExclusivity("requested")
				Expect(result.EvictedCount).To(Equal(0))
				Expect(result.NeedMore).To(BeFalse())
				Expect(pm.getShutdownCalls()).To(BeEmpty())
			})

			It("evicts a loaded model that shares a single group", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")

				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(1))
				Expect(result.NeedMore).To(BeFalse())
				Expect(pm.getShutdownCalls()).To(ConsistOf("a"))
			})

			It("evicts when groups overlap on any single name", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"x", "y"},
					"b": {"y", "z"},
				})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")

				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(1))
				Expect(pm.getShutdownCalls()).To(ConsistOf("a"))
			})

			It("evicts every conflicting loaded model", func() {
				wd = model.NewWatchDog(
					model.WithProcessManager(pm),
					model.WithForceEvictionWhenBusy(true),
				)
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
					"c": {"heavy"},
				})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")
				wd.AddAddressModelMap("addrB", "b")
				wd.Mark("addrB")
				wd.UnMark("addrB")

				result := wd.EnforceGroupExclusivity("c")
				Expect(result.EvictedCount).To(Equal(2))
				Expect(pm.getShutdownCalls()).To(ConsistOf("a", "b"))
			})

			It("skips a pinned conflicting model and reports NeedMore", func() {
				wd = model.NewWatchDog(
					model.WithProcessManager(pm),
					model.WithForceEvictionWhenBusy(true),
				)
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				wd.SetPinnedModels([]string{"a"})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")

				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(0))
				Expect(result.NeedMore).To(BeTrue())
				Expect(pm.getShutdownCalls()).To(BeEmpty())
			})

			It("skips a busy conflict when forceEvictionWhenBusy is false", func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA") // leave busy

				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(0))
				Expect(result.NeedMore).To(BeTrue())
				Expect(pm.getShutdownCalls()).To(BeEmpty())
			})

			It("evicts a busy conflict when forceEvictionWhenBusy is true", func() {
				wd = model.NewWatchDog(
					model.WithProcessManager(pm),
					model.WithForceEvictionWhenBusy(true),
				)
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA") // leave busy

				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(1))
				Expect(result.NeedMore).To(BeFalse())
				Expect(pm.getShutdownCalls()).To(ConsistOf("a"))
			})
		})

		Describe("a failed shutdown", func() {
			// Reproduces a production incident: the eviction of a model timed out
			// waiting for its lifecycle lock, the model kept running and kept its
			// VRAM, but it had already been dropped from the watchdog. The next
			// load saw no conflict and ran out of memory.
			BeforeEach(func() {
				wd = model.NewWatchDog(model.WithProcessManager(pm), model.WithLRULimit(1))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				pm.shutdownErrors["a"] = errors.New("waiting to shut down model: context deadline exceeded")
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")
			})

			It("keeps tracking a model a group eviction could not stop", func() {
				result := wd.EnforceGroupExclusivity("b")

				Expect(result.EvictedCount).To(Equal(0), "nothing was actually stopped")
				Expect(result.NeedMore).To(BeTrue(), "the conflict is unresolved, the caller must keep waiting")
				Expect(wd.GetLoadedModelCount()).To(Equal(1), "a is still running, it must stay tracked")
			})

			It("still sees the model as a conflict on the next attempt", func() {
				Expect(wd.EnforceGroupExclusivity("b").NeedMore).To(BeTrue())

				delete(pm.shutdownErrors, "a") // the lock is free this time
				result := wd.EnforceGroupExclusivity("b")
				Expect(result.EvictedCount).To(Equal(1))
				Expect(result.NeedMore).To(BeFalse())
				Expect(wd.GetLoadedModelCount()).To(Equal(0))
			})

			It("keeps tracking a model an LRU eviction could not stop", func() {
				wd.AddAddressModelMap("addrB", "b")
				wd.Mark("addrB")
				wd.UnMark("addrB")

				// Limit 1 with two models loaded: both must go to make room.
				// a refuses to stop, b stops normally.
				result := wd.EnforceLRULimit(0)

				Expect(result.EvictedCount).To(Equal(1), "only b was really stopped")
				Expect(result.NeedMore).To(BeTrue(), "a still holds its memory, the room does not exist yet")
				Expect(wd.GetLoadedModelCount()).To(Equal(1), "the model that would not stop is still tracked")
				Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addrA", "a"))
			})

			It("untracks a model that really was stopped", func() {
				delete(pm.shutdownErrors, "a")
				Expect(wd.EnforceGroupExclusivity("b").EvictedCount).To(Equal(1))
				Expect(wd.GetLoadedModelCount()).To(Equal(0))
			})

			It("treats a model that is already gone as stopped", func() {
				pm.shutdownErrors["a"] = model.ErrModelNotFound
				Expect(wd.EnforceGroupExclusivity("b").EvictedCount).To(Equal(1))
				Expect(wd.GetLoadedModelCount()).To(Equal(0))
			})
		})

		Describe("a shutdown error after the loader dropped the model", func() {
			// deleteProcess removes the store entry before it returns an error
			// for a failed remote unload (or the remote half of a local+remote
			// placement). Tracking such a model again would leave a phantom
			// entry that is counted as loaded and as a group conflict.
			var rpm *residencyProcessManager

			BeforeEach(func() {
				rpm = newResidencyProcessManager()
				wd = model.NewWatchDog(model.WithProcessManager(rpm), model.WithLRULimit(1))
				wd.ReplaceModelGroups(map[string][]string{
					"a": {"heavy"},
					"b": {"heavy"},
				})
				rpm.shutdownErrors["a"] = errors.New("unloading remote placements for model \"a\": node unreachable")
				wd.AddAddressModelMap("addrA", "a")
				wd.Mark("addrA")
				wd.UnMark("addrA")
			})

			It("does not track it again when the loader no longer holds it", func() {
				rpm.setResident("a", "addrA", false)

				result := wd.EnforceGroupExclusivity("b")

				Expect(result.NeedMore).To(BeTrue(), "the stop did report an error")
				Expect(wd.GetLoadedModelCount()).To(Equal(0), "the loader dropped a, no phantom entry may remain")
			})

			It("still tracks it again when the loader still holds it", func() {
				rpm.setResident("a", "addrA", true)

				result := wd.EnforceGroupExclusivity("b")

				Expect(result.EvictedCount).To(Equal(0))
				Expect(result.NeedMore).To(BeTrue())
				Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addrA", "a"),
					"a is still held by the loader and may still run")
			})
		})

	})

	Context("Functional Options", func() {
		It("should use default options when none provided", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			Expect(wd.GetLRULimit()).To(Equal(0))

			enabled, threshold := wd.GetMemoryReclaimerSettings()
			Expect(enabled).To(BeFalse())
			Expect(threshold).To(Equal(model.DefaultMemoryReclaimerThreshold))
		})

		It("should allow combining multiple options", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithBusyTimeout(10*time.Minute),
				model.WithIdleTimeout(30*time.Minute),
				model.WithBusyCheck(true),
				model.WithIdleCheck(true),
				model.WithLRULimit(5),
				model.WithMemoryReclaimerEnabled(true),
				model.WithMemoryReclaimerThreshold(0.80),
				model.WithForceEvictionWhenBusy(true),
			)

			Expect(wd.GetLRULimit()).To(Equal(5))

			enabled, threshold := wd.GetMemoryReclaimerSettings()
			Expect(enabled).To(BeTrue())
			Expect(threshold).To(Equal(0.80))
		})

		It("should use default forceEvictionWhenBusy (false) when not specified", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
			)
			// Default should be false - we can test this by checking behavior
			// Add a busy model and verify it's skipped
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.Mark("addr1") // Keep model1 busy

			wd.SetLRULimit(1)
			result := wd.EnforceLRULimit(0)
			// Should skip busy model1, evict model2, but NeedMore=true
			Expect(result.NeedMore).To(BeTrue())
		})

		It("should allow setting forceEvictionWhenBusy via option", func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(true),
			)

			// Add models
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			// Keep model1 busy (already marked from first Mark call)

			// Should evict busy model1
			result := wd.EnforceLRULimit(0)
			Expect(result.NeedMore).To(BeFalse())
			Expect(result.EvictedCount).To(Equal(1))
			Expect(pm.getShutdownCalls()).To(ContainElement("model1"))
		})
	})

	Context("Size-Aware Eviction", func() {
		BeforeEach(func() {
			wd = model.NewWatchDog(
				model.WithProcessManager(pm),
				model.WithLRULimit(2),
				model.WithForceEvictionWhenBusy(true),
				model.WithSizeAwareEviction(true),
			)
		})

		It("should enable size-aware eviction via option", func() {
			Expect(wd.GetSizeAwareEviction()).To(BeTrue())
		})

		It("should allow toggling size-aware eviction dynamically", func() {
			wd.SetSizeAwareEviction(false)
			Expect(wd.GetSizeAwareEviction()).To(BeFalse())
			wd.SetSizeAwareEviction(true)
			Expect(wd.GetSizeAwareEviction()).To(BeTrue())
		})

		It("should evict the largest model first when size-aware eviction is enabled", func() {
			// Register sizes: model1=100MB, model2=400MB
			wd.RegisterModelSize("model1", 100*1024*1024)
			wd.RegisterModelSize("model2", 400*1024*1024)

			// Add models — model1 older, model2 newer
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")
			time.Sleep(10 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			// With limit=2 and 2 loaded, adding a 3rd triggers eviction.
			// LRU order: model1 (oldest) would be evicted first.
			// Size order: model2 (400MB) should be evicted first.
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(result.NeedMore).To(BeFalse())
			Expect(pm.getShutdownCalls()).To(ContainElement("model2")) // largest first
			Expect(pm.getShutdownCalls()).ToNot(ContainElement("model1"))
		})

		It("should use LRU time as tiebreaker for equal-size models", func() {
			// Register equal sizes for both models
			wd.RegisterModelSize("model1", 200*1024*1024)
			wd.RegisterModelSize("model2", 200*1024*1024)

			// Add model1 first (older)
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")
			time.Sleep(20 * time.Millisecond)

			// Add model2 (newer)
			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			// Equal size → LRU tiebreaker: model1 (older) should be evicted
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(pm.getShutdownCalls()).To(ContainElement("model1"))
			Expect(pm.getShutdownCalls()).ToNot(ContainElement("model2"))
		})

		It("should fall back to LRU when no size is registered", func() {
			// No sizes registered — should behave like standard LRU
			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")
			time.Sleep(20 * time.Millisecond)

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			// Both have size 0 → LRU tiebreaker: model1 (older) evicted
			result := wd.EnforceLRULimit(0)
			Expect(result.EvictedCount).To(Equal(1))
			Expect(pm.getShutdownCalls()).To(ContainElement("model1"))
		})

		It("should clean up model size on eviction", func() {
			wd.RegisterModelSize("model1", 200*1024*1024)

			wd.AddAddressModelMap("addr1", "model1")
			wd.Mark("addr1")
			wd.UnMark("addr1")

			wd.AddAddressModelMap("addr2", "model2")
			wd.Mark("addr2")
			wd.UnMark("addr2")

			wd.EnforceLRULimit(0)

			// model1 was evicted; registering a new model with the same name
			// should start from a clean state (size not inherited)
			wd.RegisterModelSize("model1", 50*1024*1024)
			// Just verifying no panic and size can be re-registered
		})
	})
})