From ce19e9bdb72678d8fa5bfc57cbe1f9198ba6ab56 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 23 Mar 2026 17:46:32 +0800 Subject: [PATCH] feat(robot): enhance manager and execution handling - Added GetManager function to retrieve the global manager instance, returning nil if not started. - Introduced ExecRecovered event to notify about recovered non-terminal executions during manager startup. - Updated execution store to support querying by multiple statuses with ListByStatuses method. - Enhanced integration tests to accommodate longer tick intervals and added sleep delays for stability. - Improved cleanup logic in integration tests to prevent interference from previous execution records. --- agent/robot/api/lifecycle.go | 10 + agent/robot/events/events.go | 1 + agent/robot/manager/integration_clock_test.go | 31 +- .../manager/integration_concurrent_test.go | 22 +- agent/robot/manager/integration_test.go | 33 +- agent/robot/manager/manager.go | 19 ++ agent/robot/manager/recovery.go | 87 +++++ agent/robot/manager/recovery_test.go | 301 ++++++++++++++++++ agent/robot/store/execution.go | 84 ++++- agent/robot/store/execution_test.go | 35 ++ agent/robot/watcher.go | 248 +++++++++++++++ agent/robot/watcher_test.go | 276 ++++++++++++++++ monitor/service.go | 10 + 13 files changed, 1135 insertions(+), 22 deletions(-) create mode 100644 agent/robot/manager/recovery.go create mode 100644 agent/robot/manager/recovery_test.go create mode 100644 agent/robot/watcher.go create mode 100644 agent/robot/watcher_test.go diff --git a/agent/robot/api/lifecycle.go b/agent/robot/api/lifecycle.go index 62343f44..233896fd 100644 --- a/agent/robot/api/lifecycle.go +++ b/agent/robot/api/lifecycle.go @@ -135,6 +135,16 @@ func getManager() (*manager.Manager, error) { return globalManager, nil } +// GetManager returns the global manager instance, or nil if not started. +func GetManager() *manager.Manager { + managerMu.RLock() + defer managerMu.RUnlock() + if globalManager == nil || !globalManager.IsStarted() { + return nil + } + return globalManager +} + // SetManager sets the global manager instance (for testing) func SetManager(m *manager.Manager) { managerMu.Lock() diff --git a/agent/robot/events/events.go b/agent/robot/events/events.go index 8c18b3d6..26e32f72 100644 --- a/agent/robot/events/events.go +++ b/agent/robot/events/events.go @@ -67,6 +67,7 @@ const ( ExecCompleted = "robot.exec.completed" ExecFailed = "robot.exec.failed" ExecCancelled = "robot.exec.cancelled" + ExecRecovered = "robot.exec.recovered" Delivery = "robot.delivery" Message = "robot.message" ) diff --git a/agent/robot/manager/integration_clock_test.go b/agent/robot/manager/integration_clock_test.go index 2bf770f3..892a817b 100644 --- a/agent/robot/manager/integration_clock_test.go +++ b/agent/robot/manager/integration_clock_test.go @@ -64,6 +64,8 @@ func TestIntegrationClockTimesMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + // Verify robot is loaded into cache robot := m.Cache().Get("robot_integ_clock_times1") require.NotNil(t, robot, "Robot should be loaded into cache") @@ -99,6 +101,7 @@ func TestIntegrationClockTimesMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) exec.Reset() // Trigger at 10:30 (not configured) @@ -130,6 +133,7 @@ func TestIntegrationClockTimesMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) exec.Reset() // Trigger at 09:00 on Saturday (not configured) @@ -161,6 +165,8 @@ func TestIntegrationClockTimesMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() // Trigger at 09:00 on Saturday @@ -192,6 +198,7 @@ func TestIntegrationClockTimesMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) exec.Reset() loc, _ := time.LoadLocation("Asia/Shanghai") @@ -239,12 +246,14 @@ func TestIntegrationClockIntervalMode(t *testing.T) { "every": "30m", }) - m, exec := createClockTestManager(t, 100*time.Millisecond, 3, 20) + m, exec := createClockTestManager(t, 10*time.Second, 3, 20) err := m.Start() require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -271,6 +280,8 @@ func TestIntegrationClockIntervalMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -305,12 +316,14 @@ func TestIntegrationClockIntervalMode(t *testing.T) { "every": "1h", // Long interval }) - m, exec := createClockTestManager(t, 100*time.Millisecond, 3, 20) + m, exec := createClockTestManager(t, 10*time.Second, 3, 20) err := m.Start() require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -360,6 +373,8 @@ func TestIntegrationClockDaemonMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -385,6 +400,8 @@ func TestIntegrationClockDaemonMode(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -433,6 +450,8 @@ func TestIntegrationClockTimezone(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -471,6 +490,8 @@ func TestIntegrationClockTimezone(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) @@ -543,6 +564,8 @@ func TestIntegrationClockEdgeCases(t *testing.T) { require.NoError(t, err) defer mgr.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() // Trigger at matching time @@ -597,6 +620,8 @@ func TestIntegrationClockEdgeCases(t *testing.T) { require.NoError(t, err) defer mgr.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() // Trigger at matching time @@ -647,6 +672,8 @@ func TestIntegrationClockEdgeCases(t *testing.T) { require.NoError(t, err) defer mgr.Stop() + time.Sleep(500 * time.Millisecond) + exec.Reset() ctx := types.NewContext(context.Background(), nil) diff --git a/agent/robot/manager/integration_concurrent_test.go b/agent/robot/manager/integration_concurrent_test.go index 835b31ae..58182286 100644 --- a/agent/robot/manager/integration_concurrent_test.go +++ b/agent/robot/manager/integration_concurrent_test.go @@ -66,7 +66,7 @@ func TestIntegrationConcurrentExecution(t *testing.T) { ) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 5, QueueSize: 50}, } m := manager.NewWithConfig(config) @@ -112,7 +112,7 @@ func TestIntegrationConcurrentExecution(t *testing.T) { exec := executor.NewDryRunWithDelay(50 * time.Millisecond) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 5, QueueSize: 50}, } m := manager.NewWithConfig(config) @@ -176,7 +176,7 @@ func TestIntegrationQuotaEnforcement(t *testing.T) { ) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 10, QueueSize: 50}, // Many workers } m := manager.NewWithConfig(config) @@ -213,7 +213,7 @@ func TestIntegrationQuotaEnforcement(t *testing.T) { exec := executor.NewDryRunWithDelay(300 * time.Millisecond) // Slow execution config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 10, QueueSize: 100}, } m := manager.NewWithConfig(config) @@ -282,7 +282,7 @@ func TestIntegrationQuotaEnforcement(t *testing.T) { } config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 10, QueueSize: 50}, } m := manager.NewWithConfig(config) @@ -353,7 +353,7 @@ func TestIntegrationGlobalPoolLimit(t *testing.T) { ) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 3, QueueSize: 100}, // Only 3 workers } m := manager.NewWithConfig(config) @@ -394,7 +394,7 @@ func TestIntegrationGlobalPoolLimit(t *testing.T) { exec := executor.NewDryRunWithDelay(500 * time.Millisecond) // Slow execution config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 1, QueueSize: 5}, // Small queue } m := manager.NewWithConfig(config) @@ -456,7 +456,7 @@ func TestIntegrationPriorityExecution(t *testing.T) { } config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 1, QueueSize: 50}, // Single worker for ordering } m := manager.NewWithConfig(config) @@ -466,6 +466,8 @@ func TestIntegrationPriorityExecution(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + ctx := types.NewContext(context.Background(), nil) // Submit in low-to-high priority order @@ -505,7 +507,7 @@ func TestIntegrationPriorityExecution(t *testing.T) { } config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 1, QueueSize: 50}, } m := manager.NewWithConfig(config) @@ -515,6 +517,8 @@ func TestIntegrationPriorityExecution(t *testing.T) { require.NoError(t, err) defer m.Stop() + time.Sleep(500 * time.Millisecond) + ctx := types.NewContext(context.Background(), nil) // Submit clock first, then human diff --git a/agent/robot/manager/integration_test.go b/agent/robot/manager/integration_test.go index 81131504..12d69d43 100644 --- a/agent/robot/manager/integration_test.go +++ b/agent/robot/manager/integration_test.go @@ -52,9 +52,9 @@ func TestIntegrationSchedulingFlow(t *testing.T) { // Setup: Create a robot with times mode clock config setupIntegrationRobotTimes(t, "robot_integ_flow_clock", "team_integ_flow") - // Create manager with fast tick interval for testing + // Create manager with slow tick interval to avoid auto-tick interference config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 5, QueueSize: 50}, } m := manager.NewWithConfig(config) @@ -246,7 +246,7 @@ func TestIntegrationPhaseProgression(t *testing.T) { }) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 2, QueueSize: 20}, Executor: exec, } @@ -255,6 +255,8 @@ func TestIntegrationPhaseProgression(t *testing.T) { err := m.Start() require.NoError(t, err) + time.Sleep(500 * time.Millisecond) + // Trigger execution ctx := types.NewContext(context.Background(), nil) _, err = m.TriggerManual(ctx, "robot_integ_phases_clock", types.TriggerClock, nil) @@ -287,7 +289,7 @@ func TestIntegrationPhaseProgression(t *testing.T) { }) config := &manager.Config{ - TickInterval: 100 * time.Millisecond, + TickInterval: 10 * time.Second, PoolConfig: &pool.Config{WorkerSize: 2, QueueSize: 20}, Executor: exec, } @@ -296,6 +298,8 @@ func TestIntegrationPhaseProgression(t *testing.T) { err := m.Start() require.NoError(t, err) + time.Sleep(500 * time.Millisecond) + // Trigger execution via human trigger ctx := types.NewContext(context.Background(), nil) _, err = m.TriggerManual(ctx, "robot_integ_phases_human", types.TriggerHuman, nil) @@ -598,17 +602,28 @@ func setupIntegrationRobotNonAutonomous(t *testing.T, memberID, teamID string) { } } -// cleanupIntegrationRobots removes all integration test robots +// cleanupIntegrationRobots removes all integration test robots and their +// non-terminal execution records to prevent recovery interference. func cleanupIntegrationRobots(t *testing.T) { qb := capsule.Query() + + // Clean up execution records for integration robots to prevent + // recoverExecutions from picking them up during Start(). + execModel := model.Select("__yao.agent.execution") + if execModel != nil { + _, err := qb.Table(execModel.MetaData.Table.Name). + Where("member_id", "like", "robot_integ_%"). + WhereIn("status", []interface{}{"running", "paused", "pending", "waiting", "confirming"}). + Delete() + if err != nil { + t.Logf("Warning: execution cleanup error: %v", err) + } + } + m := model.Select("__yao.member") tableName := m.MetaData.Table.Name - - // Delete all robots with member_id starting with "robot_integ_" - // Using LIKE pattern for cleanup _, err := qb.Table(tableName).Where("member_id", "like", "robot_integ_%").Delete() if err != nil { - // Log but don't fail - cleanup errors are not critical t.Logf("Warning: cleanup error: %v", err) } } diff --git a/agent/robot/manager/manager.go b/agent/robot/manager/manager.go index 6def0ac1..79147c9b 100644 --- a/agent/robot/manager/manager.go +++ b/agent/robot/manager/manager.go @@ -7,10 +7,12 @@ import ( "time" "github.com/yaoapp/yao/agent/robot/cache" + "github.com/yaoapp/yao/agent/robot/events" "github.com/yaoapp/yao/agent/robot/executor" "github.com/yaoapp/yao/agent/robot/pool" "github.com/yaoapp/yao/agent/robot/trigger" "github.com/yaoapp/yao/agent/robot/types" + "github.com/yaoapp/yao/event" oauthtypes "github.com/yaoapp/yao/openapi/oauth/types" ) @@ -138,6 +140,9 @@ func (m *Manager) Start() error { return fmt.Errorf("failed to load robots: %w", err) } + // Recover non-terminal executions from previous server lifecycle + pendingNotifications := m.recoverExecutions(m.ctx) + // Set completion callback to clean up ExecutionController when execution finishes m.pool.SetOnComplete(func(execID, memberID string, status types.ExecStatus) { // Remove from ExecutionController (cleans up in-memory tracking) @@ -163,6 +168,15 @@ func (m *Manager) Start() error { m.cache.StartAutoRefresh(ctx, nil) m.started = true + + if len(pendingNotifications) > 0 { + go func() { + for _, n := range pendingNotifications { + _, _ = event.Push(context.Background(), events.ExecRecovered, n) + } + }() + } + return nil } @@ -794,6 +808,11 @@ func (m *Manager) Executor() types.Executor { return m.executor } +// ExecController returns the internal execution controller +func (m *Manager) ExecController() *trigger.ExecutionController { + return m.execController +} + // IsStarted returns true if manager is started func (m *Manager) IsStarted() bool { m.mu.RLock() diff --git a/agent/robot/manager/recovery.go b/agent/robot/manager/recovery.go new file mode 100644 index 00000000..74b5b88a --- /dev/null +++ b/agent/robot/manager/recovery.go @@ -0,0 +1,87 @@ +package manager + +import ( + "context" + "log" + + "github.com/yaoapp/yao/agent/robot/events" + "github.com/yaoapp/yao/agent/robot/store" + "github.com/yaoapp/yao/agent/robot/types" +) + +var nonTerminalStatuses = []types.ExecStatus{ + types.ExecRunning, types.ExecPaused, types.ExecPending, + types.ExecWaiting, types.ExecConfirming, +} + +// recoverExecutions scans the DB for non-terminal executions left by a prior +// server crash. Running/paused/pending records are marked failed; waiting/confirming +// records are kept as-is and returned for notification. +func (m *Manager) recoverExecutions(ctx context.Context) []events.ExecPayload { + execStore := store.NewExecutionStore() + robotStore := store.NewRobotStore() + + var pendingNotifications []events.ExecPayload + affectedMembers := map[string]bool{} + + pageSize := 100 + for page := 1; ; page++ { + result, err := execStore.ListByStatuses(ctx, nonTerminalStatuses, &store.ListOptions{ + Page: page, + PageSize: pageSize, + }) + if err != nil { + log.Printf("[recovery] failed to list non-terminal executions page %d: %v", page, err) + break + } + if len(result.Data) == 0 { + break + } + + for _, record := range result.Data { + affectedMembers[record.MemberID] = true + + switch record.Status { + case types.ExecRunning, types.ExecPaused, types.ExecPending: + if err := execStore.UpdateStatus(ctx, record.ExecutionID, types.ExecFailed, + "execution interrupted by server restart"); err != nil { + log.Printf("[recovery] failed to mark %s as failed: %v", record.ExecutionID, err) + } + case types.ExecWaiting, types.ExecConfirming: + pendingNotifications = append(pendingNotifications, events.ExecPayload{ + ExecutionID: record.ExecutionID, + MemberID: record.MemberID, + TeamID: record.TeamID, + Status: string(record.Status), + }) + } + } + + if len(result.Data) < pageSize { + break + } + } + + fixRobotStatuses(ctx, execStore, robotStore, affectedMembers) + + return pendingNotifications +} + +// fixRobotStatuses sets robots to idle when they no longer have any non-terminal executions. +func fixRobotStatuses(ctx context.Context, execStore *store.ExecutionStore, robotStore *store.RobotStore, members map[string]bool) { + for memberID := range members { + result, err := execStore.ListByStatuses(ctx, nonTerminalStatuses, &store.ListOptions{ + MemberID: memberID, + PageSize: 1, + }) + if err != nil { + log.Printf("[recovery] failed to check remaining executions for %s: %v", memberID, err) + continue + } + if result.Total == 0 { + if err := robotStore.UpdateStatus(ctx, memberID, types.RobotIdle); err != nil { + log.Printf("[recovery] failed to set %s to idle: %v", memberID, err) + } + } + } +} diff --git a/agent/robot/manager/recovery_test.go b/agent/robot/manager/recovery_test.go new file mode 100644 index 00000000..05409f40 --- /dev/null +++ b/agent/robot/manager/recovery_test.go @@ -0,0 +1,301 @@ +package manager_test + +import ( + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/model" + "github.com/yaoapp/xun/capsule" + "github.com/yaoapp/yao/agent/robot/manager" + "github.com/yaoapp/yao/agent/testutils" +) + +const recoveryTestPrefix = "_test_recovery_" + +func TestRecoveryOnRestart(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + t.Run("marks_running_as_failed_on_restart", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"run_001", recoveryTestPrefix+"member_001", "team_r", "running") + insertRecoveryRobot(t, recoveryTestPrefix+"member_001", "team_r") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"run_001") + require.NotNil(t, rec) + assert.Equal(t, "failed", rec["status"]) + errMsg, _ := rec["error"].(string) + assert.Contains(t, errMsg, "server restart") + }) + + t.Run("keeps_waiting_on_restart", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"wait_001", recoveryTestPrefix+"member_002", "team_r", "waiting") + insertRecoveryRobot(t, recoveryTestPrefix+"member_002", "team_r") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"wait_001") + require.NotNil(t, rec) + assert.Equal(t, "waiting", rec["status"]) + }) + + t.Run("keeps_confirming_on_restart", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"conf_001", recoveryTestPrefix+"member_003", "team_r", "confirming") + insertRecoveryRobot(t, recoveryTestPrefix+"member_003", "team_r") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"conf_001") + require.NotNil(t, rec) + assert.Equal(t, "confirming", rec["status"]) + }) + + t.Run("marks_paused_as_failed_on_restart", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"pause_001", recoveryTestPrefix+"member_004", "team_r", "paused") + insertRecoveryRobot(t, recoveryTestPrefix+"member_004", "team_r") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"pause_001") + require.NotNil(t, rec) + assert.Equal(t, "failed", rec["status"]) + }) + + t.Run("no_active_executions_starts_normally", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + assert.True(t, m.IsStarted()) + }) + + t.Run("updates_robot_status_to_idle_after_fail", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"run_002", recoveryTestPrefix+"member_005", "team_r", "running") + insertRecoveryRobot(t, recoveryTestPrefix+"member_005", "team_r") + setRobotStatus(t, recoveryTestPrefix+"member_005", "working") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"run_002") + require.NotNil(t, rec) + assert.Equal(t, "failed", rec["status"]) + + robot := getRobotRecord(t, recoveryTestPrefix+"member_005") + require.NotNil(t, robot) + assert.Equal(t, "idle", robot["robot_status"]) + }) + + t.Run("keeps_robot_status_if_other_waiting", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"run_003", recoveryTestPrefix+"member_006", "team_r", "running") + insertRecoveryExec(t, recoveryTestPrefix+"wait_003", recoveryTestPrefix+"member_006", "team_r", "waiting") + insertRecoveryRobot(t, recoveryTestPrefix+"member_006", "team_r") + setRobotStatus(t, recoveryTestPrefix+"member_006", "working") + + m := manager.New() + err := m.Start() + require.NoError(t, err) + defer m.Stop() + + // Running should be failed + rec := getExecRecord(t, recoveryTestPrefix+"run_003") + require.NotNil(t, rec) + assert.Equal(t, "failed", rec["status"]) + + // Waiting should remain + rec2 := getExecRecord(t, recoveryTestPrefix+"wait_003") + require.NotNil(t, rec2) + assert.Equal(t, "waiting", rec2["status"]) + + // Robot should NOT be set to idle because waiting exec still exists + robot := getRobotRecord(t, recoveryTestPrefix+"member_006") + require.NotNil(t, robot) + assert.NotEqual(t, "idle", robot["robot_status"], + "robot should not be idle when waiting execution exists") + }) + + t.Run("idempotent_on_double_restart", func(t *testing.T) { + cleanupRecoveryData(t) + defer cleanupRecoveryData(t) + + insertRecoveryExec(t, recoveryTestPrefix+"run_004", recoveryTestPrefix+"member_007", "team_r", "running") + insertRecoveryRobot(t, recoveryTestPrefix+"member_007", "team_r") + + // First start + m1 := manager.New() + err := m1.Start() + require.NoError(t, err) + m1.Stop() + + rec := getExecRecord(t, recoveryTestPrefix+"run_004") + require.NotNil(t, rec) + assert.Equal(t, "failed", rec["status"]) + + // Second start — should not panic or error + m2 := manager.New() + err = m2.Start() + require.NoError(t, err) + defer m2.Stop() + + rec2 := getExecRecord(t, recoveryTestPrefix+"run_004") + require.NotNil(t, rec2) + assert.Equal(t, "failed", rec2["status"]) + }) +} + +// ==================== Helpers ==================== + +func insertRecoveryExec(t *testing.T, execID, memberID, teamID, status string) { + t.Helper() + mod := model.Select("__yao.agent.execution") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + now := time.Now() + err := qb.Table(tableName).Insert([]map[string]interface{}{ + { + "execution_id": execID, + "member_id": memberID, + "team_id": teamID, + "trigger_type": "clock", + "status": status, + "phase": "run", + "start_time": now.Add(-1 * time.Hour), + }, + }) + require.NoError(t, err, "insert execution %s", execID) +} + +func insertRecoveryRobot(t *testing.T, memberID, teamID string) { + t.Helper() + mod := model.Select("__yao.member") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + robotConfig := map[string]interface{}{ + "identity": map[string]interface{}{"role": "Recovery Test Robot"}, + "triggers": map[string]interface{}{ + "clock": map[string]interface{}{"enabled": false}, + }, + } + configJSON, _ := json.Marshal(robotConfig) + + err := qb.Table(tableName).Insert([]map[string]interface{}{ + { + "member_id": memberID, + "team_id": teamID, + "member_type": "robot", + "display_name": "Recovery Test " + memberID, + "status": "active", + "role_id": "member", + "autonomous_mode": true, + "robot_status": "idle", + "robot_config": string(configJSON), + }, + }) + require.NoError(t, err, "insert robot %s", memberID) +} + +func setRobotStatus(t *testing.T, memberID, status string) { + t.Helper() + mod := model.Select("__yao.member") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + _, err := qb.Table(tableName).Where("member_id", memberID).Update(map[string]interface{}{ + "robot_status": status, + }) + require.NoError(t, err) +} + +func getExecRecord(t *testing.T, execID string) map[string]interface{} { + t.Helper() + mod := model.Select("__yao.agent.execution") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + rows, err := qb.Table(tableName).Where("execution_id", execID).Limit(1).Get() + require.NoError(t, err) + if len(rows) == 0 { + return nil + } + return map[string]interface{}(rows[0]) +} + +func getRobotRecord(t *testing.T, memberID string) map[string]interface{} { + t.Helper() + mod := model.Select("__yao.member") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + rows, err := qb.Table(tableName).Where("member_id", memberID).Limit(1).Get() + require.NoError(t, err) + if len(rows) == 0 { + return nil + } + return map[string]interface{}(rows[0]) +} + +func cleanupRecoveryData(t *testing.T) { + t.Helper() + + // Clean executions + execMod := model.Select("__yao.agent.execution") + execTable := execMod.MetaData.Table.Name + qb := capsule.Query() + qb.Table(execTable).Where("execution_id", "like", recoveryTestPrefix+"%").Delete() + + // Clean robots + memberMod := model.Select("__yao.member") + memberTable := memberMod.MetaData.Table.Name + qb.Table(memberTable).Where("member_id", "like", recoveryTestPrefix+"%").Delete() + // Also clean via model (soft delete) + memberMod.DeleteWhere(model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "member_id", OP: "like", Value: recoveryTestPrefix + "%"}, + }, + }) +} diff --git a/agent/robot/store/execution.go b/agent/robot/store/execution.go index c65526f0..da6033a9 100644 --- a/agent/robot/store/execution.go +++ b/agent/robot/store/execution.go @@ -7,6 +7,7 @@ import ( "time" "github.com/yaoapp/gou/model" + "github.com/yaoapp/xun/capsule" "github.com/yaoapp/yao/agent/robot/types" ) @@ -65,7 +66,8 @@ type ListOptions struct { MemberID string `json:"member_id,omitempty"` TeamID string `json:"team_id,omitempty"` Status types.ExecStatus `json:"status,omitempty"` - ExcludeStatuses []types.ExecStatus `json:"exclude_statuses,omitempty"` + Statuses []types.ExecStatus `json:"statuses,omitempty"` // Multi-status IN query; takes priority over Status when non-empty + ExcludeStatuses []types.ExecStatus `json:"exclude_statuses,omitempty"` // Exclude these statuses (ne) TriggerType types.TriggerType `json:"trigger_type,omitempty"` Page int `json:"page,omitempty"` PageSize int `json:"pagesize,omitempty"` @@ -171,7 +173,11 @@ func (s *ExecutionStore) List(ctx context.Context, opts *ListOptions) (*ListResu if opts.TeamID != "" { wheres = append(wheres, model.QueryWhere{Column: "team_id", Value: opts.TeamID}) } - if opts.Status != "" { + if len(opts.Statuses) > 0 { + // Backward compat: use the first status for simple equality filter. + // For multi-status IN queries, use ListByStatuses() instead. + wheres = append(wheres, model.QueryWhere{Column: "status", Value: string(opts.Statuses[0])}) + } else if opts.Status != "" { wheres = append(wheres, model.QueryWhere{Column: "status", Value: string(opts.Status)}) } for _, es := range opts.ExcludeStatuses { @@ -232,6 +238,80 @@ func (s *ExecutionStore) List(ctx context.Context, opts *ListOptions) (*ListResu }, nil } +// ListByStatuses queries executions matching any of the given statuses using +// capsule.Query() with WhereIn, which works reliably (unlike model.Paginate +// with OP:"in" or multiple "ne" conditions). +func (s *ExecutionStore) ListByStatuses(ctx context.Context, statuses []types.ExecStatus, opts *ListOptions) (*ListResult, error) { + if len(statuses) == 0 { + return &ListResult{Data: []*ExecutionRecord{}, Total: 0, Page: 1, PageSize: 20}, nil + } + + mod := model.Select(s.modelID) + if mod == nil { + return nil, fmt.Errorf("model %s not found", s.modelID) + } + tableName := mod.MetaData.Table.Name + + statusStrs := make([]interface{}, len(statuses)) + for i, st := range statuses { + statusStrs[i] = string(st) + } + + page := 1 + pageSize := 20 + if opts != nil { + if opts.Page > 0 { + page = opts.Page + } + if opts.PageSize > 0 { + pageSize = opts.PageSize + if pageSize > 100 { + pageSize = 100 + } + } + } + offset := (page - 1) * pageSize + + qb := capsule.Query() + + // Count query + countQB := qb.Table(tableName).WhereIn("status", statusStrs) + if opts != nil && opts.MemberID != "" { + countQB = countQB.Where("member_id", opts.MemberID) + } + total, err := countQB.Count() + if err != nil { + return nil, fmt.Errorf("failed to count executions by statuses: %w", err) + } + + // Data query + dataQB := qb.Table(tableName).WhereIn("status", statusStrs) + if opts != nil && opts.MemberID != "" { + dataQB = dataQB.Where("member_id", opts.MemberID) + } + rows, err := dataQB.OrderBy("start_time", "desc").Limit(pageSize).Offset(offset).Get() + if err != nil { + return nil, fmt.Errorf("failed to list executions by statuses: %w", err) + } + + records := make([]*ExecutionRecord, 0, len(rows)) + for _, row := range rows { + rowMap := map[string]interface{}(row) + record, err := s.mapToRecord(rowMap) + if err != nil { + continue + } + records = append(records, record) + } + + return &ListResult{ + Data: records, + Total: int(total), + Page: page, + PageSize: pageSize, + }, nil +} + // UpdatePhase updates the current phase and its data func (s *ExecutionStore) UpdatePhase(ctx context.Context, executionID string, phase types.Phase, data interface{}) error { mod := model.Select(s.modelID) diff --git a/agent/robot/store/execution_test.go b/agent/robot/store/execution_test.go index 033f77f2..9260f75b 100644 --- a/agent/robot/store/execution_test.go +++ b/agent/robot/store/execution_test.go @@ -231,6 +231,41 @@ func TestExecutionStoreList(t *testing.T) { assert.Equal(t, types.ExecCompleted, r.Status) } }) + + t.Run("list_with_statuses_returns_matching", func(t *testing.T) { + result, err := s.ListByStatuses(ctx, + []types.ExecStatus{types.ExecRunning, types.ExecFailed}, + &store.ListOptions{MemberID: "member_list_002"}) + require.NoError(t, err) + assert.Equal(t, 2, len(result.Data)) + for _, r := range result.Data { + assert.True(t, r.Status == types.ExecRunning || r.Status == types.ExecFailed, + "expected running or failed, got %s", r.Status) + } + }) + + t.Run("list_with_statuses_empty_result", func(t *testing.T) { + result, err := s.ListByStatuses(ctx, + []types.ExecStatus{types.ExecWaiting, types.ExecConfirming}, + &store.ListOptions{MemberID: "member_list_001"}) + require.NoError(t, err) + assert.NotNil(t, result.Data) + assert.Equal(t, 0, len(result.Data)) + }) + + t.Run("list_with_statuses_single_status", func(t *testing.T) { + resultStatuses, err := s.ListByStatuses(ctx, + []types.ExecStatus{types.ExecCompleted}, + &store.ListOptions{MemberID: "member_list_001"}) + require.NoError(t, err) + + resultStatus, err := s.List(ctx, &store.ListOptions{ + Status: types.ExecCompleted, + MemberID: "member_list_001", + }) + require.NoError(t, err) + assert.Equal(t, len(resultStatus.Data), len(resultStatuses.Data)) + }) } // TestExecutionStoreUpdatePhase tests updating phase and phase data diff --git a/agent/robot/watcher.go b/agent/robot/watcher.go new file mode 100644 index 00000000..aad001cb --- /dev/null +++ b/agent/robot/watcher.go @@ -0,0 +1,248 @@ +package robot + +import ( + "context" + "fmt" + "time" + + "github.com/yaoapp/gou/model" + "github.com/yaoapp/yao/agent/robot/api" + "github.com/yaoapp/yao/agent/robot/store" + "github.com/yaoapp/yao/agent/robot/types" + "github.com/yaoapp/yao/monitor" +) + +func init() { + monitor.Register(&robotTasksWatcher{}) +} + +// WatcherConfig holds tuning knobs for the robot-tasks watcher. +type WatcherConfig struct { + Interval time.Duration + MaxRunDuration time.Duration + WaitingTimeout time.Duration + ConfirmTimeout time.Duration +} + +var defaultConfig = WatcherConfig{ + Interval: 5 * time.Minute, + MaxRunDuration: 4 * time.Hour, + WaitingTimeout: 24 * time.Hour, + ConfirmTimeout: 1 * time.Hour, +} + +type robotTasksWatcher struct { + config WatcherConfig +} + +func (w *robotTasksWatcher) Name() string { return "robot-tasks" } + +func (w *robotTasksWatcher) Interval() time.Duration { + if w.config.Interval > 0 { + return w.config.Interval + } + return defaultConfig.Interval +} + +func (w *robotTasksWatcher) Check(ctx context.Context) []monitor.Alert { + mgr := api.GetManager() + if mgr == nil || !mgr.IsStarted() { + return nil + } + + var alerts []monitor.Alert + execStore := store.NewExecutionStore() + now := time.Now() + + alerts = append(alerts, w.checkZombieRunning(ctx, execStore, now)...) + alerts = append(alerts, w.checkWaitingTimeout(ctx, execStore, now)...) + alerts = append(alerts, w.checkConfirmingTimeout(ctx, execStore, now)...) + + return alerts +} + +func (w *robotTasksWatcher) maxRunDuration() time.Duration { + if w.config.MaxRunDuration > 0 { + return w.config.MaxRunDuration + } + return defaultConfig.MaxRunDuration +} + +func (w *robotTasksWatcher) waitingTimeout() time.Duration { + if w.config.WaitingTimeout > 0 { + return w.config.WaitingTimeout + } + return defaultConfig.WaitingTimeout +} + +func (w *robotTasksWatcher) confirmTimeout() time.Duration { + if w.config.ConfirmTimeout > 0 { + return w.config.ConfirmTimeout + } + return defaultConfig.ConfirmTimeout +} + +// checkZombieRunning finds running executions that exceeded maxRunDuration +// and are not tracked by the in-memory execController. +func (w *robotTasksWatcher) checkZombieRunning(ctx context.Context, execStore *store.ExecutionStore, now time.Time) []monitor.Alert { + var alerts []monitor.Alert + maxDur := w.maxRunDuration() + + result, err := execStore.List(ctx, &store.ListOptions{ + Status: types.ExecRunning, + PageSize: 100, + }) + if err != nil { + return nil + } + + mgr := api.GetManager() + + for _, rec := range result.Data { + if rec.StartTime == nil { + continue + } + deadline := rec.StartTime.Add(maxDur) + if now.Before(deadline) { + continue + } + + // Skip if still tracked by execController (genuinely running) + if mgr != nil { + if _, err := mgr.GetExecutionStatus(rec.ExecutionID); err == nil { + continue + } + } + + execID := rec.ExecutionID + alerts = append(alerts, monitor.Alert{ + Level: monitor.Warn, + Target: fmt.Sprintf("execution:%s", execID), + Message: fmt.Sprintf("zombie running execution %s (started %s, exceeded %v)", execID, rec.StartTime.Format(time.RFC3339), maxDur), + Action: func(ctx context.Context) { + mod := model.Select("__yao.agent.execution") + if mod == nil { + return + } + // CAS: only update if still running + mod.UpdateWhere( + model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "execution_id", Value: execID}, + {Column: "status", Value: string(types.ExecRunning)}, + }, + }, + map[string]interface{}{ + "status": string(types.ExecFailed), + "error": "killed by watcher: exceeded max run duration", + "end_time": time.Now(), + }, + ) + }, + }) + } + + return alerts +} + +// checkWaitingTimeout finds waiting executions past the waiting timeout. +func (w *robotTasksWatcher) checkWaitingTimeout(ctx context.Context, execStore *store.ExecutionStore, now time.Time) []monitor.Alert { + var alerts []monitor.Alert + timeout := w.waitingTimeout() + + result, err := execStore.List(ctx, &store.ListOptions{ + Status: types.ExecWaiting, + PageSize: 100, + }) + if err != nil { + return nil + } + + for _, rec := range result.Data { + if rec.UpdatedAt == nil { + continue + } + if now.Before(rec.UpdatedAt.Add(timeout)) { + continue + } + + execID := rec.ExecutionID + alerts = append(alerts, monitor.Alert{ + Level: monitor.Warn, + Target: fmt.Sprintf("execution:%s", execID), + Message: fmt.Sprintf("waiting execution %s timed out (last updated %s, timeout %v)", execID, rec.UpdatedAt.Format(time.RFC3339), timeout), + Action: func(ctx context.Context) { + mod := model.Select("__yao.agent.execution") + if mod == nil { + return + } + mod.UpdateWhere( + model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "execution_id", Value: execID}, + {Column: "status", Value: string(types.ExecWaiting)}, + }, + }, + map[string]interface{}{ + "status": string(types.ExecCancelled), + "error": "cancelled by watcher: waiting timeout exceeded", + "end_time": time.Now(), + }, + ) + }, + }) + } + + return alerts +} + +// checkConfirmingTimeout finds confirming executions past the confirm timeout. +func (w *robotTasksWatcher) checkConfirmingTimeout(ctx context.Context, execStore *store.ExecutionStore, now time.Time) []monitor.Alert { + var alerts []monitor.Alert + timeout := w.confirmTimeout() + + result, err := execStore.List(ctx, &store.ListOptions{ + Status: types.ExecConfirming, + PageSize: 100, + }) + if err != nil { + return nil + } + + for _, rec := range result.Data { + if rec.UpdatedAt == nil { + continue + } + if now.Before(rec.UpdatedAt.Add(timeout)) { + continue + } + + execID := rec.ExecutionID + alerts = append(alerts, monitor.Alert{ + Level: monitor.Info, + Target: fmt.Sprintf("execution:%s", execID), + Message: fmt.Sprintf("confirming execution %s timed out (last updated %s, timeout %v)", execID, rec.UpdatedAt.Format(time.RFC3339), timeout), + Action: func(ctx context.Context) { + mod := model.Select("__yao.agent.execution") + if mod == nil { + return + } + mod.UpdateWhere( + model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "execution_id", Value: execID}, + {Column: "status", Value: string(types.ExecConfirming)}, + }, + }, + map[string]interface{}{ + "status": string(types.ExecCancelled), + "error": "cancelled by watcher: confirmation timeout exceeded", + "end_time": time.Now(), + }, + ) + }, + }) + } + + return alerts +} diff --git a/agent/robot/watcher_test.go b/agent/robot/watcher_test.go new file mode 100644 index 00000000..4c352a89 --- /dev/null +++ b/agent/robot/watcher_test.go @@ -0,0 +1,276 @@ +package robot_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/model" + "github.com/yaoapp/xun/capsule" + "github.com/yaoapp/yao/agent/robot/api" + "github.com/yaoapp/yao/agent/robot/manager" + "github.com/yaoapp/yao/agent/testutils" + "github.com/yaoapp/yao/monitor" + + // Trigger watcher registration via init() + _ "github.com/yaoapp/yao/agent/robot" +) + +const watcherTestPrefix = "_test_watcher_" + +func TestRobotTasksWatcher(t *testing.T) { + if testing.Short() { + t.Skip("Skipping integration test") + } + + testutils.Prepare(t) + defer testutils.Clean(t) + + t.Run("detects_zombie_running_execution", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_001", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + // Insert AFTER manager start so recovery doesn't touch it + oldStart := time.Now().Add(-5 * time.Hour) + insertWatcherExec(t, watcherTestPrefix+"zombie_001", watcherTestPrefix+"member_001", "team_w", "running", &oldStart, nil) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + found := false + for _, a := range alerts { + if a.Target == "execution:"+watcherTestPrefix+"zombie_001" { + found = true + assert.Equal(t, monitor.Warn, a.Level) + assert.Contains(t, a.Message, "zombie") + assert.NotNil(t, a.Action) + } + } + assert.True(t, found, "should detect zombie running execution") + }) + + t.Run("ignores_recent_running_execution", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_002", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + // Insert AFTER start so recovery doesn't mark it failed + recentStart := time.Now().Add(-30 * time.Minute) + insertWatcherExec(t, watcherTestPrefix+"recent_001", watcherTestPrefix+"member_002", "team_w", "running", &recentStart, nil) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + for _, a := range alerts { + assert.NotEqual(t, "execution:"+watcherTestPrefix+"recent_001", a.Target, + "should not alert for recent running execution") + } + }) + + t.Run("detects_waiting_timeout", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_003", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + startTime := time.Now().Add(-25 * time.Hour) + oldUpdated := time.Now().Add(-25 * time.Hour) + insertWatcherExec(t, watcherTestPrefix+"wait_001", watcherTestPrefix+"member_003", "team_w", "waiting", &startTime, &oldUpdated) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + found := false + for _, a := range alerts { + if a.Target == "execution:"+watcherTestPrefix+"wait_001" { + found = true + assert.Equal(t, monitor.Warn, a.Level) + assert.Contains(t, a.Message, "waiting") + assert.NotNil(t, a.Action) + } + } + assert.True(t, found, "should detect waiting timeout") + }) + + t.Run("ignores_recent_waiting_execution", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_004", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + startTime := time.Now().Add(-1 * time.Hour) + recentUpdated := time.Now().Add(-30 * time.Minute) + insertWatcherExec(t, watcherTestPrefix+"wait_002", watcherTestPrefix+"member_004", "team_w", "waiting", &startTime, &recentUpdated) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + for _, a := range alerts { + assert.NotEqual(t, "execution:"+watcherTestPrefix+"wait_002", a.Target, + "should not alert for recent waiting execution") + } + }) + + t.Run("detects_confirming_timeout", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_005", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + startTime := time.Now().Add(-2 * time.Hour) + oldUpdated := time.Now().Add(-2 * time.Hour) + insertWatcherExec(t, watcherTestPrefix+"conf_001", watcherTestPrefix+"member_005", "team_w", "confirming", &startTime, &oldUpdated) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + found := false + for _, a := range alerts { + if a.Target == "execution:"+watcherTestPrefix+"conf_001" { + found = true + assert.Equal(t, monitor.Info, a.Level) + assert.Contains(t, a.Message, "confirming") + assert.NotNil(t, a.Action) + } + } + assert.True(t, found, "should detect confirming timeout") + }) + + t.Run("returns_empty_when_no_issues", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + insertWatcherRobot(t, watcherTestPrefix+"member_006", "team_w") + + m := startWatcherManager(t) + defer m.Stop() + + startTime := time.Now().Add(-1 * time.Hour) + insertWatcherExec(t, watcherTestPrefix+"done_001", watcherTestPrefix+"member_006", "team_w", "completed", &startTime, nil) + insertWatcherExec(t, watcherTestPrefix+"done_002", watcherTestPrefix+"member_006", "team_w", "failed", &startTime, nil) + + watcher := findWatcher(t, "robot-tasks") + alerts := watcher.Check(context.Background()) + + for _, a := range alerts { + assert.NotContains(t, a.Target, watcherTestPrefix, + "should not have alerts for terminal-state executions") + } + }) + + t.Run("handles_nil_manager_gracefully", func(t *testing.T) { + cleanupWatcherData(t) + defer cleanupWatcherData(t) + + // Do NOT start a manager — api.GetManager() returns nil + api.SetManager(nil) + + watcher := findWatcher(t, "robot-tasks") + assert.NotPanics(t, func() { + alerts := watcher.Check(context.Background()) + assert.Empty(t, alerts) + }) + }) +} + +// ==================== Helpers ==================== + +func startWatcherManager(t *testing.T) *manager.Manager { + t.Helper() + m := manager.New() + err := m.Start() + require.NoError(t, err) + api.SetManager(m) + return m +} + +func findWatcher(t *testing.T, name string) monitor.Watcher { + t.Helper() + w := monitor.GetWatcher(name) + require.NotNil(t, w, "watcher %q should be registered", name) + return w +} + +func insertWatcherExec(t *testing.T, execID, memberID, teamID, status string, startTime *time.Time, updatedAt *time.Time) { + t.Helper() + mod := model.Select("__yao.agent.execution") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + data := map[string]interface{}{ + "execution_id": execID, + "member_id": memberID, + "team_id": teamID, + "trigger_type": "clock", + "status": status, + "phase": "run", + } + if startTime != nil { + data["start_time"] = *startTime + } + if updatedAt != nil { + data["updated_at"] = *updatedAt + } + + err := qb.Table(tableName).Insert([]map[string]interface{}{data}) + require.NoError(t, err, "insert execution %s", execID) +} + +func insertWatcherRobot(t *testing.T, memberID, teamID string) { + t.Helper() + mod := model.Select("__yao.member") + tableName := mod.MetaData.Table.Name + qb := capsule.Query() + + err := qb.Table(tableName).Insert([]map[string]interface{}{ + { + "member_id": memberID, + "team_id": teamID, + "member_type": "robot", + "display_name": "Watcher Test " + memberID, + "status": "active", + "role_id": "member", + "autonomous_mode": true, + "robot_status": "idle", + }, + }) + require.NoError(t, err, "insert robot %s", memberID) +} + +func cleanupWatcherData(t *testing.T) { + t.Helper() + + execMod := model.Select("__yao.agent.execution") + execTable := execMod.MetaData.Table.Name + qb := capsule.Query() + qb.Table(execTable).Where("execution_id", "like", watcherTestPrefix+"%").Delete() + + memberMod := model.Select("__yao.member") + memberTable := memberMod.MetaData.Table.Name + qb.Table(memberTable).Where("member_id", "like", watcherTestPrefix+"%").Delete() + memberMod.DeleteWhere(model.QueryParam{ + Wheres: []model.QueryWhere{ + {Column: "member_id", OP: "like", Value: watcherTestPrefix + "%"}, + }, + }) +} diff --git a/monitor/service.go b/monitor/service.go index 0243031c..1e51efd8 100644 --- a/monitor/service.go +++ b/monitor/service.go @@ -35,6 +35,16 @@ type monitorService struct { started bool } +// GetWatcher returns a registered watcher by name, or nil if not found. +func GetWatcher(name string) Watcher { + svc.mu.Lock() + defer svc.mu.Unlock() + if entry, ok := svc.watchers[name]; ok { + return entry.watcher + } + return nil +} + // Register adds a watcher. Call before Start (typically in init). // Registering a watcher with the same name replaces the previous one. func Register(w Watcher) {