spark-mcp/internal/mcp/tools/spark_submit.go

package tools

import (
	"context"
	"fmt"
	"log/slog"
	"path/filepath"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"

	"spark-mcp-go/internal/executor"
	"spark-mcp-go/internal/uploads"
)

const SparkSubmitName = "spark_submit"

// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
func NewSparkSubmitTool() mcp.Tool {
	return mcp.NewTool(SparkSubmitName,
		mcp.WithDescription("Submit a Spark application to a configured cluster using the cluster's local spark-submit binary. The `args` array is appended to the cluster's default_submit_args; pass one argv element per CLI flag. If the script to run was just uploaded via upload_file, put the `path` field from that response as one of the args — do NOT construct your own path. Cluster-local paths (e.g. /opt/spark/examples/pi.py) are accepted but will be logged as a warning. To avoid the warning, always use the `path` field returned by upload_file. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. Binary is invoked as a child process — no shell, no command injection."),
		mcp.WithString("cluster_id",
			mcp.Required(),
			mcp.Description("ID of the configured cluster (from list_clusters)"),
		),
		// mcp-go v0.56.0 uses PropertyOption for array item schema.
		mcp.WithArray("args",
			mcp.Description("Spark-submit CLI args (after default_submit_args). Each element becomes one argv entry — no shell interpretation."),
			mcp.Items(map[string]any{"type": "string"}),
		),
	)
}

// SparkSubmitHandler runs spark-submit against the requested cluster.
func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	clusterID, err := req.RequireString("cluster_id")
	if err != nil {
		return errResult("spark_submit: " + err.Error()), nil
	}

	args := req.GetArguments()
	var userArgs []string
	if rawArgs, ok := args["args"]; ok && rawArgs != nil {
		arr, ok := rawArgs.([]any)
		if !ok {
			return errResult("spark_submit: args is not an array"), nil
		}
		for i, item := range arr {
			s, ok := item.(string)
			if !ok {
				return errResult(fmt.Sprintf("spark_submit: args[%d] is not a string", i)), nil
			}
			userArgs = append(userArgs, s)
		}
	}

	callLog := startToolCall(ctx, d.Logger, SparkSubmitName, map[string]any{
		"cluster_id": clusterID,
		"args":       userArgs,
	})
	defer callLog.End()

	cl, err := d.ClusterRepo.Get(ctx, clusterID)
	if err != nil {
		callLog.WithError(err)
		return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil
	}

	fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
	fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore, d.Logger)...)

	result, err := executor.Run(ctx, executor.SparkSubmitOpts{
		Binary:  cl.SparkSubmitExecuteBin,
		Args:    fullArgs,
		Timeout: d.SparkSubmitTimeout,
	})
	if err != nil {
		callLog.WithError(err)
		// Even on error, result carries ExitCode/Stderr for the LLM.
		if result != nil {
			callLog.WithResult(map[string]any{
				"app_id":      result.AppID,
				"exit_code":   result.ExitCode,
				"duration_ms": result.DurationMS,
			})
			return textResult(encodeJSON(map[string]any{
				"app_id":      result.AppID,
				"exit_code":   result.ExitCode,
				"stdout_tail": result.StdoutTail,
				"stderr_tail": result.StderrTail,
				"duration_ms": result.DurationMS,
				"error":       err.Error(),
			})), nil
		}
		return errResult("spark_submit: " + err.Error()), nil
	}

	callLog.WithResult(map[string]any{
		"app_id":      result.AppID,
		"exit_code":   result.ExitCode,
		"duration_ms": result.DurationMS,
	})
	return textResult(encodeJSON(result)), nil
}

// translateArgs canonicalizes server-minted upload paths and leaves
// cluster-local paths untouched. It is a code-level guardrail, not a
// validator: non-minted absolute paths are intentionally accepted.
func translateArgs(userArgs []string, store *uploads.Store, logger *slog.Logger) []string {
	if store == nil {
		return userArgs
	}
	out := make([]string, len(userArgs))
	for i, arg := range userArgs {
		pathPart, isPath := pathLikeArg(arg)
		if isPath {
			if fileID, err := store.Validate(pathPart); err == nil {
				if canonical, err := store.AbsPath(fileID); err == nil {
					out[i] = canonical
					continue
				}
			} else if logger != nil && filepath.IsAbs(pathPart) {
				// Absolute path that this server did not mint. Still allowed,
				// but warn so operators can spot LLM-invented remote paths.
				logger.Warn("spark_submit.unminted_path",
					"path", pathPart,
					"note", "not produced by upload_file on this server; cluster-local paths are valid, but if this came from a remote agent the path is unreachable")
			}
		}
		out[i] = arg
	}
	return out
}

// pathLikeArg extracts an absolute path from a spark-submit argv element.
// It handles plain absolute paths ("/opt/spark/pi.py") and flag-style values
// ("--conf=/x/y"). Non-path args like "--class" or "Main" return ok=false.
func pathLikeArg(arg string) (string, bool) {
	if arg == "" {
		return "", false
	}
	if strings.HasPrefix(arg, "/") {
		return arg, true
	}
	if strings.HasPrefix(arg, "-") {
		if eq := strings.Index(arg, "="); eq > 0 && eq < len(arg)-1 {
			val := arg[eq+1:]
			if strings.HasPrefix(val, "/") {
				return val, true
			}
		}
	}
	return "", false
}