update subprocess package

subprocess/subprocess.go

@@ -2,25 +2,36 @@ package subprocess
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io"
 	"os"
 	"os/exec"
 	"sync"
 )
 
+// LogFormatter is a function type that formats a single line of output
+type LogFormatter func(line string) string
+
 type Subprocess struct {
-	Command   string
-	Args      []string
-	LogPrefix string
-	Stdin     io.Reader
-	Stdout    io.Writer
-	Stderr    io.Writer
-	Dir       string
-	Env       []string
+	Command         string
+	Args            []string
+	StdoutFormatter LogFormatter // Function to format stdout lines
+	StderrFormatter LogFormatter // Function to format stderr lines
+	Stdin           io.Reader
+	Stdout          io.Writer
+	Stderr          io.Writer
+	Dir             string
+	Env             []string
 }
 
+// Run executes the subprocess without context (for backward compatibility)
 func Run(sb *Subprocess) error {
+	return RunWithContext(context.Background(), sb)
+}
+
+// RunWithContext executes the subprocess with the provided context
+func RunWithContext(ctx context.Context, sb *Subprocess) error {
 	if sb.Stdin == nil {
 		sb.Stdin = os.Stdin
 	}
@@ -31,26 +42,35 @@ func Run(sb *Subprocess) error {
 		sb.Stderr = os.Stderr
 	}
 
-	cmd := exec.Command(sb.Command, sb.Args...)
+	// Use Background context if ctx is nil
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	// Use exec.CommandContext to enable context control
+	cmd := exec.CommandContext(ctx, sb.Command, sb.Args...)
 	cmd.Dir = sb.Dir
 	cmd.Env = sb.Env
 	cmd.Stdin = sb.Stdin
 
 	m := new(sync.Mutex)
 	wg := &sync.WaitGroup{}
 
-	// use pipes to add a prefix to each line.
+	// Variables to collect errors from goroutines
+	var stdoutErr, stderrErr error
+
+	// Use pipes to add a prefix to each line
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
 		return fmt.Errorf("failed to get stdout pipe: %w", err)
 	}
 
 	wg.Add(1)
 	go func() {
-		if err := scanLines(stdout, sb.Stdout, sb.LogPrefix, m); err != nil {
-			_, _ = fmt.Fprintf(sb.Stderr, "failed to scan stdout: %v", err)
+		defer wg.Done()
+		if err := scanLines(stdout, sb.Stdout, sb.StdoutFormatter, m); err != nil {
+			stdoutErr = fmt.Errorf("failed to scan stdout: %w", err)
 		}
-		wg.Done()
 	}()
 
 	stderr, err := cmd.StderrPipe()
@@ -59,31 +79,66 @@ func Run(sb *Subprocess) error {
 	}
 	wg.Add(1)
 	go func() {
-		if err := scanLines(stderr, sb.Stderr, sb.LogPrefix, m); err != nil {
-			_, _ = fmt.Fprintf(sb.Stderr, "failed to scan stderr: %v", err)
+		defer wg.Done()
+		if err := scanLines(stderr, sb.Stderr, sb.StderrFormatter, m); err != nil {
+			stderrErr = fmt.Errorf("failed to scan stderr: %w", err)
 		}
-		wg.Done()
 	}()
 
-	err = cmd.Run()
+	// Start the process
+	err = cmd.Start()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to start command: %w", err)
 	}
+
+	// Wait for process completion (ensures process is fully terminated)
+	processErr := cmd.Wait()
+
+	// Wait for stdout/stderr processing completion
 	wg.Wait()
 
+	// Return process error first, then scan errors
+	if processErr != nil {
+		return processErr
+	}
+
+	if stdoutErr != nil {
+		return stdoutErr
+	}
+	if stderrErr != nil {
+		return stderrErr
+	}
+
 	return nil
 }
 
