From 5d7208883787e96511d02b4f79f013f54cef2bc7 Mon Sep 17 00:00:00 2001
From: mrhid6 <mrhid6@gmail.com>
Date: Mon, 20 Jul 2026 12:36:37 +0100
Subject: [PATCH] feat(agent): stream step output chunks over CommandStream

---
 agent/internal/exec/exec.go | 41 ++++++++++++++++++++++++++++---------
 agent/internal/sync/sync.go | 15 +++++++++++++-
 2 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/agent/internal/exec/exec.go b/agent/internal/exec/exec.go
index 322ef8f..72f59e7 100644
--- a/agent/internal/exec/exec.go
+++ b/agent/internal/exec/exec.go
@@ -2,22 +2,44 @@ package exec
 
 import (
 	"bufio"
-	"bytes"
 	"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
+}
+
 // 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 {
+// interpreter, and streams output via emit, returning the terminal result
+// with empty stdout/stderr but populated exit_code/output_env.
+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-")
@@ -74,21 +96,20 @@ func RunStep(cmd *pb.RunStepCmd) *pb.StepResult {
 		c.Env = append(c.Env, k+"="+v)
 	}
 
-	var stdout, stderr bytes.Buffer
-	c.Stdout = &stdout
-	c.Stderr = &stderr
+	sw := &streamWriter{emit: emit}
+	c.Stdout = sw
+	c.Stderr = sw
 	runErr := c.Run()
 
-	res.Stdout = stdout.String()
-	res.Stderr = stderr.String()
+	// stdout/stderr are streamed via emit, not returned in the result.
 	if ctx.Err() == context.DeadlineExceeded {
 		res.ExitCode = 124
-		res.Stderr += "\n[vantage] step timed out"
+		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 += "\n[vantage] " + runErr.Error()
+		res.Stderr = "[vantage] " + runErr.Error()
 	}
 
 	res.OutputEnv = parseEnvFile(envFile)
diff --git a/agent/internal/sync/sync.go b/agent/internal/sync/sync.go
index 6d63b8b..99a0430 100644
--- a/agent/internal/sync/sync.go
+++ b/agent/internal/sync/sync.go
@@ -200,8 +200,21 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
 		}
 		if cmd.RunStep != nil {
 			go func(rc *pb.RunStepCmd, cid string) {
-				res := agentexec.RunStep(rc)
+				emit := func(seq uint64, data []byte) {
+					_ = send(&pb.AgentMessage{
+						ServerId:   cfg.ServerID,
+						AgentToken: cfg.AgentToken,
+						StepOutput: &pb.StepOutputChunk{CommandId: cid, Seq: seq, Data: data},
+					})
+				}
+				res := agentexec.RunStep(rc, emit)
 				res.CommandId = cid
+				// Final eof marker so the server closes the log file.
+				_ = send(&pb.AgentMessage{
+					ServerId:   cfg.ServerID,
+					AgentToken: cfg.AgentToken,
+					StepOutput: &pb.StepOutputChunk{CommandId: cid, Eof: true},
+				})
 				_ = send(&pb.AgentMessage{
 					ServerId:   cfg.ServerID,
 					AgentToken: cfg.AgentToken,