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

package tools

import (
	"context"
	"fmt"

	"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 also accepted. 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)...)

	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) []string {
	if store == nil {
		return userArgs
	}
	out := make([]string, len(userArgs))
	for i, arg := range userArgs {
		if fileID, err := store.Validate(arg); err == nil {
			if canonical, err := store.AbsPath(fileID); err == nil {
				out[i] = canonical
				continue
			}
		}
		out[i] = arg
	}
	return out
}