-func scanLines(src io.ReadCloser, dest io.Writer, prefix string, m *sync.Mutex) error {
+func scanLines(src io.ReadCloser, dest io.Writer, formatter LogFormatter, m *sync.Mutex) error {
+	defer src.Close()
 	scanner := bufio.NewScanner(src)
 	for scanner.Scan() {
-		// prevent mixing data in a line.
+		line := scanner.Text()
+
+		// Apply formatter if provided
+		if formatter != nil {
+			line = formatter(line)
+		}
+
+		// Prevent mixing data in a line
 		m.Lock()
-		_, _ = fmt.Fprintf(dest, "%s%s\n", prefix, scanner.Text())
+		_, _ = fmt.Fprintf(dest, "%s\n", line)
 		m.Unlock()
 	}
 	if err := scanner.Err(); err != nil {
 		return fmt.Errorf("failed to scan lines: %w", err)
 	}
 	return nil
 }
+
+// Built-in formatter functions
+
+// PrefixFormatter creates a formatter that adds a prefix to each line
+func PrefixFormatter(prefix string) LogFormatter {
+	return func(line string) string {
+		return prefix + line
+	}
+}

subprocess/subprocess_test.go

@@ -0,0 +1,40 @@
+package subprocess
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestRunWithContext(t *testing.T) {
+	var stdout bytes.Buffer
+	err := RunWithContext(context.Background(), &Subprocess{
+		Command: "echo",
+		Args:    []string{"hello world"},
+		Stdout:  &stdout,
+	})
+	if err != nil {
+		t.Fatalf("Failed: %v", err)
+	}
+
+	if !strings.Contains(stdout.String(), "hello world") {
+		t.Errorf("Expected 'hello world', got: %s", stdout.String())
+	}
+}
+
+func TestRunWithContext_Timeout(t *testing.T) {
+	var stdout bytes.Buffer
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	err := RunWithContext(ctx, &Subprocess{
+		Command: "sleep",
+		Args:    []string{"10"},
+		Stdout:  &stdout,
+	})
+	if err == nil {
+		t.Fatal("Expected timeout error")
+	}
+	t.Logf("Expected timeout error, got: %v", err)
+}

viewkit.go

@@ -82,12 +82,13 @@ type ViewKit struct {
 	// ViteDevServerStdout is a writer for the Vite dev server stdout.
 	// The default value is os.Stdout.
 	ViteDevServerStdout io.Writer
+	// ViteDevServerStdoutFormatter is a formatter for the Vite dev server stdout.
+	ViteDevServerStdoutFormatter subprocess.LogFormatter
 	// ViteDevServerStderr is a writer for the Vite dev server stderr.
 	// The default value is os.Stderr.
 	ViteDevServerStderr io.Writer
-	// ViteDevServerLogPrefix is a prefix for the Vite dev server log.
-	// The default value is "[vite] ".
-	ViteDevServerLogPrefix string
+	// ViteDevServerStderrFormatter is a formatter for the Vite dev server stderr.
+	ViteDevServerStderrFormatter subprocess.LogFormatter
 	// ViteManifest is a Vite manifest.
 	// This is needed to resolve the asset paths in production environment.
 	ViteManifest ViteManifest
@@ -122,8 +123,9 @@ func New() *ViewKit {
 		ViteDevServerURL:                      "http://localhost:5173",
 		ViteDevServerCommand:                  []string{"npx", "vite", "--clearScreen=false"},
 		ViteDevServerStdout:                   os.Stdout,
+		ViteDevServerStdoutFormatter:          subprocess.PrefixFormatter("[echo-viewkit:vite] "),
 		ViteDevServerStderr:                   os.Stderr,
-		ViteDevServerLogPrefix:                "[echo-viewkit:vite] ",
+		ViteDevServerStderrFormatter:          subprocess.PrefixFormatter("[echo-viewkit:vite] "),
 		ViteManifest:                          nil,
 		ViteBasePath:                          "",
 	}
@@ -272,11 +274,12 @@ func (v *ViewKit) StartViteDevServer() error {
 	}
 
 	return subprocess.Run(&subprocess.Subprocess{
-		Command:   v.ViteDevServerCommand[0],
-		Args:      v.ViteDevServerCommand[1:],
-		Stdout:    v.ViteDevServerStdout,
-		Stderr:    v.ViteDevServerStderr,
-		LogPrefix: v.ViteDevServerLogPrefix,
+		Command:         v.ViteDevServerCommand[0],
+		Args:            v.ViteDevServerCommand[1:],
+		Stdout:          v.ViteDevServerStdout,
+		Stderr:          v.ViteDevServerStderr,
+		StdoutFormatter: v.ViteDevServerStdoutFormatter,
+		StderrFormatter: v.ViteDevServerStderrFormatter,
 	})
 }