vantage-agent/internal/exec/exec.go

package exec

import (
	"bufio"
	"context"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/mrhid6/vantage/agent/internal/grpc/pb"
)

// streamWriter forwards every write to emit() as an ordered chunk. Used as both
// Stdout and Stderr so output interleaves in real execution order. The mutex
// ensures a single stdout/stderr write is not interleaved mid-slice with another.
type streamWriter struct {
	mu   sync.Mutex
	seq  uint64
	emit func(seq uint64, data []byte)
}

func (w *streamWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.emit != nil {
		buf := make([]byte, len(p))
		copy(buf, p)
		w.emit(w.seq, buf)
		w.seq++
	}
	return len(p), nil
}

// WorkspacePath returns the per-run working directory for a workspace id. The
// same id always maps to the same path so RunStep and the cleanup command agree.
func WorkspacePath(workspaceID string) string {
	return filepath.Join(os.TempDir(), "vantage-run-"+workspaceID)
}

// 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 streams output via emit, returning the terminal result
// with empty stdout/stderr but populated exit_code/output_env.
//
// When the command carries a WorkspaceId the step runs with that per-run working
// directory as its cwd (created here if missing); the server removes it once the
// run finishes. The script and env files always live in a private temp dir so
// they never leak into the shared workspace.
func RunStep(cmd *pb.RunStepCmd, emit func(seq uint64, data []byte)) *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)

	workDir := ""
	if cmd.WorkspaceId != "" {
		workDir = WorkspacePath(cmd.WorkspaceId)
		if err := os.MkdirAll(workDir, 0700); err != nil {
			res.ExitCode = 1
			res.Stderr = "create workspace: " + err.Error()
			return res
		}
	}

	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)
	}

	if workDir != "" {
		c.Dir = workDir
	}

	c.Env = append(os.Environ(), "WORKFLOW_ENV="+envFile)
	for k, v := range cmd.Env {
		c.Env = append(c.Env, k+"="+v)
	}

	sw := &streamWriter{emit: emit}
	c.Stdout = sw
	c.Stderr = sw
	runErr := c.Run()

	// stdout/stderr are streamed via emit, not returned in the result.
	if ctx.Err() == context.DeadlineExceeded {
		res.ExitCode = 124
		res.Stderr = "[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 = "[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
}