diff --git a/internal/mcp/tools/spark_submit.go b/internal/mcp/tools/spark_submit.go index a88ac3c..3b3883d 100644 --- a/internal/mcp/tools/spark_submit.go +++ b/internal/mcp/tools/spark_submit.go @@ -3,14 +3,12 @@ package tools import ( "context" "fmt" - "log/slog" - "path/filepath" - "strings" + "sort" + "strconv" "github.com/mark3labs/mcp-go/mcp" "spark-mcp-go/internal/executor" - "spark-mcp-go/internal/uploads" ) const SparkSubmitName = "spark_submit" @@ -18,15 +16,47 @@ 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.WithDescription("Submit a Spark application to a configured cluster. The command is built in this order: --master, --queue, --executor-memory, --executor-cores, --num-executors, then --conf per spark_conf entry, then --flag value per extra_args entry, then script_path last. For script_path, pass the `path` field returned by upload_file (do not construct your own path). 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"}), + mcp.WithString("master", + mcp.Required(), + mcp.Description("Spark master URL, e.g. yarn, k8s://https://..."), + ), + mcp.WithString("deploy_mode", + mcp.Required(), + mcp.Description("Required for forward compatibility; not currently emitted in the command by the builder (matches Python reference)."), + mcp.Enum("client", "cluster"), + ), + mcp.WithString("script_path", + mcp.Required(), + mcp.Description("Absolute path to the script. Pass the path field returned by upload_file — do not construct your own path. The script is always the last argv element."), + ), + mcp.WithString("queue", + mcp.Required(), + mcp.Description("YARN queue name"), + ), + mcp.WithString("executor_memory", + mcp.Required(), + mcp.Description("e.g. 4G"), + ), + mcp.WithNumber("executor_cores", + mcp.Required(), + mcp.Description("Cores per executor (integer)"), + ), + mcp.WithNumber("num_executors", + mcp.Required(), + mcp.Description("Static executor count (integer)"), + ), + mcp.WithObject("spark_conf", + mcp.Description("Map of key=value pairs, each emitted as --conf key=value"), + mcp.AdditionalProperties(map[string]any{"type": "string"}), + ), + mcp.WithObject("extra_args", + mcp.Description("Map of flag -> value, each emitted as --flag value"), + mcp.AdditionalProperties(map[string]any{"type": "string"}), ), ) } @@ -37,26 +67,65 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) if err != nil { return errResult("spark_submit: " + err.Error()), nil } + master, err := req.RequireString("master") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + deployMode, err := req.RequireString("deploy_mode") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + scriptPath, err := req.RequireString("script_path") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + queue, err := req.RequireString("queue") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + executorMemory, err := req.RequireString("executor_memory") + 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) - } + + // deploy_mode is required for forward compatibility but the builder does not + // emit --deploy-mode, matching the Python reference's actual behavior. If the + // Go side needs --deploy-mode later, add it here and document the divergence. + if deployMode != "client" && deployMode != "cluster" { + return errResult("spark_submit: deploy_mode must be client or cluster"), nil + } + + sparkConf, err := parseStringMap(args, "spark_conf") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + extraArgs, err := parseStringMap(args, "extra_args") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + + executorCores, err := requireNonNegativeInt(args, "executor_cores") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil + } + numExecutors, err := requireNonNegativeInt(args, "num_executors") + if err != nil { + return errResult("spark_submit: " + err.Error()), nil } callLog := startToolCall(ctx, d.Logger, SparkSubmitName, map[string]any{ - "cluster_id": clusterID, - "args": userArgs, + "cluster_id": clusterID, + "master": master, + "deploy_mode": deployMode, + "script_path": scriptPath, + "queue": queue, + "executor_memory": executorMemory, + "executor_cores": executorCores, + "num_executors": numExecutors, + "spark_conf": sparkConf, + "extra_args": extraArgs, }) defer callLog.End() @@ -66,12 +135,34 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil } - fullArgs := append([]string{}, cl.DefaultSubmitArgs...) - fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore, d.Logger)...) + binary := cl.SparkSubmitExecuteBin + + // Build argv inline, mirroring the Python reference's output order exactly. + // DefaultSubmitArgs is no longer prepended; the field is currently unused at + // runtime but kept for backward compatibility with the admin API. Removal is + // a separate follow-up. + cmd := []string{binary} + cmd = append(cmd, "--master", master) + cmd = append(cmd, "--queue", queue) + cmd = append(cmd, "--executor-memory", executorMemory) + cmd = append(cmd, "--executor-cores", strconv.Itoa(executorCores)) + cmd = append(cmd, "--num-executors", strconv.Itoa(numExecutors)) + + for _, k := range sortedStringKeys(sparkConf) { + cmd = append(cmd, "--conf", k+"="+sparkConf[k]) + } + for _, k := range sortedStringKeys(extraArgs) { + cmd = append(cmd, "--"+k, extraArgs[k]) + } + cmd = append(cmd, scriptPath) + + if d.Logger != nil { + d.Logger.Debug("spark_submit.built", "cluster_id", clusterID, "argv", cmd) + } result, err := executor.Run(ctx, executor.SparkSubmitOpts{ - Binary: cl.SparkSubmitExecuteBin, - Args: fullArgs, + Binary: binary, + Args: cmd, Timeout: d.SparkSubmitTimeout, }) if err != nil { @@ -103,52 +194,57 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) 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 +// requireNonNegativeInt extracts an integer parameter from the request and +// rejects negative or non-integer values. +func requireNonNegativeInt(args map[string]any, key string) (int, error) { + v, ok := args[key] + if !ok || v == nil { + return 0, fmt.Errorf("%s is required", key) } - 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") - } + switch n := v.(type) { + case int: + if n < 0 { + return 0, fmt.Errorf("%s must be a non-negative integer", key) } - out[i] = arg + return n, nil + case float64: + ni := int(n) + if n < 0 || float64(ni) != n { + return 0, fmt.Errorf("%s must be a non-negative integer", key) + } + return ni, nil + default: + return 0, fmt.Errorf("%s must be a number", key) } - 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 +// parseStringMap extracts an optional object whose values are strings. +func parseStringMap(args map[string]any, key string) (map[string]string, error) { + out := make(map[string]string) + v, ok := args[key] + if !ok || v == nil { + return out, nil } - if strings.HasPrefix(arg, "/") { - return arg, true + m, ok := v.(map[string]any) + if !ok { + return nil, fmt.Errorf("%s must be an object", key) } - 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 - } + for k, val := range m { + s, ok := val.(string) + if !ok { + return nil, fmt.Errorf("%s[%q] must be a string", key, k) } + out[k] = s } - return "", false + return out, nil +} + +// sortedStringKeys returns the keys of m in deterministic order. +func sortedStringKeys(m map[string]string) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys } diff --git a/internal/mcp/tools/spark_submit_path_test.go b/internal/mcp/tools/spark_submit_path_test.go index 1eb26fa..e8cc638 100644 --- a/internal/mcp/tools/spark_submit_path_test.go +++ b/internal/mcp/tools/spark_submit_path_test.go @@ -1,20 +1,19 @@ package tools import ( - "bytes" "context" - "encoding/json" - "log/slog" "os" "path/filepath" "strings" "testing" "time" + "github.com/mark3labs/mcp-go/mcp" + "spark-mcp-go/internal/cluster" ) -func TestSparkSubmit_PathPassthrough(t *testing.T) { +func TestSparkSubmit_StructuredCommand(t *testing.T) { deps, repo := testDepsWithDataDir(t) deps.SparkSubmitTimeout = 5 * time.Second store := deps.UploadStore @@ -41,8 +40,16 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) { }) req := newToolRequest(SparkSubmitName, map[string]any{ - "cluster_id": "cluster-echo", - "args": []any{absPath, "--class", "Main"}, + "cluster_id": "cluster-echo", + "master": "yarn", + "deploy_mode": "cluster", + "script_path": absPath, + "queue": "default", + "executor_memory": "2G", + "executor_cores": 1, + "num_executors": 4, + "spark_conf": map[string]any{"a": "1", "b": "2"}, + "extra_args": map[string]any{"name": "wordcount-job"}, }) res, err := deps.SparkSubmitHandler(context.Background(), req) if err != nil { @@ -57,9 +64,21 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) { t.Fatalf("read echoed args: %v", err) } lines := strings.Split(strings.TrimSpace(string(got)), "\n") - want := []string{absPath, "--class", "Main"} + + want := []string{ + echoScript, + "--master", "yarn", + "--queue", "default", + "--executor-memory", "2G", + "--executor-cores", "1", + "--num-executors", "4", + "--conf", "a=1", + "--conf", "b=2", + "--name", "wordcount-job", + absPath, + } if len(lines) != len(want) { - t.Fatalf("echoed lines=%v, want %v", lines, want) + t.Fatalf("lines=%v\nwant=%v", lines, want) } for i, l := range lines { if l != want[i] { @@ -67,148 +86,69 @@ func TestSparkSubmit_PathPassthrough(t *testing.T) { } } - // Sanity: the canonical path still contains the minted file ID. + if lines[len(lines)-1] != absPath { + t.Errorf("last line=%q, want script_path %q", lines[len(lines)-1], absPath) + } + if !strings.HasSuffix(absPath, "/"+fileID) { t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID) } } -func TestTranslateArgs_Logging(t *testing.T) { +func TestSparkSubmit_MissingRequiredField(t *testing.T) { deps, _ := testDepsWithDataDir(t) - t.Run("minted_path_no_warn", func(t *testing.T) { - store := deps.UploadStore - fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py") - if err != nil { - t.Fatalf("save upload: %v", err) - } - - logger, buf := capturingLogger(t) - out := translateArgs([]string{absPath, "--class", "Main"}, deps.UploadStore, logger) - if out[0] != absPath { - t.Errorf("out[0]=%q, want %q", out[0], absPath) - } - - for _, rec := range parseLogRecords(t, buf) { - if rec["msg"] == "spark_submit.unminted_path" { - t.Errorf("unexpected warn for minted path: %v", rec) - } - } - - if !strings.HasSuffix(absPath, "/"+fileID) { - t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID) - } - }) - - t.Run("cluster_local_path_warns", func(t *testing.T) { - logger, buf := capturingLogger(t) - out := translateArgs([]string{"/opt/spark/examples/pi.py"}, deps.UploadStore, logger) - if len(out) != 1 || out[0] != "/opt/spark/examples/pi.py" { - t.Errorf("out=%v, want [/opt/spark/examples/pi.py]", out) - } - - var warns int - for _, rec := range parseLogRecords(t, buf) { - if rec["msg"] == "spark_submit.unminted_path" { - warns++ - if rec["path"] != "/opt/spark/examples/pi.py" { - t.Errorf("warn path=%v, want /opt/spark/examples/pi.py", rec["path"]) - } - } - } - if warns != 1 { - t.Errorf("warns=%d, want 1; log=%s", warns, buf.String()) - } - }) - - t.Run("non_path_args_never_warn", func(t *testing.T) { - logger, buf := capturingLogger(t) - out := translateArgs([]string{"--class", "Main", "--master", "yarn"}, deps.UploadStore, logger) - want := []string{"--class", "Main", "--master", "yarn"} - if len(out) != len(want) { - t.Fatalf("out=%v, want %v", out, want) - } - for i := range out { - if out[i] != want[i] { - t.Errorf("out[%d]=%q, want %q", i, out[i], want[i]) - } - } - - for _, rec := range parseLogRecords(t, buf) { - if rec["msg"] == "spark_submit.unminted_path" { - t.Errorf("unexpected warn for non-path args: %v", rec) - } - } - }) -} - -func capturingLogger(t *testing.T) (*slog.Logger, *bytes.Buffer) { - t.Helper() - buf := &bytes.Buffer{} - return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug})), buf -} - -func parseLogRecords(t *testing.T, buf *bytes.Buffer) []map[string]any { - t.Helper() - var recs []map[string]any - for _, line := range strings.Split(strings.TrimSpace(buf.String()), "\n") { - if line == "" { - continue - } - var rec map[string]any - if err := json.Unmarshal([]byte(line), &rec); err != nil { - t.Fatalf("invalid log JSON: %q (err=%v)", line, err) - } - recs = append(recs, rec) - } - return recs -} - -func TestSparkSubmit_ClusterLocalPathPassthrough(t *testing.T) { - deps, repo := testDepsWithDataDir(t) - deps.SparkSubmitTimeout = 5 * time.Second - - echoed := filepath.Join(t.TempDir(), "echoed") - echoScript := filepath.Join(t.TempDir(), "echo-args.sh") - if err := os.WriteFile(echoScript, []byte("#!/bin/sh\nfor arg do\n echo \"$arg\" >> \"$ECHO_FILE\"\ndone\n"), 0o755); err != nil { - t.Fatalf("write echo script: %v", err) - } - - t.Setenv("ECHO_FILE", echoed) - - createCluster(t, repo, &cluster.Cluster{ - ID: "cluster-echo", - Name: "Echo", - IsActive: true, - AuthType: cluster.AuthNone, - SparkSubmitExecuteBin: echoScript, - }) - - clusterLocalPath := "/usr/bin/env" req := newToolRequest(SparkSubmitName, map[string]any{ - "cluster_id": "cluster-echo", - "args": []any{clusterLocalPath, "--class", "Main"}, + "cluster_id": "cluster-echo", + "deploy_mode": "cluster", + "script_path": "/tmp/script.py", + "queue": "default", + "executor_memory": "2G", + "executor_cores": 1, + "num_executors": 4, }) res, err := deps.SparkSubmitHandler(context.Background(), req) if err != nil { t.Fatalf("handler error: %v", err) } - if res.IsError { - t.Fatalf("unexpected error result: %v", res.Content) + if !res.IsError { + t.Fatalf("expected error result, got: %v", res.Content) } - - got, err := os.ReadFile(echoed) - if err != nil { - t.Fatalf("read echoed args: %v", err) + text, ok := mcp.AsTextContent(res.Content[0]) + if !ok { + t.Fatalf("content is not text: %T", res.Content[0]) } - lines := strings.Split(strings.TrimSpace(string(got)), "\n") - want := []string{clusterLocalPath, "--class", "Main"} - if len(lines) != len(want) { - t.Fatalf("echoed lines=%v, want %v", lines, want) - } - for i, l := range lines { - if l != want[i] { - t.Errorf("line[%d]=%q, want %q", i, l, want[i]) - } + if !strings.Contains(text.Text, "master") { + t.Errorf("error text=%q, want mention of master", text.Text) + } +} + +func TestSparkSubmit_BadSparkConfValue(t *testing.T) { + deps, _ := testDepsWithDataDir(t) + + req := newToolRequest(SparkSubmitName, map[string]any{ + "cluster_id": "cluster-echo", + "master": "yarn", + "deploy_mode": "cluster", + "script_path": "/tmp/script.py", + "queue": "default", + "executor_memory": "2G", + "executor_cores": 1, + "num_executors": 4, + "spark_conf": map[string]any{"a": 1}, + }) + res, err := deps.SparkSubmitHandler(context.Background(), req) + if err != nil { + t.Fatalf("handler error: %v", err) + } + if !res.IsError { + t.Fatalf("expected error result, got: %v", res.Content) + } + text, ok := mcp.AsTextContent(res.Content[0]) + if !ok { + t.Fatalf("content is not text: %T", res.Content[0]) + } + if !strings.Contains(text.Text, "spark_conf") { + t.Errorf("error text=%q, want mention of spark_conf", text.Text) } }