tao.chen/spark-mcp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

uploads,tools,config,main: 将 upload_file 返回路径改为服务器端绝对路径并透传给 spark_submit

Browse Source 
之前 upload_file 把文件写到 DataDir/uploads/<filename>，返回相对路径
uploads/hello.txt。MCP 服务器与 agent 不在同一台机器，agent 无法构造服务
器本地路径，因此 spark_submit 无法定位脚本。

改为由新的 internal/uploads 包 mint 一个 32 字符十六进制 file_id，文件落盘为
DataDir/uploads/<file_id>，元数据写入 <file_id>.meta.json（原始文件名只存在
sidecar 里）。upload_file 返回 {file_id, path, name, size, sha256}，其中 path
是绝对路径。spark_submit 的 description 明确要求 LLM 直接把 upload_file 返回
的 path 放进 args，不要自己构造路径。

为什么只在描述里约束而不在代码里拒绝非 mint 的绝对路径：集群本地已有路径
（如 /opt/spark/examples/pi.py）是合法的 spark-submit 参数，代码不能替
LLM 拒绝。

测试锁定：
- internal/uploads: Save 往返、AbsPath/Validate 非法路径、Sweep 过期/未过期/
  孤立 sidecar
- internal/mcp/tools: upload_file 新响应字段、spark_submit 透传 mint 路径与
  集群本地路径

刻意未做：S3/HDFS 上传、给 spark_submit 新增 file_id 参数、在代码层面拒绝
集群本地绝对路径。

