Add lifecycle hooks for tasks and the overall workflow graph

internal/run.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"sync"
@@ -201,6 +202,7 @@ func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openB
 	}
 
 	allRunning := false
+	graphCompleted := false
 
 	for {
 		select {
@@ -230,9 +232,13 @@ func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openB
 			}
 
 			if len(failures) > 0 {
+				runLifecycleHook(context.Background(), types.Task{}, wf.Lifecycle.GetOnFailureHook(), os.Stdout, logger)
 				return fmt.Errorf("failed tasks: %v", failures)
 			}
 
+			if graphCompleted {
+				runLifecycleHook(context.Background(), types.Task{}, wf.Lifecycle.GetOnSuccessHook(), os.Stdout, logger)
+			}
 			return nil
 		case event := <-events:
 			switch x := event.(type) {
@@ -275,6 +281,7 @@ func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openB
 
 				if len(pendingTasks) == 0 {
 					logger.Println("✅ exiting because all requested tasks completed and none should be restarted")
+					graphCompleted = true
 					cancel()
 				} else if len(remainingTasks) == 0 {
 					if !allRunning {
@@ -507,13 +514,15 @@ func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openB
 
 					if err != nil {
 						setNodeStatus(node, "failed", fmt.Sprint(err))
+						runLifecycleHook(ctx, t, t.GetOnFailureHook(), out, logger)
 						if t.GetRestartPolicy() != "Never" {
 							restart()
 						}
 						return
 					}
 
 					setNodeStatus(node, "succeeded", "")
+					runLifecycleHook(ctx, t, t.GetOnSuccessHook(), out, logger)
 					if t.GetRestartPolicy() == "Always" {
 						restart()
 					}
@@ -526,3 +535,29 @@ func RunSubgraph(ctx context.Context, cancel context.CancelFunc, port int, openB
 		}
 	}
 }
+
+// runLifecycleHook runs the given lifecycle hook command, logging any errors.
+// It is a best-effort operation: if the hook command fails, the error is logged
+// but does not affect the task's outcome.
+func runLifecycleHook(ctx context.Context, t types.Task, hook *types.LifecycleHook, out io.Writer, logger *log.Logger) {
+	if hook == nil {
+		return
+	}
+	cmd := hook.GetCommand()
+	if len(cmd) == 0 {
+		return
+	}
+	environ, err := types.Environ(types.Spec{}, t)
+	if err != nil {
+		logger.Printf("lifecycle hook: failed to get environment: %v", err)
+		return
+	}
+	c := exec.CommandContext(ctx, cmd[0], cmd[1:]...)
+	c.Dir = t.WorkingDir
+	c.Stdout = out
+	c.Stderr = out
+	c.Env = append(environ, os.Environ()...)
+	if err := c.Run(); err != nil {
+		logger.Printf("lifecycle hook failed: %v", err)
+	}
+}

internal/types/lifecycle.go

@@ -0,0 +1,47 @@
+package types
+
+// LifecycleHook defines a command to run at a specific point in the task lifecycle.
+type LifecycleHook struct {
+	// The command to run.
+	Command Strings `json:"command,omitempty"`
+	// The shell script to run, instead of command.
+	Sh string `json:"sh,omitempty"`
+}
+
+// GetCommand returns the command to run, handling both command and sh forms.
+func (h *LifecycleHook) GetCommand() Strings {
+	if h == nil {
+		return nil
+	}
+	if len(h.Command) > 0 {
+		return h.Command
+	}
+	if h.Sh != "" {
+		return []string{"sh", "-c", h.Sh}
+	}
+	return nil
+}
+
+// Lifecycle describes actions that the system should take in response to lifecycle events.
+type Lifecycle struct {
+	// OnSuccess is the hook to run after the task succeeds.
+	OnSuccess *LifecycleHook `json:"onSuccess,omitempty"`
+	// OnFailure is the hook to run after the task fails.
+	OnFailure *LifecycleHook `json:"onFailure,omitempty"`
+}
+
+// GetOnSuccessHook returns the OnSuccess hook, or nil if the Lifecycle is nil.
+func (l *Lifecycle) GetOnSuccessHook() *LifecycleHook {
+	if l == nil {
+		return nil
+	}
+	return l.OnSuccess
+}
+
+// GetOnFailureHook returns the OnFailure hook, or nil if the Lifecycle is nil.
+func (l *Lifecycle) GetOnFailureHook() *LifecycleHook {
+	if l == nil {
+		return nil
+	}
+	return l.OnFailure
+}

internal/types/lifecycle_test.go

@@ -0,0 +1,86 @@
+package types
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLifecycleHook_GetCommand(t *testing.T) {
+	t.Run("Nil", func(t *testing.T) {
+		var h *LifecycleHook
+		assert.Nil(t, h.GetCommand())
+	})
+	t.Run("Empty", func(t *testing.T) {
+		h := &LifecycleHook{}
+		assert.Nil(t, h.GetCommand())
+	})
+	t.Run("Command", func(t *testing.T) {
+		h := &LifecycleHook{Command: Strings{"echo", "hello"}}
+		assert.Equal(t, Strings{"echo", "hello"}, h.GetCommand())
+	})
+	t.Run("Sh", func(t *testing.T) {
+		h := &LifecycleHook{Sh: "echo hello"}
+		assert.Equal(t, Strings{"sh", "-c", "echo hello"}, h.GetCommand())
+	})
+	t.Run("CommandPreferredOverSh", func(t *testing.T) {
+		h := &LifecycleHook{Command: Strings{"echo", "hi"}, Sh: "echo hello"}
+		assert.Equal(t, Strings{"echo", "hi"}, h.GetCommand())
+	})
+}
+
+func TestLifecycle_GetOnSuccessHook(t *testing.T) {
+	t.Run("NilLifecycle", func(t *testing.T) {
+		var l *Lifecycle
+		assert.Nil(t, l.GetOnSuccessHook())
+	})
+	t.Run("NoOnSuccess", func(t *testing.T) {
+		l := &Lifecycle{}
+		assert.Nil(t, l.GetOnSuccessHook())
+	})
+	t.Run("WithOnSuccess", func(t *testing.T) {
+		hook := &LifecycleHook{Sh: "echo success"}
+		l := &Lifecycle{OnSuccess: hook}
+		assert.Equal(t, hook, l.GetOnSuccessHook())
+	})
+}
+
+func TestLifecycle_GetOnFailureHook(t *testing.T) {
+	t.Run("NilLifecycle", func(t *testing.T) {
+		var l *Lifecycle
+		assert.Nil(t, l.GetOnFailureHook())
+	})
+	t.Run("NoOnFailure", func(t *testing.T) {
+		l := &Lifecycle{}
+		assert.Nil(t, l.GetOnFailureHook())
+	})
+	t.Run("WithOnFailure", func(t *testing.T) {
+		hook := &LifecycleHook{Sh: "echo failed"}
+		l := &Lifecycle{OnFailure: hook}
+		assert.Equal(t, hook, l.GetOnFailureHook())
+	})
+}
+
+func TestTask_GetOnSuccessHook(t *testing.T) {
+	t.Run("NoLifecycle", func(t *testing.T) {
+		task := &Task{}
+		assert.Nil(t, task.GetOnSuccessHook())
+	})
+	t.Run("WithOnSuccess", func(t *testing.T) {
+		hook := &LifecycleHook{Sh: "echo success"}
+		task := &Task{Lifecycle: &Lifecycle{OnSuccess: hook}}
+		assert.Equal(t, hook, task.GetOnSuccessHook())
+	})
+}
+
+func TestTask_GetOnFailureHook(t *testing.T) {
+	t.Run("NoLifecycle", func(t *testing.T) {
+		task := &Task{}
+		assert.Nil(t, task.GetOnFailureHook())
+	})
+	t.Run("WithOnFailure", func(t *testing.T) {
+		hook := &LifecycleHook{Sh: "echo failed"}
+		task := &Task{Lifecycle: &Lifecycle{OnFailure: hook}}
+		assert.Equal(t, hook, task.GetOnFailureHook())
+	})
+}

internal/types/spec.go

@@ -18,6 +18,8 @@ type Spec struct {
 	Env EnvVars `json:"env,omitempty"`
 	// Environment file (e.g. .env) to use
 	Envfile Envfile `json:"envfile,omitempty"`
+	// Lifecycle describes actions that the system should take in response to graph-level lifecycle events.
+	Lifecycle *Lifecycle `json:"lifecycle,omitempty"`
 }
 
 func (s *Spec) GetTerminationGracePeriod() time.Duration {

internal/types/task.go

@@ -90,6 +90,8 @@ type Task struct {
 	Group string `json:"group,omitempty"`
 	// Whether this is the default task to run if no task is specified.
 	Default bool `json:"default,omitempty"`
+	// Lifecycle describes actions that the system should take in response to task lifecycle events.
+	Lifecycle *Lifecycle `json:"lifecycle,omitempty"`
 }
 
 func (t *Task) GetHostPorts() []uint16 {
@@ -216,3 +218,19 @@ func (t *Task) GetStalledTimeout() time.Duration {
 	}
 	return 30 * time.Second
 }
+
+// GetOnSuccessHook returns the lifecycle hook to run when the task succeeds, or nil if none.
+func (t *Task) GetOnSuccessHook() *LifecycleHook {
+	if t.Lifecycle == nil {
+		return nil
+	}
+	return t.Lifecycle.GetOnSuccessHook()
+}
+
+// GetOnFailureHook returns the lifecycle hook to run when the task fails, or nil if none.
+func (t *Task) GetOnFailureHook() *LifecycleHook {
+	if t.Lifecycle == nil {
+		return nil
+	}
+	return t.Lifecycle.GetOnFailureHook()
+}

schema/workflow.schema.json

@@ -72,6 +72,42 @@
       ],
       "title": "HostPath"
     },
+    "Lifecycle": {
+      "properties": {
+        "onSuccess": {
+          "$ref": "#/$defs/LifecycleHook",
+          "title": "onSuccess",
+          "description": "OnSuccess is the hook to run after the task succeeds."
+        },
+        "onFailure": {
+          "$ref": "#/$defs/LifecycleHook",
+          "title": "onFailure",
+          "description": "OnFailure is the hook to run after the task fails."
+        }
+      },
+      "additionalProperties": false,
+      "type": "object",
+      "title": "Lifecycle",
+      "description": "Lifecycle describes actions that the system should take in response to lifecycle events."
+    },
+    "LifecycleHook": {
+      "properties": {
+        "command": {
+          "$ref": "#/$defs/Strings",
+          "title": "command",
+          "description": "The command to run."
+        },
+        "sh": {
+          "type": "string",
+          "title": "sh",
+          "description": "The shell script to run, instead of command."
+        }
+      },
+      "additionalProperties": false,
+      "type": "object",
+      "title": "LifecycleHook",
+      "description": "LifecycleHook defines a command to run at a specific point in the task lifecycle."
+    },
     "Port": {
       "properties": {
         "containerPort": {
@@ -298,6 +334,11 @@
           "type": "boolean",
           "title": "default",
           "description": "Whether this is the default task to run if no task is specified."
+        },
+        "lifecycle": {
+          "$ref": "#/$defs/Lifecycle",
+          "title": "lifecycle",
+          "description": "Lifecycle describes actions that the system should take in response to task lifecycle events."
         }
       },
       "additionalProperties": false,
@@ -394,6 +435,10 @@
         "envfile": {
           "$ref": "#/$defs/Envfile",
           "title": "envfile"
+        },
+        "lifecycle": {
+          "$ref": "#/$defs/Lifecycle",
+          "title": "lifecycle"
         }
       },
       "additionalProperties": false,