Files
spark-mcp/internal/mcp/tools/upload_file.go
T
tao.chenandClaude f374ebfdb4 uploads,tools,config,main: 将 upload_file 返回路径改为服务器端绝对路径并透传给 spark_submit
之前 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>
2026-07-13 19:13:14 +08:00

108 lines
3.1 KiB
Go

package tools
import (
"context"
"encoding/base64"
"fmt"
"path/filepath"
"regexp"
"github.com/mark3labs/mcp-go/mcp"
)
const UploadFileName = "upload_file"
// filenameRegex restricts upload names to a safe, portable character set.
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.WithString("filename",
mcp.Required(),
mcp.Description("Plain file name without path separators (1-128 chars, [a-zA-Z0-9._-])"),
),
mcp.WithString("content",
mcp.Required(),
mcp.Description("File contents; text or base64-encoded binary"),
),
mcp.WithString("encoding",
mcp.Description("Encoding of content"),
mcp.Enum("text", "base64"),
mcp.DefaultString("text"),
),
)
}
// UploadFileHandler writes user-provided content to DataDir/uploads/filename.
func (d *Deps) UploadFileHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
filename, err := req.RequireString("filename")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
content, err := req.RequireString("content")
if err != nil {
return errResult("upload_file: " + err.Error()), nil
}
args := req.GetArguments()
encoding := "text"
if v, ok := args["encoding"].(string); ok && v != "" {
encoding = v
}
if encoding != "text" && encoding != "base64" {
return errResult(fmt.Sprintf("upload_file: invalid encoding %q", encoding)), nil
}
if err := validateUploadFilename(filename); err != nil {
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":
decoded, err := base64.StdEncoding.DecodeString(content)
if err != nil {
return errResult("upload_file: decode base64: " + err.Error()), nil
}
data = decoded
default:
data = []byte(content)
}
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{
"file_id": fileID,
"path": absPath,
"name": filename,
"size": size,
"sha256": sha256Hex,
}
return textResult(encodeJSON(result)), nil
}
func validateUploadFilename(filename string) error {
if len(filename) == 0 || len(filename) > 128 {
return fmt.Errorf("filename length must be 1-128")
}
if filepath.Base(filename) != filename {
return fmt.Errorf("filename must not contain path separators or '..'")
}
if filename == "." || filename == ".." {
return fmt.Errorf("filename must not be '.' or '..'")
}
if !filenameRegex.MatchString(filename) {
return fmt.Errorf("filename contains invalid characters")
}
return nil
}