破坏性变更：upload_file 响应从相对 path 改为绝对 path，并新增 file_id/name/
sha256 字段。该工具尚无外部调用者。

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
2026-07-13 19:13:14 +08:00
co-authored by Claude
parent 210aeb1315
commit f374ebfdb4
10 changed files with 748 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/mcp/tools/deps.go
+2
View File
@@ -7,6 +7,7 @@ import (
"spark-mcp-go/internal/analyzer"
"spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads"
)
// Deps bundles the dependencies shared by all MCP Tool handlers.
@@ -17,6 +18,7 @@ type Deps struct {
HTTPClient *httpclient.Client
MaxResponseBytes int64
DataDir string // upload_file writes to DataDir/uploads
UploadStore *uploads.Store
AnalyzerThresholds analyzer.Thresholds
}
internal/mcp/tools/fetch_upload_test.go
+39 -9
View File
@@ -8,6 +8,7 @@ import (
"net/http/httptest"
"os"
"path/filepath"
"regexp"
"strconv"
"strings"
"testing"
@@ -18,6 +19,7 @@ import (
"spark-mcp-go/internal/cluster"
"spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/storage"
"spark-mcp-go/internal/uploads"
)
// testDepsWithDataDir returns dependencies backed by an in-memory DB and a
@@ -30,6 +32,12 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
}
t.Cleanup(func() { _ = db.Close() })
dataDir := t.TempDir()
uploadStore, err := uploads.New(filepath.Join(dataDir, "uploads"))
if err != nil {
t.Fatalf("new upload store: %v", err)
}
return &Deps{
HTTPClient: httpclient.New(httpclient.Config{
Timeout: 5 * time.Second,
@@ -37,7 +45,8 @@ func testDepsWithDataDir(t *testing.T) (*Deps, *storage.ClusterRepo) {
}),
ClusterRepo: db.Clusters(),
MaxResponseBytes: 1 << 20,
DataDir: t.TempDir(),
DataDir: dataDir,
UploadStore: uploadStore,
}, db.Clusters()
}
@@ -324,15 +333,26 @@ func TestUploadFile(t *testing.T) {
}
payload := resultJSON(t, res)
if payload["path"] != "uploads/hello.txt" {
t.Errorf("path=%v, want uploads/hello.txt", payload["path"])
fileID, ok := payload["file_id"].(string)
if !ok || !regexp.MustCompile(`^[0-9a-f]{32}$`).MatchString(fileID) {
t.Errorf("file_id=%v, want 32 lowercase hex chars", payload["file_id"])
}
path, ok := payload["path"].(string)
if !ok || !filepath.IsAbs(path) || !strings.HasSuffix(path, "/"+fileID) {
t.Errorf("path=%v, want absolute path ending in /%s", payload["path"], fileID)
}
if payload["name"] != "hello.txt" {
t.Errorf("name=%v, want hello.txt", payload["name"])
}
if payload["size"] != float64(len("hello world")) {
t.Errorf("size=%v, want %d", payload["size"], len("hello world"))
}
sha, ok := payload["sha256"].(string)
if !ok || !regexp.MustCompile(`^[0-9a-f]{64}$`).MatchString(sha) {
t.Errorf("sha256=%v, want 64-char hex", payload["sha256"])
}
final := filepath.Join(deps.DataDir, "uploads", "hello.txt")
got, err := os.ReadFile(final)
got, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read uploaded file: %v", err)
}
@@ -340,7 +360,7 @@ func TestUploadFile(t *testing.T) {
t.Errorf("content=%q, want %q", got, "hello world")
}
info, err := os.Stat(final)
info, err := os.Stat(path)
if err != nil {
t.Fatalf("stat uploaded file: %v", err)
}
@@ -393,9 +413,19 @@ func TestUploadFile_Base64(t *testing.T) {
t.Fatalf("handler error: %v", err)
}
_ = resultJSON(t, res)
final := filepath.Join(deps.DataDir, "uploads", "hello.bin")
got, err := os.ReadFile(final)
payload := resultJSON(t, res)
path, ok := payload["path"].(string)
if !ok || !filepath.IsAbs(path) {
t.Errorf("path=%v, want absolute path", payload["path"])
}
if payload["name"] != "hello.bin" {
t.Errorf("name=%v, want hello.bin", payload["name"])
}
if payload["size"] != float64(len("hello")) {
t.Errorf("size=%v, want %d", payload["size"], len("hello"))
}
got, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read uploaded file: %v", err)
}
internal/mcp/tools/spark_submit.go
+23 -2
View File
@@ -7,6 +7,7 @@ import (
"github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/executor"
"spark-mcp-go/internal/uploads"
)
const SparkSubmitName = "spark_submit"
@@ -14,7 +15,7 @@ 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. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. The cluster's default_submit_args are prepended to your args. Binary is invoked as a child process — no shell, no command injection."),
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)"),
@@ -63,7 +64,7 @@ func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest)
}
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
fullArgs = append(fullArgs, userArgs...)
fullArgs = append(fullArgs, translateArgs(userArgs, d.UploadStore)...)
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
Binary: cl.SparkSubmitExecuteBin,
@@ -98,3 +99,23 @@ 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) []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
}
internal/mcp/tools/spark_submit_path_test.go
+120
View File
@@ -0,0 +1,120 @@
package tools
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"time"
"spark-mcp-go/internal/cluster"
)
func TestSparkSubmit_PathPassthrough(t *testing.T) {
deps, repo := testDepsWithDataDir(t)
deps.SparkSubmitTimeout = 5 * time.Second
store := deps.UploadStore
fileID, _, _, _, absPath, err := store.Save([]byte("print('hi')\n"), "hello.py")
if err != nil {
t.Fatalf("save upload: %v", err)
}
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,
})
req := newToolRequest(SparkSubmitName, map[string]any{
"cluster_id": "cluster-echo",
"args": []any{absPath, "--class", "Main"},
})
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)
}
got, err := os.ReadFile(echoed)
if err != nil {
t.Fatalf("read echoed args: %v", err)
}
lines := strings.Split(strings.TrimSpace(string(got)), "\n")
want := []string{absPath, "--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])
}
}
// Sanity: the canonical path still contains the minted file ID.
if !strings.HasSuffix(absPath, "/"+fileID) {
t.Errorf("absPath=%q does not end with fileID %q", absPath, fileID)
}
}
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"},
})
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)
}
got, err := os.ReadFile(echoed)
if err != nil {
t.Fatalf("read echoed args: %v", err)
}
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])
}
}
}
internal/mcp/tools/upload_file.go
+13 -10
View File
@@ -4,7 +4,6 @@ import (
"context"
"encoding/base64"
"fmt"
"os"
"path/filepath"
"regexp"
@@ -19,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 ./data/uploads/ so it can be referenced by spark_submit later."),
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.WithString("filename",
mcp.Required(),
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
@@ -60,6 +59,10 @@ 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":
@@ -72,17 +75,17 @@ func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (
data = []byte(content)
}
final := filepath.Join(d.DataDir, "uploads", filename)
if err := os.MkdirAll(filepath.Dir(final), 0o750); err != nil {
return errResult("upload_file: create uploads dir: " + err.Error()), nil
}
if err := os.WriteFile(final, data, 0o640); err != nil {
return errResult("upload_file: write file: " + err.Error()), nil
fileID, _, size, sha256Hex, absPath, err := d.UploadStore.Save(data, filename)
if err != nil {
return errResult("upload_file: save upload: " + err.Error()), nil
}
result := map[string]any{
"path": fmt.Sprintf("uploads/%s", filename),
"size": len(data),
"file_id": fileID,
"path": absPath,
"name": filename,
"size": size,
"sha256": sha256Hex,
}
return textResult(encodeJSON(result)), nil
}