tao.chenandClaude
f374ebfdb4
之前 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>
121 lines
3.3 KiB
Go
121 lines
3.3 KiB
Go
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])
|
|
}
|
|
}
|
|
}
|