From f0c86a3bdf257ea203919f77c8b174268747f194 Mon Sep 17 00:00:00 2001
From: mrhid6 <mrhid6@gmail.com>
Date: Mon, 20 Jul 2026 11:29:14 +0100
Subject: [PATCH] feat(agent): execute RunStepCmd with WORKFLOW_ENV capture

---
 agent/internal/exec/exec.go | 118 ++++++++++++++++++++++++++++++++++++
 agent/internal/sync/sync.go |  11 ++++
 2 files changed, 129 insertions(+)
 create mode 100644 agent/internal/exec/exec.go

diff --git a/agent/internal/exec/exec.go b/agent/internal/exec/exec.go
new file mode 100644
index 0000000..322ef8f
--- /dev/null
+++ b/agent/internal/exec/exec.go
@@ -0,0 +1,118 @@
+package exec
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/mrhid6/vantage/agent/internal/grpc/pb"
+)
+
+// RunStep writes the script to a temp file, provides a WORKFLOW_ENV file for
+// the script to append KEY=value output to, executes it under the requested
+// interpreter, and returns captured output plus parsed output env.
+func RunStep(cmd *pb.RunStepCmd) *pb.StepResult {
+	res := &pb.StepResult{CommandId: "", OutputEnv: map[string]string{}}
+
+	dir, err := os.MkdirTemp("", "vantage-step-")
+	if err != nil {
+		res.ExitCode = 1
+		res.Stderr = "create temp dir: " + err.Error()
+		return res
+	}
+	defer os.RemoveAll(dir)
+
+	envFile := filepath.Join(dir, "workflow_env")
+	if err := os.WriteFile(envFile, nil, 0600); err != nil {
+		res.ExitCode = 1
+		res.Stderr = "create env file: " + err.Error()
+		return res
+	}
+
+	var scriptPath string
+	var c *exec.Cmd
+	timeout := time.Duration(cmd.TimeoutSeconds) * time.Second
+	if timeout <= 0 {
+		timeout = 30 * time.Minute
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	switch cmd.Interpreter {
+	case "powershell":
+		scriptPath = filepath.Join(dir, "step.ps1")
+		if err := os.WriteFile(scriptPath, []byte(cmd.Script), 0600); err != nil {
+			res.ExitCode = 1
+			res.Stderr = err.Error()
+			return res
+		}
+		shell := "pwsh"
+		if runtime.GOOS == "windows" {
+			if _, err := exec.LookPath("pwsh"); err != nil {
+				shell = "powershell.exe"
+			}
+		}
+		c = exec.CommandContext(ctx, shell, "-NoProfile", "-NonInteractive", "-File", scriptPath)
+	default: // "bash"
+		scriptPath = filepath.Join(dir, "step.sh")
+		if err := os.WriteFile(scriptPath, []byte(cmd.Script), 0700); err != nil {
+			res.ExitCode = 1
+			res.Stderr = err.Error()
+			return res
+		}
+		c = exec.CommandContext(ctx, "bash", scriptPath)
+	}
+
+	c.Env = append(os.Environ(), "WORKFLOW_ENV="+envFile)
+	for k, v := range cmd.Env {
+		c.Env = append(c.Env, k+"="+v)
+	}
+
+	var stdout, stderr bytes.Buffer
+	c.Stdout = &stdout
+	c.Stderr = &stderr
+	runErr := c.Run()
+
+	res.Stdout = stdout.String()
+	res.Stderr = stderr.String()
+	if ctx.Err() == context.DeadlineExceeded {
+		res.ExitCode = 124
+		res.Stderr += "\n[vantage] step timed out"
+	} else if ee, ok := runErr.(*exec.ExitError); ok {
+		res.ExitCode = ee.ExitCode()
+	} else if runErr != nil {
+		res.ExitCode = 1
+		res.Stderr += "\n[vantage] " + runErr.Error()
+	}
+
+	res.OutputEnv = parseEnvFile(envFile)
+	return res
+}
+
+// parseEnvFile reads KEY=value lines (last write wins). Blank lines and lines
+// without '=' are ignored.
+func parseEnvFile(path string) map[string]string {
+	out := map[string]string{}
+	f, err := os.Open(path)
+	if err != nil {
+		return out
+	}
+	defer f.Close()
+	sc := bufio.NewScanner(f)
+	sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
+	for sc.Scan() {
+		line := sc.Text()
+		i := strings.IndexByte(line, '=')
+		if i <= 0 {
+			continue
+		}
+		out[line[:i]] = line[i+1:]
+	}
+	return out
+}
diff --git a/agent/internal/sync/sync.go b/agent/internal/sync/sync.go
index b6a8cc8..f9828e8 100644
--- a/agent/internal/sync/sync.go
+++ b/agent/internal/sync/sync.go
@@ -17,6 +17,7 @@ import (
 	"time"
 
 	"github.com/mrhid6/vantage/agent/internal/config"
+	agentexec "github.com/mrhid6/vantage/agent/internal/exec"
 	grpcclient "github.com/mrhid6/vantage/agent/internal/grpc"
 	"github.com/mrhid6/vantage/agent/internal/grpc/pb"
 	"github.com/mrhid6/vantage/agent/internal/keys"
@@ -186,6 +187,16 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
 		if cmd.ApplyUpdates != nil {
 			go handleApplyUpdates(cfg, cmd)
 		}
+		if cmd.RunStep != nil {
+			res := agentexec.RunStep(cmd.RunStep)
+			res.CommandId = cmd.CommandId
+			_ = stream.Send(&pb.AgentMessage{
+				ServerId:   cfg.ServerID,
+				AgentToken: cfg.AgentToken,
+				StepResult: res,
+			})
+			continue
+		}
 	}
 }