Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
97f6184f65 | ||
|
|
dab4cd062e |
@@ -5,6 +5,7 @@ package admin
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
@@ -170,6 +171,10 @@ func updateCluster(repo *storage.ClusterRepo, auditRepo *audit.Repo) gin.Handler
|
||||
merged.DefaultSubmitArgs = decodeStringSlice(v)
|
||||
}
|
||||
|
||||
if len(merged.DefaultSubmitArgs) > 0 {
|
||||
slog.Default().Info("admin: cluster sets default_submit_args which is deprecated; the spark_submit Tool no longer prepends these args. Configure flags via spark_conf / extra_args in the structured tool call.", "cluster_id", merged.ID)
|
||||
}
|
||||
|
||||
// Re-validate after merging so a missing auth_type on a fresh
|
||||
// cluster (defaulted by DB to "none") doesn't get clobbered.
|
||||
if err := validateClusterUpdate(&merged); err != nil {
|
||||
|
||||
@@ -34,7 +34,11 @@ const (
|
||||
// password" — keeps the existing encrypted blob intact.
|
||||
// - URLAllowlist: extra glob patterns beyond the RM/SHS hosts; fetch_url
|
||||
// uses this to permit long-tail SHS endpoints the agent may construct.
|
||||
// - DefaultSubmitArgs: prepended to every spark_submit Tool call's args.
|
||||
// - DefaultSubmitArgs: deprecated; no longer prepended by spark_submit
|
||||
// (post-7920dc9 the builder constructs argv from structured fields).
|
||||
// Kept for backward compatibility with the admin API and DB schema;
|
||||
// actual removal is a follow-up. New cluster configurations should
|
||||
// leave this empty.
|
||||
// - RateLimitPerMin: 0 disables the per-cluster limiter.
|
||||
type Cluster struct {
|
||||
ID string `json:"id"`
|
||||
@@ -49,6 +53,7 @@ type Cluster struct {
|
||||
SSLVerify bool `json:"ssl_verify"`
|
||||
SSLCABundle string `json:"ssl_ca_bundle,omitempty"`
|
||||
URLAllowlist []string `json:"url_allowlist,omitempty"`
|
||||
// Deprecated: see godoc.
|
||||
DefaultSubmitArgs []string `json:"default_submit_args,omitempty"`
|
||||
RateLimitPerMin int `json:"rate_limit_per_min"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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=[<path>, '--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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user