diff --git a/internal/mcp/tools/deps.go b/internal/mcp/tools/deps.go index 8e230b0..aa4bb17 100644 --- a/internal/mcp/tools/deps.go +++ b/internal/mcp/tools/deps.go @@ -18,7 +18,7 @@ type Deps struct { HTTPClient *httpclient.Client MaxResponseBytes int64 DataDir string // upload_file writes to DataDir/uploads - UploadStore *uploads.Store + UploadStore uploads.Store AnalyzerThresholds analyzer.Thresholds } diff --git a/internal/mcp/tools/spark_submit.go b/internal/mcp/tools/spark_submit.go index 33c83d7..c9e6489 100644 --- a/internal/mcp/tools/spark_submit.go +++ b/internal/mcp/tools/spark_submit.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "log/slog" "sort" "strconv" @@ -63,11 +64,13 @@ func NewSparkSubmitTool() mcp.Tool { // SparkSubmitHandler runs spark-submit against the requested cluster. func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + args := req.GetArguments() + clusterID, err := req.RequireString("cluster_id") if err != nil { return errResult("spark_submit: " + err.Error()), nil } - master, err := req.RequireString("master") + master, err := requireNonEmptyString(args, "master") if err != nil { return errResult("spark_submit: " + err.Error()), nil } @@ -80,27 +83,26 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) return errResult("spark_submit: " + err.Error()), nil } - executorMemory, err := req.RequireString("executor_memory") + queue, err := requireNonEmptyString(args, "queue") if err != nil { return errResult("spark_submit: " + err.Error()), nil } - queue, err := req.RequireString("queue") + executorMemory, err := requireNonEmptyString(args, "executor_memory") if err != nil { return errResult("spark_submit: " + err.Error()), nil } - args := req.GetArguments() - // Reject paths that didn't come through upload_file on this server. // Cluster-local paths and arbitrary host paths are unreachable from the MCP // server's spark-submit, so we fail fast with a clear message rather than // letting spark-submit produce a confusing FileNotFoundException later. // Placed after all required-string parses so error messages surface in // field order (cluster_id, master, ..., script_path, queue, ...). - if d.UploadStore != nil { - if _, err := d.UploadStore.Validate(scriptPath); err != nil { - return errResult(fmt.Sprintf("spark_submit: script_path %q is not from upload_file on this server; upload the file first via upload_file and pass back the path field", scriptPath)), nil + if _, err := d.UploadStore.Validate(scriptPath); err != nil { + if d.Logger != nil { + d.Logger.Warn("spark_submit.unminted_path_rejected", slog.String("path", scriptPath), slog.String("err", err.Error())) } + return errResult(fmt.Sprintf("spark_submit: script_path %q is not from upload_file on this server; upload the file first via upload_file and pass back the path field", scriptPath)), nil } // deploy_mode is required for forward compatibility but the builder does not @@ -149,25 +151,16 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) } 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) + cmd := buildSparkSubmitCommand(SparkSubmitCommandOpts{ + Master: master, + Queue: queue, + ExecutorMemory: executorMemory, + ExecutorCores: executorCores, + NumExecutors: numExecutors, + SparkConf: sparkConf, + ExtraArgs: extraArgs, + ScriptPath: scriptPath, + }) if d.Logger != nil { d.Logger.Debug("spark_submit.built", "cluster_id", clusterID, "argv", cmd) @@ -207,6 +200,40 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) return textResult(encodeJSON(result)), nil } +// SparkSubmitCommandOpts holds the structured arguments used to build the +// spark-submit argv. The returned argv does NOT include the binary itself. +type SparkSubmitCommandOpts struct { + Master string + Queue string + ExecutorMemory string + ExecutorCores int + NumExecutors int + SparkConf map[string]string + ExtraArgs map[string]string + ScriptPath string +} + +// buildSparkSubmitCommand builds the post-binary argv for spark-submit. +// 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. +func buildSparkSubmitCommand(opts SparkSubmitCommandOpts) []string { + cmd := []string{"--master", opts.Master} + cmd = append(cmd, "--queue", opts.Queue) + cmd = append(cmd, "--executor-memory", opts.ExecutorMemory) + cmd = append(cmd, "--executor-cores", strconv.Itoa(opts.ExecutorCores)) + cmd = append(cmd, "--num-executors", strconv.Itoa(opts.NumExecutors)) + + for _, k := range sortedStringKeys(opts.SparkConf) { + cmd = append(cmd, "--conf", k+"="+opts.SparkConf[k]) + } + for _, k := range sortedStringKeys(opts.ExtraArgs) { + cmd = append(cmd, "--"+k, opts.ExtraArgs[k]) + } + cmd = append(cmd, opts.ScriptPath) + return cmd +} + // 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) { @@ -261,3 +288,17 @@ func sortedStringKeys(m map[string]string) []string { sort.Strings(keys) return keys } + +// requireNonEmptyString extracts a required string parameter and rejects empty +// values. MCP's RequireString already enforces presence/type; this enforces +// that the field is not blank. +func requireNonEmptyString(args map[string]any, key string) (string, error) { + s, ok := args[key].(string) + if !ok { + return "", fmt.Errorf("%s is required", key) + } + if s == "" { + return "", fmt.Errorf("%s is required", key) + } + return s, nil +} diff --git a/internal/mcp/tools/spark_submit_path_test.go b/internal/mcp/tools/spark_submit_path_test.go index 5cdfff9..647cd86 100644 --- a/internal/mcp/tools/spark_submit_path_test.go +++ b/internal/mcp/tools/spark_submit_path_test.go @@ -66,7 +66,6 @@ func TestSparkSubmit_StructuredCommand(t *testing.T) { lines := strings.Split(strings.TrimSpace(string(got)), "\n") want := []string{ - echoScript, "--master", "yarn", "--queue", "default", "--executor-memory", "2G", @@ -80,6 +79,9 @@ func TestSparkSubmit_StructuredCommand(t *testing.T) { if len(lines) != len(want) { t.Fatalf("lines=%v\nwant=%v", lines, want) } + if lines[0] != "--master" { + t.Errorf("line[0]=%q, want \"--master\"", lines[0]) + } for i, l := range lines { if l != want[i] { t.Errorf("line[%d]=%q, want %q", i, l, want[i]) @@ -196,3 +198,37 @@ func TestSparkSubmit_RejectsNonMintedPath(t *testing.T) { t.Errorf("error text=%q, want mention of not from upload_file", text.Text) } } + +func TestSparkSubmit_EmptyMaster(t *testing.T) { + deps, _ := testDepsWithDataDir(t) + store := deps.UploadStore + _, _, _, _, mintedPath, err := store.Save([]byte("# dummy\n"), "dummy.py") + if err != nil { + t.Fatalf("save upload: %v", err) + } + + req := newToolRequest(SparkSubmitName, map[string]any{ + "cluster_id": "cluster-echo", + "master": "", + "deploy_mode": "cluster", + "script_path": mintedPath, + "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("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, "master is required") && !strings.Contains(text.Text, "master") { + t.Errorf("error text=%q, want mention of master", text.Text) + } +} diff --git a/internal/mcp/tools/upload_file.go b/internal/mcp/tools/upload_file.go index 0f956ee..ff94791 100644 --- a/internal/mcp/tools/upload_file.go +++ b/internal/mcp/tools/upload_file.go @@ -18,7 +18,7 @@ var filenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`) // NewUploadFileTool returns the schema for the upload_file MCP Tool. func NewUploadFileTool() mcp.Tool { return mcp.NewTool(UploadFileName, - mcp.WithDescription("Upload a script or config file to the server's data/uploads/ directory. Returns {file_id, path, name, size, sha256}. To submit it, pass the returned `path` as one of the args in spark_submit (e.g. args=[, '--class', 'Main']). The path is server-side; the agent does not need to construct one."), + mcp.WithDescription("Upload a script or config file to the server's data/uploads/ directory. Returns {file_id, path, name, size, sha256}. To submit it, pass the returned `path` as the `script_path` field in spark_submit, along with the structured fields (master, deploy_mode, queue, executor_memory, executor_cores, num_executors; optional spark_conf, extra_args)."), mcp.WithString("filename", mcp.Required(), mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"), @@ -59,10 +59,6 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) ( return errResult("upload_file: " + err.Error()), nil } - if d.UploadStore == nil { - return errResult("upload_file: upload store not configured"), nil - } - var data []byte switch encoding { case "base64": diff --git a/internal/uploads/store.go b/internal/uploads/store.go index a59f7e6..13e1bea 100644 --- a/internal/uploads/store.go +++ b/internal/uploads/store.go @@ -25,31 +25,31 @@ type Store struct { // New creates a Store rooted at root. It creates root with mode 0o750 if it // does not exist. The root must be an absolute path. -func New(root string) (*Store, error) { +func New(root string) (Store, error) { if root == "" { - return nil, fmt.Errorf("uploads: root must not be empty") + return Store{}, fmt.Errorf("uploads: root must not be empty") } absRoot, err := filepath.Abs(root) if err != nil { - return nil, fmt.Errorf("uploads: resolve root: %w", err) + return Store{}, fmt.Errorf("uploads: resolve root: %w", err) } if !filepath.IsAbs(absRoot) { - return nil, fmt.Errorf("uploads: root %q is not an absolute path", absRoot) + return Store{}, fmt.Errorf("uploads: root %q is not an absolute path", absRoot) } if err := os.MkdirAll(absRoot, 0o750); err != nil { - return nil, fmt.Errorf("uploads: create root: %w", err) + return Store{}, fmt.Errorf("uploads: create root: %w", err) } info, err := os.Stat(absRoot) if err != nil { - return nil, fmt.Errorf("uploads: stat root: %w", err) + return Store{}, fmt.Errorf("uploads: stat root: %w", err) } if !info.IsDir() { - return nil, fmt.Errorf("uploads: root %q is not a directory", absRoot) + return Store{}, fmt.Errorf("uploads: root %q is not a directory", absRoot) } - return &Store{Root: absRoot}, nil + return Store{Root: absRoot}, nil } type sidecar struct { @@ -137,6 +137,10 @@ func (s *Store) Validate(absPath string) (string, error) { return "", fmt.Errorf("uploads: path %q is not under Root", absPath) } + if strings.HasSuffix(rel, ".meta.json") { + return "", fmt.Errorf("uploads: path %q is a sidecar, pass the data file path", absPath) + } + fileID := strings.TrimSuffix(rel, ".meta.json") if !fileIDRegex.MatchString(fileID) { return "", fmt.Errorf("uploads: path %q does not denote a valid upload", absPath) @@ -198,13 +202,16 @@ func (s *Store) Sweep(ttl time.Duration) (deleted int, err error) { if it.data { expired := false + corruptSidecar := false if it.meta { uploadedAt, parseErr := readUploadedAt(metaPath) if parseErr == nil && uploadedAt.Before(cutoff) { expired = true + } else if parseErr != nil { + corruptSidecar = true } } - if !expired && !it.meta { + if !expired && (corruptSidecar || !it.meta) { info, statErr := os.Stat(dataPath) if statErr == nil && info.ModTime().Before(cutoff) { expired = true diff --git a/internal/uploads/store_test.go b/internal/uploads/store_test.go index 893c40e..336bf4e 100644 --- a/internal/uploads/store_test.go +++ b/internal/uploads/store_test.go @@ -256,3 +256,61 @@ func TestStore_Sweep_OrphanSidecarReaped(t *testing.T) { func timePtr(t time.Time) *time.Time { return &t } + +func TestStore_Validate_RejectsSidecarPath(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + + fileID := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + if err := os.WriteFile(filepath.Join(store.Root, fileID), []byte("x"), 0o640); err != nil { + t.Fatalf("create data file: %v", err) + } + if err := os.WriteFile(filepath.Join(store.Root, fileID+".meta.json"), []byte(`{"uploaded_at":"`+time.Now().Format(time.RFC3339Nano)+`"}`), 0o600); err != nil { + t.Fatalf("create sidecar: %v", err) + } + + _, err = store.Validate(filepath.Join(store.Root, fileID+".meta.json")) + if err == nil { + t.Errorf("Validate(sidecar) succeeded, want error") + } +} + +func TestStore_Sweep_DeletesDataWithCorruptSidecar(t *testing.T) { + store, err := New(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + + id := "00000000000000000000000000000005" + dataPath := filepath.Join(store.Root, id) + metaPath := dataPath + ".meta.json" + + if err := os.WriteFile(dataPath, []byte("stale data"), 0o640); err != nil { + t.Fatalf("create data file: %v", err) + } + // Corrupt sidecar: not valid JSON. + if err := os.WriteFile(metaPath, []byte("not valid json"), 0o600); err != nil { + t.Fatalf("create sidecar: %v", err) + } + // Set mtime well in the past so the fallback triggers. + past := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(dataPath, past, past); err != nil { + t.Fatalf("set mtime: %v", err) + } + + deleted, err := store.Sweep(1 * time.Hour) + if err != nil { + t.Fatalf("sweep: %v", err) + } + if deleted != 1 { + t.Errorf("deleted=%d, want 1", deleted) + } + if _, err := os.Stat(dataPath); !os.IsNotExist(err) { + t.Errorf("data file with corrupt sidecar still exists") + } + if _, err := os.Stat(metaPath); !os.IsNotExist(err) { + t.Errorf("corrupt sidecar still exists") + } +}