From 0c802e7083b5e811bf4e812262148ea0fb42e2c4 Mon Sep 17 00:00:00 2001 From: rcourtman Date: Sun, 1 Feb 2026 16:27:25 +0000 Subject: [PATCH] fix(patrol): improve service lifecycle, graceful shutdown, and concurrency --- internal/ai/investigation/orchestrator.go | 89 +++- .../orchestrator_concurrent_test.go | 423 ++++++++++++++++++ internal/ai/investigation_adapter.go | 10 + internal/ai/patrol.go | 4 + internal/ai/patrol_integration_test.go | 375 ++++++++++++++++ .../ai/patrol_lifecycle_integration_test.go | 285 ++++++++++++ internal/ai/patrol_run.go | 58 ++- internal/ai/service.go | 8 +- ..._handlers_investigation_additional_test.go | 8 + pkg/server/server.go | 3 + 10 files changed, 1247 insertions(+), 16 deletions(-) create mode 100644 internal/ai/investigation/orchestrator_concurrent_test.go create mode 100644 internal/ai/patrol_integration_test.go create mode 100644 internal/ai/patrol_lifecycle_integration_test.go diff --git a/internal/ai/investigation/orchestrator.go b/internal/ai/investigation/orchestrator.go index 5fbb97448..27c89a00d 100644 --- a/internal/ai/investigation/orchestrator.go +++ b/internal/ai/investigation/orchestrator.go @@ -168,12 +168,20 @@ type Orchestrator struct { // Fix verifier for post-fix verification fixVerifier FixVerifier + // License checker for defense-in-depth autonomy clamping + licenseChecker LicenseChecker + // Optional metrics callback for recording investigation outcomes metricsCallback MetricsCallback // Track running investigations runningCount int runningMu sync.Mutex + + // Shutdown coordination + shutdownCtx context.Context // Cancelled on Shutdown() to signal all investigations + shutdownFn context.CancelFunc // Triggers shutdownCtx cancellation + wg sync.WaitGroup // Tracks in-flight InvestigateFinding calls } // NewOrchestrator creates a new investigation orchestrator @@ -184,6 +192,7 @@ func NewOrchestrator( approvalStore ApprovalStore, config InvestigationConfig, ) *Orchestrator { + ctx, cancel := context.WithCancel(context.Background()) return &Orchestrator{ chatService: chatService, store: store, @@ -191,6 +200,8 @@ func NewOrchestrator( approvalStore: approvalStore, guardrails: NewGuardrails(), config: config, + shutdownCtx: ctx, + shutdownFn: cancel, } } @@ -239,6 +250,18 @@ func (o *Orchestrator) SetFixVerifier(fv FixVerifier) { o.fixVerifier = fv } +// LicenseChecker provides license feature checking for defense-in-depth validation +type LicenseChecker interface { + HasFeature(feature string) bool +} + +// SetLicenseChecker sets the license checker for defense-in-depth autonomy clamping +func (o *Orchestrator) SetLicenseChecker(checker LicenseChecker) { + o.mu.Lock() + defer o.mu.Unlock() + o.licenseChecker = checker +} + // GetConfig returns the current configuration func (o *Orchestrator) GetConfig() InvestigationConfig { o.mu.RLock() @@ -253,13 +276,58 @@ func (o *Orchestrator) CanStartInvestigation() bool { return o.runningCount < o.config.MaxConcurrent } +// Shutdown signals all running investigations to stop, persists state, +// and waits for them to finish up to the context deadline. +func (o *Orchestrator) Shutdown(ctx context.Context) error { + log.Info().Msg("Investigation orchestrator: starting shutdown") + + // Signal all running investigations to cancel + o.shutdownFn() + + // Force-save investigation state before waiting + if o.store != nil { + if err := o.store.ForceSave(); err != nil { + log.Error().Err(err).Msg("Investigation orchestrator: failed to force-save store during shutdown") + } + } + + // Wait for in-flight investigations to finish (or context to expire) + done := make(chan struct{}) + go func() { + o.wg.Wait() + close(done) + }() + + select { + case <-done: + log.Info().Msg("Investigation orchestrator: all investigations completed") + return nil + case <-ctx.Done(): + o.runningMu.Lock() + remaining := o.runningCount + o.runningMu.Unlock() + log.Warn().Int("remaining", remaining).Msg("Investigation orchestrator: shutdown timed out with running investigations") + return fmt.Errorf("shutdown timed out with %d investigations still running", remaining) + } +} + // InvestigateFinding starts an investigation for a finding func (o *Orchestrator) InvestigateFinding(ctx context.Context, finding *Finding, autonomyLevel string) error { + // Check if shutdown is in progress + select { + case <-o.shutdownCtx.Done(): + return fmt.Errorf("orchestrator is shutting down") + default: + } + // Check if we can start a new investigation if !o.CanStartInvestigation() { return fmt.Errorf("maximum concurrent investigations reached (%d)", o.config.MaxConcurrent) } + // Track this investigation for graceful shutdown + o.wg.Add(1) + // Increment running count o.runningMu.Lock() o.runningCount++ @@ -270,8 +338,17 @@ func (o *Orchestrator) InvestigateFinding(ctx context.Context, finding *Finding, o.runningMu.Lock() o.runningCount-- o.runningMu.Unlock() + o.wg.Done() }() + // Make ctx shutdown-aware: if the orchestrator's shutdownCtx is cancelled, + // ctx will also be cancelled, propagating to all downstream operations + // (session creation, LLM calls, tool execution). + ctx, ctxCancel := context.WithCancel(ctx) + defer ctxCancel() + stop := context.AfterFunc(o.shutdownCtx, func() { ctxCancel() }) + defer stop() + // Only enable autonomous mode if patrol autonomy is "full". // For all other levels, write commands (pulse_control, pulse_file_edit) // must go through the normal approval gate. The investigation AI should @@ -445,7 +522,9 @@ func formatOptional(label, value string) string { // executeWithLimits runs the investigation with turn and timeout limits func (o *Orchestrator) executeWithLimits(ctx context.Context, investigation *InvestigationSession, prompt string) error { - // Create context with timeout + // Create context with timeout. The passed-in ctx is already shutdown-aware + // (derived via context.AfterFunc in InvestigateFinding), so the timeout + // context inherits both the caller's deadline and the shutdown signal. timeoutCtx, cancel := context.WithTimeout(ctx, o.config.Timeout) defer cancel() @@ -584,6 +663,14 @@ func (o *Orchestrator) processResult(ctx context.Context, investigation *Investi } } + // Defense-in-depth: clamp autonomy if no auto-fix license + if o.licenseChecker != nil && !o.licenseChecker.HasFeature("ai_autofix") { + if currentLevel == "assisted" || currentLevel == "full" { + currentLevel = "approval" + log.Warn().Str("finding_id", finding.ID).Msg("Auto-fix requires Pro license - clamping to approval mode") + } + } + // Check if fix requires approval based on current autonomy level requiresApproval := o.guardrails.RequiresApproval( finding.Severity, diff --git a/internal/ai/investigation/orchestrator_concurrent_test.go b/internal/ai/investigation/orchestrator_concurrent_test.go new file mode 100644 index 000000000..4904efada --- /dev/null +++ b/internal/ai/investigation/orchestrator_concurrent_test.go @@ -0,0 +1,423 @@ +package investigation + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/rcourtman/pulse-go-rewrite/internal/ai/approval" +) + +// concurrentChatService is a chat service stub that supports concurrent usage +// with configurable delays and tracking. +type concurrentChatService struct { + mu sync.Mutex + sessionCounter int + executionDelay time.Duration + executeFunc func(ctx context.Context, cb StreamCallback) error +} + +func (s *concurrentChatService) CreateSession(ctx context.Context) (*Session, error) { + s.mu.Lock() + s.sessionCounter++ + id := fmt.Sprintf("session-%d", s.sessionCounter) + s.mu.Unlock() + return &Session{ID: id}, nil +} + +func (s *concurrentChatService) ExecuteStream(ctx context.Context, req ExecuteRequest, callback StreamCallback) error { + if s.executeFunc != nil { + return s.executeFunc(ctx, callback) + } + if s.executionDelay > 0 { + select { + case <-time.After(s.executionDelay): + case <-ctx.Done(): + return ctx.Err() + } + } + payload, _ := json.Marshal(map[string]string{"text": "CANNOT_FIX: simulated"}) + callback(StreamEvent{Type: "content", Data: payload}) + return nil +} + +func (s *concurrentChatService) GetMessages(ctx context.Context, sessionID string) ([]Message, error) { + return nil, nil +} + +func (s *concurrentChatService) DeleteSession(ctx context.Context, sessionID string) error { + return nil +} + +func (s *concurrentChatService) ListAvailableTools(ctx context.Context, prompt string) []string { + return nil +} + +func (s *concurrentChatService) SetAutonomousMode(enabled bool) {} + +// concurrentFindingsStore is a thread-safe findings store for concurrent tests. +type concurrentFindingsStore struct { + mu sync.RWMutex + findings map[string]*Finding +} + +func newConcurrentFindingsStore() *concurrentFindingsStore { + return &concurrentFindingsStore{ + findings: make(map[string]*Finding), + } +} + +func (s *concurrentFindingsStore) Add(f *Finding) { + s.mu.Lock() + defer s.mu.Unlock() + s.findings[f.ID] = f +} + +func (s *concurrentFindingsStore) Get(id string) *Finding { + s.mu.RLock() + defer s.mu.RUnlock() + return s.findings[id] +} + +func (s *concurrentFindingsStore) Update(f *Finding) bool { + s.mu.Lock() + defer s.mu.Unlock() + s.findings[f.ID] = f + return true +} + +// concurrentApprovalStore is a thread-safe stub for the ApprovalStore interface +// used in the investigation package (Create only). +type concurrentApprovalStore struct { + mu sync.Mutex + approvals []*Approval +} + +func (s *concurrentApprovalStore) Create(appr *Approval) error { + s.mu.Lock() + defer s.mu.Unlock() + s.approvals = append(s.approvals, appr) + return nil +} + +func (s *concurrentApprovalStore) count() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.approvals) +} + +// --------------------------------------------------------------------------- +// Test 1: TestConcurrent_MultipleInvestigations +// --------------------------------------------------------------------------- + +func TestConcurrent_MultipleInvestigations(t *testing.T) { + const N = 8 + const maxConcurrent = 3 + + var activeCount int64 // tracks currently executing investigations + var peakCount int64 + + chatService := &concurrentChatService{ + executeFunc: func(ctx context.Context, cb StreamCallback) error { + cur := atomic.AddInt64(&activeCount, 1) + // Track peak concurrency + for { + old := atomic.LoadInt64(&peakCount) + if cur <= old || atomic.CompareAndSwapInt64(&peakCount, old, cur) { + break + } + } + defer atomic.AddInt64(&activeCount, -1) + + // Simulate work + select { + case <-time.After(50 * time.Millisecond): + case <-ctx.Done(): + return ctx.Err() + } + + payload, _ := json.Marshal(map[string]string{"text": "CANNOT_FIX: simulated investigation"}) + cb(StreamEvent{Type: "content", Data: payload}) + return nil + }, + } + + store := NewStore("") + findings := newConcurrentFindingsStore() + + cfg := DefaultConfig() + cfg.MaxConcurrent = maxConcurrent + cfg.Timeout = 10 * time.Second + + orchestrator := NewOrchestrator(chatService, store, findings, nil, cfg) + + // Pre-create findings + for i := 0; i < N; i++ { + findings.Add(&Finding{ + ID: fmt.Sprintf("finding-%d", i), + Title: fmt.Sprintf("Test Finding %d", i), + Severity: "warning", + Category: "performance", + ResourceID: fmt.Sprintf("vm-%d", i), + ResourceName: fmt.Sprintf("test-vm-%d", i), + ResourceType: "vm", + Description: "Test finding for concurrency", + }) + } + + var wg sync.WaitGroup + errors := make([]error, N) + + for i := 0; i < N; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + finding := findings.Get(fmt.Sprintf("finding-%d", idx)) + errors[idx] = orchestrator.InvestigateFinding(context.Background(), finding, "approval") + }(i) + } + + wg.Wait() + + // Count successes vs max-concurrent rejections + succeeded := 0 + rejected := 0 + for _, err := range errors { + if err == nil { + succeeded++ + } else { + rejected++ + } + } + + t.Logf("Results: %d succeeded, %d rejected, peak concurrency: %d", succeeded, rejected, atomic.LoadInt64(&peakCount)) + + // At least maxConcurrent should succeed (the first batch gets through) + if succeeded < maxConcurrent { + t.Errorf("expected at least %d investigations to succeed, got %d", maxConcurrent, succeeded) + } + + // The running count should be back to zero after all complete + orchestrator.runningMu.Lock() + finalCount := orchestrator.runningCount + orchestrator.runningMu.Unlock() + if finalCount != 0 { + t.Errorf("expected running count to return to 0, got %d", finalCount) + } + + // Verify CanStartInvestigation returns true after all complete + if !orchestrator.CanStartInvestigation() { + t.Error("expected CanStartInvestigation to return true after all investigations complete") + } +} + +// --------------------------------------------------------------------------- +// Test 2: TestConcurrent_InvestigationDuringShutdown +// --------------------------------------------------------------------------- + +func TestConcurrent_InvestigationDuringShutdown(t *testing.T) { + investigationStarted := make(chan struct{}) + investigationBlocked := make(chan struct{}) + + chatService := &concurrentChatService{ + executeFunc: func(ctx context.Context, cb StreamCallback) error { + close(investigationStarted) + // Block until context is cancelled (simulating a long-running investigation) + <-investigationBlocked + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(5 * time.Second): + // Fallback so test doesn't hang forever + payload, _ := json.Marshal(map[string]string{"text": "CANNOT_FIX: timed out"}) + cb(StreamEvent{Type: "content", Data: payload}) + return nil + } + }, + } + + store := NewStore("") + findings := newConcurrentFindingsStore() + findings.Add(&Finding{ + ID: "finding-shutdown-1", + Title: "Shutdown Test Finding", + Severity: "warning", + Category: "performance", + ResourceID: "vm-1", + ResourceName: "test-vm", + ResourceType: "vm", + Description: "Finding for shutdown test", + }) + + cfg := DefaultConfig() + cfg.MaxConcurrent = 5 + cfg.Timeout = 30 * time.Second + + orchestrator := NewOrchestrator(chatService, store, findings, nil, cfg) + + // Start an investigation in the background + investigationDone := make(chan error, 1) + go func() { + finding := findings.Get("finding-shutdown-1") + investigationDone <- orchestrator.InvestigateFinding(context.Background(), finding, "approval") + }() + + // Wait for investigation to actually start executing + select { + case <-investigationStarted: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for investigation to start") + } + + // Now trigger shutdown with a reasonable timeout + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 3*time.Second) + defer shutdownCancel() + + // Unblock the investigation so it can observe the context cancellation + close(investigationBlocked) + + shutdownErr := orchestrator.Shutdown(shutdownCtx) + + // Shutdown should complete (not deadlock) + // It may return nil (all completed) or timeout error -- both are acceptable + // as long as we did not deadlock. + t.Logf("Shutdown returned: %v", shutdownErr) + + // The investigation should also finish + select { + case err := <-investigationDone: + t.Logf("Investigation finished with: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("investigation did not finish after shutdown -- possible deadlock") + } + + // After shutdown, new investigations should be rejected + finding := findings.Get("finding-shutdown-1") + err := orchestrator.InvestigateFinding(context.Background(), finding, "approval") + if err == nil { + t.Error("expected error when starting investigation after shutdown") + } +} + +// --------------------------------------------------------------------------- +// Test 3: TestConcurrent_ApprovalStoreConcurrency +// --------------------------------------------------------------------------- + +func TestConcurrent_ApprovalStoreConcurrency(t *testing.T) { + tmpDir := t.TempDir() + + store, err := approval.NewStore(approval.StoreConfig{ + DataDir: tmpDir, + DefaultTimeout: 5 * time.Minute, + MaxApprovals: 200, + DisablePersistence: true, + }) + if err != nil { + t.Fatalf("failed to create approval store: %v", err) + } + + const N = 20 + var wg sync.WaitGroup + + // Phase 1: Create N approval requests concurrently + ids := make([]string, N) + createErrs := make([]error, N) + + for i := 0; i < N; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + id := fmt.Sprintf("approval-%d", idx) + ids[idx] = id + createErrs[idx] = store.CreateApproval(&approval.ApprovalRequest{ + ID: id, + Command: fmt.Sprintf("echo test-%d", idx), + TargetType: "vm", + TargetID: fmt.Sprintf("vm-%d", idx), + TargetName: fmt.Sprintf("test-vm-%d", idx), + Context: "concurrent test", + }) + }(i) + } + wg.Wait() + + for i, err := range createErrs { + if err != nil { + t.Errorf("create approval %d failed: %v", i, err) + } + } + + // Verify all were created + pending := store.GetPendingApprovals() + if len(pending) != N { + t.Errorf("expected %d pending approvals, got %d", N, len(pending)) + } + + // Phase 2: Concurrently approve the first half and deny the second half + approveErrs := make([]error, N) + denyErrs := make([]error, N) + + for i := 0; i < N; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + if idx < N/2 { + _, approveErrs[idx] = store.Approve(ids[idx], "admin") + } else { + _, denyErrs[idx] = store.Deny(ids[idx], "admin", "test denial") + } + }(i) + } + wg.Wait() + + for i := 0; i < N/2; i++ { + if approveErrs[i] != nil { + t.Errorf("approve %d failed: %v", i, approveErrs[i]) + } + } + for i := N / 2; i < N; i++ { + if denyErrs[i] != nil { + t.Errorf("deny %d failed: %v", i, denyErrs[i]) + } + } + + // Phase 3: Verify final state concurrently + var verifyWg sync.WaitGroup + for i := 0; i < N; i++ { + verifyWg.Add(1) + go func(idx int) { + defer verifyWg.Done() + req, ok := store.GetApproval(ids[idx]) + if !ok { + t.Errorf("approval %d not found", idx) + return + } + if idx < N/2 { + if req.Status != approval.StatusApproved { + t.Errorf("approval %d: expected approved, got %s", idx, req.Status) + } + } else { + if req.Status != approval.StatusDenied { + t.Errorf("approval %d: expected denied, got %s", idx, req.Status) + } + } + }(i) + } + verifyWg.Wait() + + // Verify stats + stats := store.GetStats() + if stats["approved"] != N/2 { + t.Errorf("expected %d approved, got %d", N/2, stats["approved"]) + } + if stats["denied"] != N-N/2 { + t.Errorf("expected %d denied, got %d", N-N/2, stats["denied"]) + } + if stats["pending"] != 0 { + t.Errorf("expected 0 pending, got %d", stats["pending"]) + } +} diff --git a/internal/ai/investigation_adapter.go b/internal/ai/investigation_adapter.go index 28593630f..56b9ff97b 100644 --- a/internal/ai/investigation_adapter.go +++ b/internal/ai/investigation_adapter.go @@ -101,6 +101,11 @@ func (a *InvestigationOrchestratorAdapter) ReinvestigateFinding(ctx context.Cont return a.orchestrator.ReinvestigateFinding(ctx, findingID, autonomyLevel) } +// Shutdown delegates to the underlying orchestrator's Shutdown method. +func (a *InvestigationOrchestratorAdapter) Shutdown(ctx context.Context) error { + return a.orchestrator.Shutdown(ctx) +} + // SetFixVerifier registers a PatrolService as the fix verifier on the underlying orchestrator. func (a *InvestigationOrchestratorAdapter) SetFixVerifier(patrol *PatrolService) { a.orchestrator.SetFixVerifier(&patrolFixVerifier{patrol: patrol}) @@ -111,6 +116,11 @@ func (a *InvestigationOrchestratorAdapter) SetMetricsCallback() { a.orchestrator.SetMetricsCallback(&patrolMetricsCallback{}) } +// SetLicenseChecker wires license checking into the orchestrator for defense-in-depth. +func (a *InvestigationOrchestratorAdapter) SetLicenseChecker(checker investigation.LicenseChecker) { + a.orchestrator.SetLicenseChecker(checker) +} + // patrolMetricsCallback adapts PatrolMetrics to the investigation.MetricsCallback interface. type patrolMetricsCallback struct{} diff --git a/internal/ai/patrol.go b/internal/ai/patrol.go index 81d15d908..599b42c90 100644 --- a/internal/ai/patrol.go +++ b/internal/ai/patrol.go @@ -177,6 +177,9 @@ type InvestigationOrchestrator interface { CanStartInvestigation() bool // ReinvestigateFinding triggers a re-investigation of a finding ReinvestigateFinding(ctx context.Context, findingID, autonomyLevel string) error + // Shutdown signals all running investigations to stop, persists state, + // and waits for them to finish (up to the context deadline). + Shutdown(ctx context.Context) error } // InvestigationFinding is the finding type expected by the orchestrator @@ -268,6 +271,7 @@ type PatrolService struct { // Investigation orchestrator for autonomous investigation of findings investigationOrchestrator InvestigationOrchestrator + investigationWg sync.WaitGroup // Tracks in-flight investigation goroutines // Unified findings callback - pushes findings to unified store unifiedFindingCallback UnifiedFindingCallback diff --git a/internal/ai/patrol_integration_test.go b/internal/ai/patrol_integration_test.go new file mode 100644 index 000000000..860a3cce2 --- /dev/null +++ b/internal/ai/patrol_integration_test.go @@ -0,0 +1,375 @@ +package ai + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/rcourtman/pulse-go-rewrite/internal/ai/tools" + "github.com/rcourtman/pulse-go-rewrite/internal/config" + "github.com/rcourtman/pulse-go-rewrite/internal/models" +) + +// --- Integration test stubs --- + +// integrationOrchestrator records investigation triggers without running real LLM calls. +type integrationOrchestrator struct { + canStart bool + investigated []string // finding IDs passed to InvestigateFinding + investigateCh chan string + runningCount int32 +} + +func newIntegrationOrchestrator(canStart bool) *integrationOrchestrator { + return &integrationOrchestrator{ + canStart: canStart, + investigateCh: make(chan string, 10), + } +} + +func (o *integrationOrchestrator) InvestigateFinding(_ context.Context, f *InvestigationFinding, _ string) error { + o.investigated = append(o.investigated, f.ID) + o.investigateCh <- f.ID + return nil +} +func (o *integrationOrchestrator) GetInvestigationByFinding(_ string) *InvestigationSession { + return nil +} +func (o *integrationOrchestrator) GetRunningCount() int { + return int(atomic.LoadInt32(&o.runningCount)) +} +func (o *integrationOrchestrator) GetFixedCount() int { return 0 } +func (o *integrationOrchestrator) CanStartInvestigation() bool { + return o.canStart +} +func (o *integrationOrchestrator) ReinvestigateFinding(_ context.Context, _, _ string) error { + return nil +} +func (o *integrationOrchestrator) Shutdown(_ context.Context) error { + return nil +} + +// --- Integration tests --- + +// TestIntegration_PatrolRunCreatesFindings wires a real PatrolService with a mock +// chat service that reports a finding via patrol_report_finding. Verifies: +// - Finding is created in the FindingsStore +// - Investigation is triggered on the orchestrator +func TestIntegration_PatrolRunCreatesFindings(t *testing.T) { + // Setup AI service with mock provider + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + PatrolAutonomyLevel: "approval", + } + svc.provider = &mockProvider{} + + executor := tools.NewPulseToolExecutor(tools.ExecutorConfig{}) + mockCS := &patrolMockChatService{ + executor: executor, + executePatrolStreamFunc: func(ctx context.Context, req PatrolExecuteRequest, callback ChatStreamCallback) (*PatrolStreamResponse, error) { + creator := executor.GetPatrolFindingCreator() + if creator == nil { + return nil, fmt.Errorf("patrol finding creator not set") + } + // Report a finding for a resource with high CPU + _, _, err := creator.CreateFinding(tools.PatrolFindingInput{ + Key: "high-cpu", + Severity: "warning", + Category: "performance", + ResourceID: "vm-100", + ResourceName: "web-server", + ResourceType: "vm", + Title: "High CPU usage on web-server", + Description: "CPU is at 92%", + Evidence: "Current CPU: 92%", + }) + if err != nil { + return nil, fmt.Errorf("create finding: %w", err) + } + return &PatrolStreamResponse{Content: "Found high CPU on web-server"}, nil + }, + } + svc.SetChatService(mockCS) + + // State with a VM at 92% CPU + state := models.StateSnapshot{ + VMs: []models.VM{ + {ID: "vm-100", VMID: 100, Name: "web-server", Node: "pve-1", Status: "running", CPU: 0.92, Memory: models.Memory{Usage: 50}}, + }, + Nodes: []models.Node{ + {ID: "node/pve-1", Name: "pve-1", Status: "online", CPU: 0.40}, + }, + } + stateProvider := &patrolTestStateProvider{state: state} + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + Interval: 10 * time.Minute, + AnalyzeNodes: true, + AnalyzeGuests: true, + }) + + // Set up orchestrator to record investigation triggers + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Provide a mock config that returns "approval" autonomy level + mockAICfg := &config.AIConfig{ + Enabled: true, + PatrolAutonomyLevel: "approval", + } + svc.cfg = mockAICfg + + // Run patrol + ps.ForcePatrol(context.Background()) + + // Wait for the patrol run to complete and investigation goroutine to fire + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if ps.runHistoryStore.Count() > 0 { + break + } + time.Sleep(20 * time.Millisecond) + } + + // Verify finding was created + allFindings := ps.findings.GetActive(FindingSeverityInfo) + found := false + for _, f := range allFindings { + if f.ResourceID == "vm-100" && f.Key == "high-cpu" { + found = true + break + } + } + if !found { + t.Fatalf("expected finding for vm-100 with key 'high-cpu', got %d active findings", len(allFindings)) + } + + // Verify investigation was triggered (check with timeout due to async goroutine) + select { + case findingID := <-orch.investigateCh: + // Verify it's the right finding + if findingID == "" { + t.Fatal("investigation triggered with empty finding ID") + } + case <-time.After(3 * time.Second): + t.Fatal("expected investigation to be triggered within 3 seconds") + } +} + +// TestIntegration_FindingRejectedByThreshold verifies that findings for resources +// with metrics below the actionability threshold are rejected. +func TestIntegration_FindingRejectedByThreshold(t *testing.T) { + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + } + svc.provider = &mockProvider{} + + executor := tools.NewPulseToolExecutor(tools.ExecutorConfig{}) + var rejectedErr error + mockCS := &patrolMockChatService{ + executor: executor, + executePatrolStreamFunc: func(ctx context.Context, req PatrolExecuteRequest, callback ChatStreamCallback) (*PatrolStreamResponse, error) { + creator := executor.GetPatrolFindingCreator() + if creator == nil { + return nil, fmt.Errorf("patrol finding creator not set") + } + // Try to report finding for a resource with LOW CPU (30%) + _, _, err := creator.CreateFinding(tools.PatrolFindingInput{ + Key: "high-cpu", + Severity: "warning", + Category: "performance", + ResourceID: "vm-200", + ResourceName: "idle-server", + ResourceType: "vm", + Title: "High CPU usage on idle-server", + Description: "CPU is elevated", + Evidence: "Current CPU: 30%", + }) + rejectedErr = err + return &PatrolStreamResponse{Content: "Analysis complete"}, nil + }, + } + svc.SetChatService(mockCS) + + // State with a VM at only 30% CPU + state := models.StateSnapshot{ + VMs: []models.VM{ + {ID: "vm-200", VMID: 200, Name: "idle-server", Node: "pve-1", Status: "running", CPU: 0.30, Memory: models.Memory{Usage: 30}}, + }, + Nodes: []models.Node{ + {ID: "node/pve-1", Name: "pve-1", Status: "online"}, + }, + } + stateProvider := &patrolTestStateProvider{state: state} + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + AnalyzeGuests: true, + AnalyzeNodes: true, + }) + + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Run patrol + ps.ForcePatrol(context.Background()) + + // Wait for completion + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if ps.runHistoryStore.Count() > 0 { + break + } + time.Sleep(20 * time.Millisecond) + } + + // The finding should have been rejected + if rejectedErr == nil { + t.Fatal("expected finding to be rejected (CPU at 30% is below 50% threshold)") + } + + // No findings should be in the store for this resource + findings := ps.GetFindingsForResource("vm-200") + if len(findings) > 0 { + t.Fatalf("expected no findings for vm-200, got %d", len(findings)) + } + + // No investigation should have been triggered + select { + case id := <-orch.investigateCh: + t.Fatalf("expected no investigation to be triggered, but got finding %s", id) + case <-time.After(500 * time.Millisecond): + // Good - no investigation triggered + } +} + +// TestIntegration_StaleFindingReconciliation verifies that findings that are +// NOT re-reported by the LLM on a subsequent patrol run get auto-resolved. +func TestIntegration_StaleFindingReconciliation(t *testing.T) { + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + } + svc.provider = &mockProvider{} + + executor := tools.NewPulseToolExecutor(tools.ExecutorConfig{}) + mockCS := &patrolMockChatService{ + executor: executor, + executePatrolStreamFunc: func(ctx context.Context, req PatrolExecuteRequest, callback ChatStreamCallback) (*PatrolStreamResponse, error) { + // Deliberately report NO findings — the LLM "doesn't see" the issue anymore + // Also check existing findings to trigger the seeded-finding tracking + creator := executor.GetPatrolFindingCreator() + if creator != nil { + _ = creator.GetActiveFindings("", "warning") + } + return &PatrolStreamResponse{Content: "All looks healthy"}, nil + }, + } + svc.SetChatService(mockCS) + + state := models.StateSnapshot{ + VMs: []models.VM{ + {ID: "vm-300", VMID: 300, Name: "fixed-server", Node: "pve-1", Status: "running", CPU: 0.20, Memory: models.Memory{Usage: 30}}, + }, + Nodes: []models.Node{ + {ID: "node/pve-1", Name: "pve-1", Status: "online"}, + }, + } + stateProvider := &patrolTestStateProvider{state: state} + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + AnalyzeGuests: true, + AnalyzeNodes: true, + }) + + // Pre-seed a finding that was detected on a previous run + preSeedFinding := &Finding{ + ID: generateFindingID("vm-300", "performance", "high-cpu"), + Key: "high-cpu", + Severity: FindingSeverityWarning, + Category: FindingCategoryPerformance, + ResourceID: "vm-300", + ResourceName: "fixed-server", + ResourceType: "vm", + Title: "High CPU usage", + Description: "Was high before", + DetectedAt: time.Now().Add(-2 * time.Hour), + LastSeenAt: time.Now().Add(-1 * time.Hour), + TimesRaised: 3, + } + ps.findings.Add(preSeedFinding) + + // Verify finding is active before patrol + active := ps.findings.GetActive(FindingSeverityWarning) + if len(active) == 0 { + t.Fatal("expected pre-seeded finding to be active before patrol") + } + + // Run patrol (LLM doesn't re-report the finding) + ps.ForcePatrol(context.Background()) + + // Wait for completion + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if ps.runHistoryStore.Count() > 0 { + break + } + time.Sleep(20 * time.Millisecond) + } + + // The finding should have been auto-resolved by reconcileStaleFindings + stored := ps.findings.Get(preSeedFinding.ID) + if stored == nil { + t.Fatal("expected finding to still exist in store (resolved, not deleted)") + } + if !stored.IsResolved() { + t.Fatal("expected pre-seeded finding to be auto-resolved after patrol didn't re-report it") + } +} + +// TestIntegration_GracefulShutdownStopsClealy verifies that Stop() on a +// PatrolService returns cleanly and doesn't panic even with an orchestrator set. +func TestIntegration_GracefulShutdownStopsCleanly(t *testing.T) { + ps := NewPatrolService(nil, nil) + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Start and immediately stop + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + ps.SetConfig(PatrolConfig{Enabled: true, Interval: 1 * time.Hour}) + ps.Start(ctx) + + // Give it a moment to enter the patrol loop + time.Sleep(100 * time.Millisecond) + + // Stop should complete without hanging + done := make(chan struct{}) + go func() { + ps.Stop() + close(done) + }() + + select { + case <-done: + // Good - shutdown completed + case <-time.After(5 * time.Second): + t.Fatal("Stop() did not complete within 5 seconds") + } +} diff --git a/internal/ai/patrol_lifecycle_integration_test.go b/internal/ai/patrol_lifecycle_integration_test.go new file mode 100644 index 000000000..c2feed187 --- /dev/null +++ b/internal/ai/patrol_lifecycle_integration_test.go @@ -0,0 +1,285 @@ +package ai + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/rcourtman/pulse-go-rewrite/internal/ai/tools" + "github.com/rcourtman/pulse-go-rewrite/internal/config" + "github.com/rcourtman/pulse-go-rewrite/internal/models" +) + +// TestLifecycle_FindingToInvestigationToApproval verifies the end-to-end lifecycle: +// a patrol run creates a critical finding, which triggers investigation, and the +// finding's investigation status transitions from pending → running. +func TestLifecycle_FindingToInvestigationToApproval(t *testing.T) { + // Setup AI service with mock provider + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + PatrolAutonomyLevel: "approval", + } + svc.provider = &mockProvider{} + + executor := tools.NewPulseToolExecutor(tools.ExecutorConfig{}) + mockCS := &patrolMockChatService{ + executor: executor, + executePatrolStreamFunc: func(ctx context.Context, req PatrolExecuteRequest, callback ChatStreamCallback) (*PatrolStreamResponse, error) { + creator := executor.GetPatrolFindingCreator() + if creator == nil { + return nil, fmt.Errorf("patrol finding creator not set") + } + // Report a critical finding + _, _, err := creator.CreateFinding(tools.PatrolFindingInput{ + Key: "disk-full", + Severity: "critical", + Category: "storage", + ResourceID: "vm-500", + ResourceName: "db-server", + ResourceType: "vm", + Title: "Disk almost full on db-server", + Description: "Root partition is at 95%", + Evidence: "Disk usage: 95%", + }) + if err != nil { + return nil, fmt.Errorf("create finding: %w", err) + } + return &PatrolStreamResponse{Content: "Found disk full on db-server"}, nil + }, + } + svc.SetChatService(mockCS) + + // State with a VM that has high disk usage + state := models.StateSnapshot{ + VMs: []models.VM{ + {ID: "vm-500", VMID: 500, Name: "db-server", Node: "pve-1", Status: "running", CPU: 0.50, Memory: models.Memory{Usage: 60}, Disk: models.Disk{Usage: 95}}, + }, + Nodes: []models.Node{ + {ID: "node/pve-1", Name: "pve-1", Status: "online", CPU: 0.30}, + }, + } + stateProvider := &patrolTestStateProvider{state: state} + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + Interval: 10 * time.Minute, + AnalyzeNodes: true, + AnalyzeGuests: true, + }) + + // Set up orchestrator to record investigation triggers + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Run patrol + ps.ForcePatrol(context.Background()) + + // Wait for the patrol run to complete + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if ps.runHistoryStore.Count() > 0 { + break + } + time.Sleep(20 * time.Millisecond) + } + + // Verify finding was created + allFindings := ps.findings.GetActive(FindingSeverityCritical) + var createdFinding *Finding + for _, f := range allFindings { + if f.ResourceID == "vm-500" && f.Key == "disk-full" { + createdFinding = f + break + } + } + if createdFinding == nil { + t.Fatalf("expected finding for vm-500 with key 'disk-full', got %d active findings", len(allFindings)) + } + + // Verify the finding severity is critical + if createdFinding.Severity != FindingSeverityCritical { + t.Fatalf("expected finding severity 'critical', got %q", createdFinding.Severity) + } + + // Verify investigation was triggered (async goroutine) + select { + case findingID := <-orch.investigateCh: + if findingID == "" { + t.Fatal("investigation triggered with empty finding ID") + } + if findingID != createdFinding.ID { + t.Fatalf("investigation triggered for wrong finding: got %s, want %s", findingID, createdFinding.ID) + } + case <-time.After(3 * time.Second): + t.Fatal("expected investigation to be triggered within 3 seconds") + } + + // Verify the finding's investigation status transitioned (the orchestrator sets it to running + // when InvestigateFinding is called; in our mock it records the call but we can verify + // the finding was passed with correct state) + stored := ps.findings.Get(createdFinding.ID) + if stored == nil { + t.Fatal("finding not found in store after investigation trigger") + } +} + +// TestLifecycle_StuckInvestigationRecovery verifies that findings stuck in "running" +// status for too long are recovered to "failed" with outcome "timed_out". +func TestLifecycle_StuckInvestigationRecovery(t *testing.T) { + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + PatrolAutonomyLevel: "approval", + } + svc.provider = &mockProvider{} + + stateProvider := &patrolTestStateProvider{ + state: models.StateSnapshot{}, + } + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + Interval: 10 * time.Minute, + AnalyzeGuests: true, + }) + + // Set up orchestrator (required for MaybeInvestigateFinding path) + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Create a finding and manually set it to "running" with LastInvestigatedAt 20 minutes ago + twentyMinAgo := time.Now().Add(-20 * time.Minute) + stuckFinding := &Finding{ + ID: "stuck-finding-1", + Key: "high-mem", + Severity: FindingSeverityWarning, + Category: FindingCategoryPerformance, + ResourceID: "vm-600", + ResourceName: "app-server", + ResourceType: "vm", + Title: "High memory usage on app-server", + Description: "Memory at 90%", + Evidence: "Current memory: 90%", + InvestigationStatus: string(InvestigationStatusRunning), + InvestigationSessionID: "session-stuck-1", + LastInvestigatedAt: &twentyMinAgo, + InvestigationAttempts: 1, + DetectedAt: time.Now().Add(-30 * time.Minute), + LastSeenAt: time.Now().Add(-5 * time.Minute), + } + ps.findings.Add(stuckFinding) + + // Verify finding is active and in running state before recovery + stored := ps.findings.Get("stuck-finding-1") + if stored == nil { + t.Fatal("expected finding to be in store") + } + if stored.InvestigationStatus != string(InvestigationStatusRunning) { + t.Fatalf("expected investigation status 'running', got %q", stored.InvestigationStatus) + } + + // Run recovery + ps.recoverStuckInvestigations() + + // Verify the finding status is now "failed" with outcome "timed_out" + recovered := ps.findings.Get("stuck-finding-1") + if recovered == nil { + t.Fatal("expected finding to still exist after recovery") + } + if recovered.InvestigationStatus != string(InvestigationStatusFailed) { + t.Fatalf("expected investigation status 'failed' after recovery, got %q", recovered.InvestigationStatus) + } + if recovered.InvestigationOutcome != string(InvestigationOutcomeTimedOut) { + t.Fatalf("expected investigation outcome 'timed_out' after recovery, got %q", recovered.InvestigationOutcome) + } +} + +// TestLifecycle_AutonomyLevelRecheckOnRetry verifies that when autonomy is changed +// to "monitor" mode, retryTimedOutInvestigations() calls MaybeInvestigateFinding +// but the finding is NOT investigated because monitor mode blocks investigation. +func TestLifecycle_AutonomyLevelRecheckOnRetry(t *testing.T) { + persistence := config.NewConfigPersistence(t.TempDir()) + svc := NewService(persistence, nil) + // Start with "full" autonomy + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + PatrolAutonomyLevel: "full", + } + svc.provider = &mockProvider{} + + stateProvider := &patrolTestStateProvider{ + state: models.StateSnapshot{}, + } + + ps := NewPatrolService(svc, stateProvider) + ps.SetConfig(PatrolConfig{ + Enabled: true, + Interval: 10 * time.Minute, + AnalyzeGuests: true, + }) + + // Set up orchestrator to track investigation calls + orch := newIntegrationOrchestrator(true) + ps.SetInvestigationOrchestrator(orch) + + // Create a finding with InvestigationStatus="failed" and InvestigationOutcome="timed_out" + timedOutFinding := &Finding{ + ID: "timed-out-finding-1", + Key: "cpu-spike", + Severity: FindingSeverityCritical, + Category: FindingCategoryPerformance, + ResourceID: "vm-700", + ResourceName: "compute-server", + ResourceType: "vm", + Title: "CPU spike on compute-server", + Description: "CPU spiked to 99%", + Evidence: "Current CPU: 99%", + InvestigationStatus: string(InvestigationStatusFailed), + InvestigationOutcome: string(InvestigationOutcomeTimedOut), + InvestigationAttempts: 1, + InvestigationSessionID: "session-timeout-1", + DetectedAt: time.Now().Add(-1 * time.Hour), + LastSeenAt: time.Now().Add(-5 * time.Minute), + } + ps.findings.Add(timedOutFinding) + + // Now change autonomy to "monitor" (which blocks all investigation) + svc.cfg = &config.AIConfig{ + Enabled: true, + PatrolModel: "mock:model", + PatrolAutonomyLevel: "monitor", + } + + // Call retryTimedOutInvestigations — this will call MaybeInvestigateFinding + // but ShouldInvestigate should return false because autonomy is "monitor" + ps.retryTimedOutInvestigations() + + // Verify that no investigation was triggered on the orchestrator + select { + case id := <-orch.investigateCh: + t.Fatalf("expected no investigation to be triggered in monitor mode, but got finding %s", id) + case <-time.After(500 * time.Millisecond): + // Good — no investigation was triggered because monitor mode blocks it + } + + // Verify the finding's status was NOT changed (still failed/timed_out) + stored := ps.findings.Get("timed-out-finding-1") + if stored == nil { + t.Fatal("expected finding to still exist in store") + } + if stored.InvestigationStatus != string(InvestigationStatusFailed) { + t.Fatalf("expected investigation status to remain 'failed', got %q", stored.InvestigationStatus) + } + if stored.InvestigationOutcome != string(InvestigationOutcomeTimedOut) { + t.Fatalf("expected investigation outcome to remain 'timed_out', got %q", stored.InvestigationOutcome) + } +} diff --git a/internal/ai/patrol_run.go b/internal/ai/patrol_run.go index e83e174f5..4e519599d 100644 --- a/internal/ai/patrol_run.go +++ b/internal/ai/patrol_run.go @@ -43,7 +43,9 @@ func (p *PatrolService) Start(ctx context.Context) { go p.patrolLoop(ctx) } -// Stop stops the patrol service +// Stop stops the patrol service. It signals the patrol loop to exit, then +// waits up to 15 seconds for in-flight investigations to finish and +// force-saves findings/investigation state to disk. func (p *PatrolService) Stop() { p.mu.Lock() if !p.running { @@ -52,9 +54,44 @@ func (p *PatrolService) Stop() { } p.running = false close(p.stopCh) + orchestrator := p.investigationOrchestrator + findings := p.findings p.mu.Unlock() log.Info().Msg("Stopping AI Patrol Service") + + // Give investigations 15 seconds to finish (leaves headroom within server's 30s budget) + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer shutdownCancel() + + // Signal orchestrator to cancel running investigations and persist state + if orchestrator != nil { + if err := orchestrator.Shutdown(shutdownCtx); err != nil { + log.Warn().Err(err).Msg("AI Patrol: Investigation orchestrator shutdown returned error") + } + } + + // Wait for investigation goroutines tracked by PatrolService + done := make(chan struct{}) + go func() { + p.investigationWg.Wait() + close(done) + }() + select { + case <-done: + // All investigation goroutines finished + case <-shutdownCtx.Done(): + log.Warn().Msg("AI Patrol: Timed out waiting for investigation goroutines to finish") + } + + // Force-save findings store + if findings != nil { + if err := findings.ForceSave(); err != nil { + log.Error().Err(err).Msg("AI Patrol: Failed to force-save findings during shutdown") + } + } + + log.Info().Msg("AI Patrol Service stopped") } // patrolLoop is the main background loop @@ -236,16 +273,15 @@ func (p *PatrolService) runPatrolWithTrigger(ctx context.Context, trigger Trigge runStats.dockerChecked + runStats.storageChecked + runStats.pbsChecked + runStats.hostsChecked + runStats.kubernetesChecked - // Determine if we can run LLM analysis (requires AI service + license + circuit breaker not open) + // Determine if we can run LLM analysis (requires AI service + circuit breaker not open) aiServiceEnabled := p.aiService != nil && p.aiService.IsEnabled() - hasPatrolFeature := aiServiceEnabled && p.aiService.HasLicenseFeature(FeatureAIPatrol) - canRunLLM := hasPatrolFeature && llmAllowed + canRunLLM := aiServiceEnabled && llmAllowed // Check if we can run LLM analysis (AI-only patrol) if !canRunLLM { - reason := "requires Pulse Pro license for LLM analysis" + reason := "AI not configured - set up a provider in Settings > Pulse Assistant" if !aiServiceEnabled { - reason = "Pulse Assistant service not configured - configure a provider in Settings > Pulse Assistant" + reason = "AI not configured - set up a provider in Settings > Pulse Assistant" } else if !llmAllowed { reason = "circuit breaker is open" GetPatrolMetrics().RecordCircuitBlock() @@ -358,6 +394,9 @@ func (p *PatrolService) runPatrolWithTrigger(ctx context.Context, trigger Trigge log.Debug().Int("cleaned", cleaned).Msg("AI Patrol: Cleaned up old findings") } + // Recover investigations stuck in "running" state (goroutine panicked or was killed) + p.recoverStuckInvestigations() + // Retry investigations that failed due to timeout (shorter cooldown than permanent failures) p.retryTimedOutInvestigations() @@ -631,13 +670,12 @@ func (p *PatrolService) runScopedPatrol(ctx context.Context, scope PatrolScope) // Determine if we can run LLM analysis aiServiceEnabled := p.aiService != nil && p.aiService.IsEnabled() - hasPatrolFeature := aiServiceEnabled && p.aiService.HasLicenseFeature(FeatureAIPatrol) - canRunLLM := hasPatrolFeature && llmAllowed + canRunLLM := aiServiceEnabled && llmAllowed if !canRunLLM { - reason := "requires Pulse Pro license for LLM analysis" + reason := "AI not configured - set up a provider in Settings > Pulse Assistant" if !aiServiceEnabled { - reason = "Pulse Assistant service not configured - configure a provider in Settings > Pulse Assistant" + reason = "AI not configured - set up a provider in Settings > Pulse Assistant" } else if !llmAllowed { reason = "circuit breaker is open" GetPatrolMetrics().RecordCircuitBlock() diff --git a/internal/ai/service.go b/internal/ai/service.go index 5355a3016..919891206 100644 --- a/internal/ai/service.go +++ b/internal/ai/service.go @@ -696,11 +696,9 @@ func (s *Service) StartPatrol(ctx context.Context) { return } - // Check license for AI Patrol feature (Pro only) - if licenseChecker != nil && !licenseChecker.HasFeature(FeatureAIPatrol) { - log.Info().Msg("AI Patrol requires Pulse Pro license - patrol will not run LLM analysis") - // Note: We still configure patrol but it won't have LLM capability - // The patrol service itself should check HasLicenseFeature before LLM calls + // Check license for auto-fix feature (Pro only) - patrol itself is free with BYOK + if licenseChecker != nil && !licenseChecker.HasFeature(FeatureAIAutoFix) { + log.Info().Msg("AI Patrol Auto-Fix requires Pulse Pro license - fixes will require manual approval") } // Configure patrol from AI config (preserve defaults for resource types not in AI config) diff --git a/internal/api/ai_handlers_investigation_additional_test.go b/internal/api/ai_handlers_investigation_additional_test.go index 8f7c43a06..b849faf6a 100644 --- a/internal/api/ai_handlers_investigation_additional_test.go +++ b/internal/api/ai_handlers_investigation_additional_test.go @@ -62,6 +62,10 @@ func (s *stubInvestigationOrchestrator) ReinvestigateFinding(ctx context.Context return nil } +func (s *stubInvestigationOrchestrator) Shutdown(ctx context.Context) error { + return nil +} + type stubChatService struct { messages []ai.ChatMessage } @@ -86,6 +90,10 @@ func (s *stubChatService) DeleteSession(ctx context.Context, sessionID string) e return nil } +func (s *stubChatService) ReloadConfig(ctx context.Context, cfg *config.AIConfig) error { + return nil +} + func TestFindingsStoreWrapper_GetAndUpdate(t *testing.T) { store := ai.NewFindingsStore() store.Add(&ai.Finding{ diff --git a/pkg/server/server.go b/pkg/server/server.go index b4f313c16..ae430c5cd 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -352,6 +352,9 @@ shutdown: log.Error().Err(err).Msg("Server shutdown error") } + // Gracefully stop AI intelligence services (patrol, investigations, triggers) + router.ShutdownAIIntelligence() + // Stop AI chat service (kills sidecar process group) router.StopAIChat(shutdownCtx)