feat: add jobs to heartbeat (#252)

Author: jajeffries

.gitignore

@@ -2,4 +2,6 @@ diode/
 orb-configs/
 build/
 .coverage/
-.vscode/
+.vscode/
+
+docs/.cursor

agent/agent.go

@@ -66,7 +66,7 @@ func New(logger *slog.Logger, c config.Config) (Agent, error) {
 
 	restartBackendChan := make(chan string, restartBackendChanSize)
 
-	backendStateManager := backend.NewStateManager(c.OrbAgent.ConfigManager.Active, logger, restartBackendChan)
+	backendStateManager := backend.NewStateManager(c.OrbAgent.ConfigManager.Active, logger, restartBackendChan, pm.GetRepo())
 	// Pass a background context to the config manager at construction time. The
 	// manager keeps its own copy and later derives child contexts from the
 	// runtime context supplied in Agent.Start.

agent/backend/backend.go

@@ -10,6 +10,36 @@ import (
 	"github.com/netboxlabs/orb-agent/agent/policies"
 )
 
+// PolicyStatusJob represents a job in the backend status response
+type PolicyStatusJob struct {
+	ID          string    `json:"id"`
+	Status      string    `json:"status"`
+	Reason      string    `json:"reason"`
+	EntityCount *int64    `json:"entity_count,omitempty"`
+	CreatedAt   time.Time `json:"created_at"`
+	UpdatedAt   time.Time `json:"updated_at"`
+}
+
+// PolicyStatus represents policy status from backend status endpoint
+type PolicyStatus struct {
+	Name   string            `json:"name"`
+	Status string            `json:"status"`
+	Jobs   []PolicyStatusJob `json:"jobs"`
+}
+
+// StatusResponse represents the full status response from backend
+type StatusResponse struct {
+	Version       string         `json:"version"`
+	StartTime     time.Time      `json:"start_time"`
+	UpTimeSeconds float64        `json:"up_time_seconds"`
+	Policies      []PolicyStatus `json:"policies"`
+}
+
+// PolicyStatusProvider is an optional interface for backends that support policy status polling
+type PolicyStatusProvider interface {
+	GetPolicyStatus() ([]PolicyStatus, error)
+}
+
 // Running Status types
 const (
 	Unknown RunningStatus = iota

agent/backend/backend_state.go

@@ -5,6 +5,8 @@ import (
 	"log/slog"
 	"sync"
 	"time"
+
+	"github.com/netboxlabs/orb-agent/agent/policies"
 )
 
 // MinRestartTime is the minimum time to wait between restarts
@@ -33,16 +35,18 @@ type stateManager struct {
 	ticker             *time.Ticker
 	logger             *slog.Logger
 	restartBackendChan chan string
+	policyRepo         policies.PolicyRepo
 }
 
 // NewStateManager creates a new StateManager with the given logger and restart channel
-func NewStateManager(activeConfigMgr string, logger *slog.Logger, restartBackendChan chan string) StateManager {
+func NewStateManager(activeConfigMgr string, logger *slog.Logger, restartBackendChan chan string, policyRepo policies.PolicyRepo) StateManager {
 	if configMgrSupportsStateMonitoring(activeConfigMgr) {
 		return &stateManager{
 			backendState:       make(map[string]*State),
 			ticker:             time.NewTicker(BackendMonitorInterval),
 			logger:             logger,
 			restartBackendChan: restartBackendChan,
+			policyRepo:         policyRepo,
 		}
 	}
 	return nullStateManager{}
@@ -99,6 +103,21 @@ func (manager *stateManager) StartBackendMonitor(name string, be Backend) {
 				}
 			}
 			manager.mu.Unlock()
+
+			// Poll policy status if backend supports it
+			if provider, ok := be.(PolicyStatusProvider); ok && manager.policyRepo != nil {
+				statuses, err := provider.GetPolicyStatus()
+				if err != nil {
+					manager.logger.Debug("failed to get policy status", "backend", name, "error", err)
+				} else {
+					for _, ps := range statuses {
+						jobs := convertToJobData(ps.Jobs)
+						if err := manager.policyRepo.UpdateJobs(ps.Name, jobs); err != nil {
+							manager.logger.Debug("failed to update jobs for policy", "policy", ps.Name, "error", err)
+						}
+					}
+				}
+			}
 		}
 	}()
 }
@@ -138,3 +157,19 @@ func (manager *stateManager) Get() map[string]*State {
 	}
 	return result
 }
+
+// convertToJobData converts backend PolicyStatusJob to policies.JobData
+func convertToJobData(statusJobs []PolicyStatusJob) []policies.JobData {
+	jobs := make([]policies.JobData, len(statusJobs))
+	for i, sj := range statusJobs {
+		jobs[i] = policies.JobData{
+			ID:          sj.ID,
+			Status:      sj.Status,
+			Reason:      sj.Reason,
+			EntityCount: sj.EntityCount,
+			CreatedAt:   sj.CreatedAt,
+			UpdatedAt:   sj.UpdatedAt,
+		}
+	}
+	return jobs
+}