spark-mcp/internal/executor/spark_submit.go

// Package executor runs the local spark-submit binary for the spark_submit Tool.
//
// It is the single place in the codebase that shells out to user-configured
// binaries. To prevent command injection the binary is taken verbatim from
// cluster.SparkSubmitExecuteBin and arguments are passed via exec.Command's
// slice form — never as a single string, never via "sh -c".
package executor

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

// stdoutCap caps spark-submit stdout/stderr to keep the per-Tool log file
// from blowing up on verbose jobs. 10 MiB is generous for the "Submitted
// application application_xxx" line and a few progress lines.
const stdoutCap = 10 << 20

// submittedAppPattern matches the explicit success line, including the
// application id captured in group 1. This mirrors the Python reference.
var submittedAppPattern = regexp.MustCompile(`Submitted application (\S+)`)

// trackingURLPattern matches "tracking URL: <url>" on stderr. The URL tail
// is used as a fallback source for the application id.
var trackingURLPattern = regexp.MustCompile(`tracking URL:\s+(\S+)`)

// SparkSubmitOpts is the resolved command line for a single spark-submit run.
//
// Binary must be an absolute path (validated at cluster create/update time).
// Args are the post-binary CLI args, already including any cluster-default
// args prepended by the caller.
type SparkSubmitOpts struct {
	Binary string
	Args   []string
	// Timeout is the wall-clock budget for the entire run. Zero means no
	// timeout. We use context.WithTimeout to enforce it.
	Timeout time.Duration
}

// Result is what spark_submit returns to the LLM.
type Result struct {
	AppID       string `json:"app_id,omitempty"`
	TrackingURL string `json:"tracking_url,omitempty"`
	ExitCode    int    `json:"exit_code"`
	StdoutTail  string `json:"stdout_tail"`
	StderrTail  string `json:"stderr_tail"`
	DurationMS  int64  `json:"duration_ms"`
}

// ErrBinaryEmpty is returned when the cluster has no spark-submit binary
// configured. The agent should call list_clusters to surface this.
var ErrBinaryEmpty = errors.New("executor: spark_submit_execute_bin is empty")

// Run executes the configured spark-submit and returns the parsed result.
//
// Command construction: exec.Command(name, args...) — the slice form means
// no shell, so a malicious arg like "foo; rm -rf /" is passed as a single
// argument to spark-submit, never executed.
func Run(ctx context.Context, opts SparkSubmitOpts) (*Result, error) {
	if strings.TrimSpace(opts.Binary) == "" {
		return nil, ErrBinaryEmpty
	}

	if opts.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
		defer cancel()
	}

	cmd := exec.CommandContext(ctx, opts.Binary, opts.Args...)
	// Critical: do NOT call cmd.Run() with /bin/sh. The slice form above is
	// the entire injection defense.

	stdout := &capBuf{cap: stdoutCap}
	stderr := &capBuf{cap: stdoutCap}
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	start := time.Now()
	err := cmd.Run()
	dur := time.Since(start)

	res := &Result{
		ExitCode:   exitCodeFromErr(err),
		StdoutTail: stdout.String(),
		StderrTail: stderr.String(),
		DurationMS: dur.Milliseconds(),
	}

	appID, trackingURL, ok := parseSparkSubmitOutput(res.StdoutTail, res.StderrTail)
	res.AppID = appID
	res.TrackingURL = trackingURL

	if !ok {
		return res, fmt.Errorf("executor: spark-submit ran but no application_id in output (exit_code=%d)", res.ExitCode)
	}
	if err != nil {
		return res, fmt.Errorf("executor: spark-submit failed: %w", err)
	}
	return res, nil
}

// parseSparkSubmitOutput extracts the application id and tracking URL from
// spark-submit output. It mirrors the Python reference implementation: first
// look for "Submitted application <id>", then fall back to "tracking URL:",
// deriving the id from the last path segment. The returned ok is false when
// no application id could be determined.
func parseSparkSubmitOutput(stdout, stderr string) (appID, trackingURL string, ok bool) {
	// Stage 1: explicit success line, prefer stdout then stderr.
	var trackingFound bool
	if m := submittedAppPattern.FindStringSubmatch(stdout); m != nil {
		appID = m[1]
		ok = true
	} else if m := submittedAppPattern.FindStringSubmatch(stderr); m != nil {
		appID = m[1]
		ok = true
	}

	// Stage 2: tracking URL fallback (stderr only, like the Python reference).
	if !ok {
		if m := trackingURLPattern.FindStringSubmatch(stderr); m != nil {
			trackingURL = m[1]
			trackingFound = true
			appID = extractAppIDFromTrackingURL(trackingURL)
			ok = appID != ""
			if !ok {
				trackingURL = ""
			}
		}
	}

	// Tracking URL is independent: if present anywhere on stderr, capture it.
	if trackingURL == "" && !trackingFound {
		if m := trackingURLPattern.FindStringSubmatch(stderr); m != nil {
			trackingURL = m[1]
		}
	}
	return
}

// extractAppIDFromTrackingURL returns the last path segment if it starts with
// "application_", otherwise an empty string.
func extractAppIDFromTrackingURL(rawURL string) string {
	trimmed := strings.TrimRight(rawURL, "/")
	idx := strings.LastIndex(trimmed, "/")
	if idx < 0 {
		return ""
	}
	tail := trimmed[idx+1:]
	if strings.HasPrefix(tail, "application_") {
		return tail
	}
	return ""
}

// exitCodeFromErr returns 0 on nil, the process's exit code on *exec.ExitError,
// and 1 on other errors (e.g. context deadline). exec.ExitError is the typed
// way to ask the OS-level exit code without inspecting strings.
func exitCodeFromErr(err error) int {
	if err == nil {
		return 0
	}
	var exitErr *exec.ExitError
	if errors.As(err, &exitErr) {
		return exitErr.ExitCode()
	}
	// context.DeadlineExceeded or failed to start — surface as 1.
	return 1
}

// capBuf is an io.Writer that keeps at most `cap` bytes from the end.
//
// The point isn't to deliver partial stdout to the LLM (LLM rarely needs the
// full stream) — it's to bound memory and the per-Tool log file. We keep
// the tail because that's where "Submitted application application_xxx" and
// final error messages live.
type capBuf struct {
	cap int
	buf bytes.Buffer
}

func (b *capBuf) Write(p []byte) (int, error) {
	n, err := b.buf.Write(p)
	if err != nil {
		return n, err
	}
	if b.buf.Len() > b.cap {
		// Drop the front. Cheaper than a ring buffer; works fine at 10 MiB.
		overflow := b.buf.Len() - b.cap
		_ = b.buf.Next(overflow)
	}
	return len(p), nil
}

func (b *capBuf) String() string { return b.buf.String() }

// Ensure io.Discard stays referenced; we keep this in case future variants
// want to write to a logger instead of capturing.
var _ = io.Discard