From 6a289e5cf0c2bc0bed7e89b3dd76c06ad9e3c0b2 Mon Sep 17 00:00:00 2001
From: Aditya Thebe <contact@adityathebe.com>
Date: Fri, 20 Sep 2024 10:39:27 +0545
Subject: [PATCH] fix: runsOn handling

---
 cmd/playbook.go           |  6 ++----
 playbook/runner/runner.go | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/cmd/playbook.go b/cmd/playbook.go
index c46de861e..fb75807d2 100644
--- a/cmd/playbook.go
+++ b/cmd/playbook.go
@@ -148,13 +148,13 @@ var Run = &cobra.Command{
 			logger.Fatalf(err.Error())
 			return
 		}
+
 		if action == nil {
 			logger.Errorf("No actions to run")
 			return
 		}
 
 		for action != nil {
-
 			if delayed, err := runner.CheckDelay(ctx, *p, *run, action, step); err != nil {
 				ctx.Errorf("Error running action %s: %v", action.Name, err)
 				break
@@ -163,7 +163,6 @@ var Run = &cobra.Command{
 			}
 
 			runAction, err := run.StartAction(ctx.DB(), action.Name)
-
 			if err != nil {
 				ctx.Errorf("Error starting action %s: %v", action.Name, err)
 				break
@@ -175,11 +174,11 @@ var Run = &cobra.Command{
 			}
 
 			action, _, err = runner.GetNextActionToRun(ctx, *p, *run)
-
 			if action != nil && action.Name == runAction.Name {
 				ctx.Errorf("%v", ctx.Oops().Errorf("Action cycle detected for: %s", action.Name))
 				shutdown.ShutdownAndExit(1, "")
 			}
+
 			if err != nil {
 				ctx.Errorf("Error getting next action %s: %v", action.Name, err)
 				break
@@ -197,7 +196,6 @@ var Run = &cobra.Command{
 		if summary.Run.Status != models.PlaybookRunStatusCompleted {
 			shutdown.ShutdownAndExit(1, fmt.Sprintf("Playbook run status: %s ", summary.Run.Status))
 		}
-
 	},
 }
 
diff --git a/playbook/runner/runner.go b/playbook/runner/runner.go
index aaefdf2cf..1ac27f293 100644
--- a/playbook/runner/runner.go
+++ b/playbook/runner/runner.go
@@ -142,12 +142,26 @@ func CheckDelay(ctx context.Context, playbook models.Playbook, run models.Playbo
 	return false, nil
 }
 
+func canRunOnLocalAgent(spec v1.PlaybookSpec, action *v1.PlaybookAction) bool {
+	if len(spec.RunsOn) != 0 {
+		if !lo.Contains(spec.RunsOn, Main) {
+			return false
+		}
+	}
+
+	if len(action.RunsOn) != 0 {
+		if !lo.Contains(action.RunsOn, Main) {
+			return false
+		}
+	}
+
+	return true
+}
+
 // ScheduleRun finds the next action step that needs to run and
 // creates the PlaybookActionRun in a scheduled status, with an optional agentId
 func ScheduleRun(ctx context.Context, run models.PlaybookRun) error {
-
 	var playbook models.Playbook
-
 	if err := ctx.DB().First(&playbook, run.PlaybookID).Error; err != nil {
 		return ctx.Oops("db").Wrap(err)
 	}
@@ -179,7 +193,12 @@ func ScheduleRun(ctx context.Context, run models.PlaybookRun) error {
 		}, action.Name))
 	}
 
-	if run.AgentID == nil && (len(action.RunsOn) == 0 || lo.Contains(action.RunsOn, Main)) {
+	var playbookSpec v1.PlaybookSpec
+	if err := json.Unmarshal(playbook.Spec, &playbookSpec); err != nil {
+		return ctx.Oops().Wrap(err)
+	}
+
+	if canRunOnLocalAgent(playbookSpec, action) {
 		if runAction, err := run.StartAction(ctx.DB(), action.Name); err != nil {
 			return ctx.Oops("db").Wrap(err)
 		} else {
@@ -187,10 +206,15 @@ func ScheduleRun(ctx context.Context, run models.PlaybookRun) error {
 		}
 	}
 
-	if agent, err := db.FindFirstAgent(ctx, action.RunsOn...); err != nil {
+	eligibleAgents := action.RunsOn
+	if len(eligibleAgents) == 0 {
+		eligibleAgents = playbookSpec.RunsOn
+	}
+
+	if agent, err := db.FindFirstAgent(ctx, eligibleAgents...); err != nil {
 		return ctx.Oops("db").Wrap(err)
 	} else if agent == nil {
-		return ctx.Oops("db").Wrapf(err, "failed to find any agent (%s)", strings.Join(action.RunsOn, ","))
+		return ctx.Oops("db").Wrapf(err, "failed to find any agent (%s)", strings.Join(eligibleAgents, ","))
 	} else {
 		// Assign the action to an agent and step the status to Waiting
 		// When the agent polls for new actions to run, we return and then set the status to Running