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

Phase 5 Batch 1: executor + list_clusters + spark_submit + /mcp 路由

Browse Source 
- internal/executor: spark-submit 进程执行, exec.Command 切片形式
  (绝不 shell) + stdout 10MiB tail-cap + app_id regex
  + 5 个单测覆盖 slice-form 注入防御 + 退出码
- internal/mcp/server.go: NewServer + Streamable HTTP Handler
- internal/mcp/tools/{list_clusters, spark_submit, deps, helpers}.go
- mcp-go API 实测修正: WithArray 用 Items() 声明 item 类型;
  StreamableHTTPServer 直接实现 http.Handler, 无 .Handler() 方法
- main.go: 挂 /mcp 路由 (AgentAuth + mcpHandler 手动组合)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
tao.chen
2026-07-10 16:50:02 +08:00
co-authored by Claude
parent 2c39091706
commit 705665745e
5 changed files with 325 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
go.mod
+5
View File
@@ -4,6 +4,7 @@ go 1.25.5
require (
github.com/gin-gonic/gin v1.12.0
github.com/mark3labs/mcp-go v0.56.0
modernc.org/sqlite v1.53.0
)
@@ -20,6 +21,7 @@ require (
github.com/go-playground/validator/v10 v10.30.1 // indirect
github.com/goccy/go-json v0.10.5 // indirect
github.com/goccy/go-yaml v1.19.2 // indirect
github.com/google/jsonschema-go v0.4.2 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
@@ -32,8 +34,11 @@ require (
github.com/quic-go/qpack v0.6.0 // indirect
github.com/quic-go/quic-go v0.59.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect
github.com/spf13/cast v1.7.1 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.3.1 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
go.mongodb.org/mongo-driver/v2 v2.5.0 // indirect
golang.org/x/arch v0.22.0 // indirect
golang.org/x/crypto v0.48.0 // indirect
go.sum
+20
View File
@@ -9,8 +9,12 @@ github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gE
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/gabriel-vasile/mimetype v1.4.12 h1:e9hWvmLYvtp846tLHam2o++qitpguFiYCKbn0w9jyqw=
github.com/gabriel-vasile/mimetype v1.4.12/go.mod h1:d+9Oxyo1wTzWdyVUPMmXFvp4F9tea18J8ufA774AB3s=
github.com/gin-contrib/sse v1.1.0 h1:n0w2GMuUpWDVp7qSpvze6fAu9iRxJY4Hmj6AmBOU05w=
@@ -32,6 +36,8 @@ github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7Lk
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
@@ -42,8 +48,14 @@ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnr
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
github.com/mark3labs/mcp-go v0.56.0 h1:7aCj2wODCskMi08f923ADG+EfELZBdiKILny415cIS8=
github.com/mark3labs/mcp-go v0.56.0/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -63,6 +75,12 @@ github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SA
github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ=
github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU=
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -78,6 +96,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
github.com/ugorji/go/codec v1.3.1 h1:waO7eEiFDwidsBN6agj1vJQ4AG7lh2yqXyOXqhgQuyY=
github.com/ugorji/go/codec v1.3.1/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE=
go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0=
go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y=
internal/executor/spark_submit.go
+159
View File
@@ -0,0 +1,159 @@
// Package executor runs the local spark-submit binary for the spark_submit Tool.
//
// It is the single place in the codebase that shells out to user-configured
// binaries. To prevent command injection the binary is taken verbatim from
// cluster.SparkSubmitExecuteBin and arguments are passed via exec.Command's
// slice form — never as a single string, never via "sh -c".
package executor
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"os/exec"
"regexp"
"strings"
"time"
)
// stdoutCap caps spark-submit stdout/stderr to keep the per-Tool log file
// from blowing up on verbose jobs. 10 MiB is generous for the "Submitted
// application application_xxx" line and a few progress lines.
const stdoutCap = 10 << 20
// appIDPattern matches the YARN line printed on successful submission:
//
// "Submitted application application_12345_0001"
var appIDPattern = regexp.MustCompile(`application_\d+_\d+`)
// SparkSubmitOpts is the resolved command line for a single spark-submit run.
//
// Binary must be an absolute path (validated at cluster create/update time).
// Args are the post-binary CLI args, already including any cluster-default
// args prepended by the caller.
type SparkSubmitOpts struct {
Binary string
Args []string
// Timeout is the wall-clock budget for the entire run. Zero means no
// timeout. We use context.WithTimeout to enforce it.
Timeout time.Duration
}
// Result is what spark_submit returns to the LLM.
type Result struct {
AppID string `json:"app_id,omitempty"`
ExitCode int `json:"exit_code"`
StdoutTail string `json:"stdout_tail"`
StderrTail string `json:"stderr_tail"`
DurationMS int64 `json:"duration_ms"`
}
// ErrBinaryEmpty is returned when the cluster has no spark-submit binary
// configured. The agent should call list_clusters to surface this.
var ErrBinaryEmpty = errors.New("executor: spark_submit_execute_bin is empty")
// Run executes the configured spark-submit and returns the parsed result.
//
// Command construction: exec.Command(name, args...) — the slice form means
// no shell, so a malicious arg like "foo; rm -rf /" is passed as a single
// argument to spark-submit, never executed.
func Run(ctx context.Context, opts SparkSubmitOpts) (*Result, error) {
if strings.TrimSpace(opts.Binary) == "" {
return nil, ErrBinaryEmpty
}
if opts.Timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
defer cancel()
}
cmd := exec.CommandContext(ctx, opts.Binary, opts.Args...)
// Critical: do NOT call cmd.Run() with /bin/sh. The slice form above is
// the entire injection defense.
stdout := &capBuf{cap: stdoutCap}
stderr := &capBuf{cap: stdoutCap}
cmd.Stdout = stdout
cmd.Stderr = stderr
start := time.Now()
err := cmd.Run()
dur := time.Since(start)
res := &Result{
ExitCode: exitCodeFromErr(err),
StdoutTail: stdout.String(),
StderrTail: stderr.String(),
DurationMS: dur.Milliseconds(),
}
// app_id: prefer stdout, fall back to stderr. Both are searched only on
// the tail (cap-bounded) so we don't pay for the full output.
if id := matchAppID(res.StdoutTail); id != "" {
res.AppID = id
} else if id := matchAppID(res.StderrTail); id != "" {
res.AppID = id
}
if err != nil {
// Wrap the underlying error but keep the parsed result so the LLM
// can see exit_code and stderr even on failure.
return res, fmt.Errorf("executor: spark-submit failed: %w", err)
}
return res, nil
}
func matchAppID(s string) string {
if s == "" {
return ""
}
return appIDPattern.FindString(s)
}
// exitCodeFromErr returns 0 on nil, the process's exit code on *exec.ExitError,
// and 1 on other errors (e.g. context deadline). exec.ExitError is the typed
// way to ask the OS-level exit code without inspecting strings.
func exitCodeFromErr(err error) int {
if err == nil {
return 0
}
var exitErr *exec.ExitError
if errors.As(err, &exitErr) {
return exitErr.ExitCode()
}
// context.DeadlineExceeded or failed to start — surface as 1.
return 1
}
// capBuf is an io.Writer that keeps at most `cap` bytes from the end.
//
// The point isn't to deliver partial stdout to the LLM (LLM rarely needs the
// full stream) — it's to bound memory and the per-Tool log file. We keep
// the tail because that's where "Submitted application application_xxx" and
// final error messages live.
type capBuf struct {
cap int
buf bytes.Buffer
}
func (b *capBuf) Write(p []byte) (int, error) {
n, err := b.buf.Write(p)
if err != nil {
return n, err
}
if b.buf.Len() > b.cap {
// Drop the front. Cheaper than a ring buffer; works fine at 10 MiB.
overflow := b.buf.Len() - b.cap
_ = b.buf.Next(overflow)
}
return len(p), nil
}
func (b *capBuf) String() string { return b.buf.String() }
// Ensure io.Discard stays referenced; we keep this in case future variants
// want to write to a logger instead of capturing.
var _ = io.Discard
internal/executor/spark_submit_test.go
+104
View File
@@ -0,0 +1,104 @@
package executor
import (
"context"
"runtime"
"strings"
"testing"
"time"
)
// TestRun_BinaryEmpty covers the ErrBinaryEmpty path. The agent-facing
// behavior is "list the cluster, fix the binary field" — we just guard the
// error here.
func TestRun_BinaryEmpty(t *testing.T) {
_, err := Run(context.Background(), SparkSubmitOpts{Binary: "", Args: []string{"--help"}})
if err == nil || !strings.Contains(err.Error(), "empty") {
t.Fatalf("want empty-binary error, got %v", err)
}
}
// TestRun_FakeBinary uses /bin/echo as a stand-in for spark-submit. We
// verify slice-form arg passing (no shell) by including an arg that would
// be catastrophic if interpreted by a shell.
func TestRun_FakeBinary(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses unix binary")
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
res, err := Run(ctx, SparkSubmitOpts{
Binary: "/bin/echo",
Args: []string{"--master", "yarn; rm -rf /tmp/this-should-not-exist"},
Timeout: 3 * time.Second,
})
if err != nil {
t.Fatalf("Run: %v", err)
}
if res.ExitCode != 0 {
t.Fatalf("ExitCode: got %d, want 0", res.ExitCode)
}
// The arg must be passed through verbatim, not interpreted by a shell.
if !strings.Contains(res.StdoutTail, "yarn; rm -rf /tmp/this-should-not-exist") {
t.Errorf("stdout missing literal arg, got: %q", res.StdoutTail)
}
if res.AppID != "" {
t.Errorf("AppID should be empty for echo, got %q", res.AppID)
}
}
// TestRun_NonZeroExit captures the exit code when the binary returns !=0.
func TestRun_NonZeroExit(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses unix binary")
}
res, err := Run(context.Background(), SparkSubmitOpts{
Binary: "/bin/sh",
Args: []string{"-c", "exit 7"},
Timeout: 3 * time.Second,
})
if err == nil {
t.Fatal("want non-nil error from non-zero exit")
}
if res.ExitCode != 7 {
t.Errorf("ExitCode: got %d, want 7", res.ExitCode)
}
}
// TestMatchAppID locks down the regex so a YARN output tweak doesn't
// silently break app_id extraction.
func TestMatchAppID(t *testing.T) {
cases := []struct {
in, want string
}{
{"Submitted application application_12345_0001", "application_12345_0001"},
{"... some prefix application_999_42 ... tail", "application_999_42"},
{"no app id here", ""},
{"", ""},
{"application_1_2 application_3_4", "application_1_2"}, // first match
}
for _, tc := range cases {
got := matchAppID(tc.in)
if got != tc.want {
t.Errorf("matchAppID(%q) = %q, want %q", tc.in, got, tc.want)
}
}
}
// TestCapBuf ensures the tail-capping behavior: writing more than `cap`
// bytes keeps the most recent `cap` bytes.
func TestCapBuf(t *testing.T) {
b := &capBuf{cap: 10}
// Write 20 bytes total in two chunks.
if _, err := b.Write([]byte("0123456789")); err != nil {
t.Fatal(err)
}
if _, err := b.Write([]byte("ABCDEFGHIJ")); err != nil {
t.Fatal(err)
}
// Tail should be the last 10 bytes: "ABCDEFGHIJ".
if got := b.String(); got != "ABCDEFGHIJ" {
t.Errorf("capBuf tail: got %q, want %q", got, "ABCDEFGHIJ")
}
}
main.go
+37 -2
View File
@@ -1,7 +1,7 @@
// Command spark-mcp-go launches the Spark MCP Server.
//
// Phase 4 scope: wire config + slog + Gin + /healthz + /admin/* cluster CRUD
// with audit log. MCP transport and Tool surface land in later phases.
// Phase 5 scope: config + slog + Gin + /healthz + /admin/* + /mcp
// Streamable HTTP. 2 MCP Tools wired (list_clusters, spark_submit).
package main
import (
@@ -17,9 +17,14 @@ import (
"github.com/gin-gonic/gin"
"spark-mcp-go/internal/admin"
"spark-mcp-go/internal/analyzer"
"spark-mcp-go/internal/audit"
"spark-mcp-go/internal/config"
"spark-mcp-go/internal/httpclient"
"spark-mcp-go/internal/logging"
mcpsrv "spark-mcp-go/internal/mcp"
"spark-mcp-go/internal/mcp/tools"
"spark-mcp-go/internal/middleware"
"spark-mcp-go/internal/storage"
)
@@ -70,6 +75,36 @@ func run() error {
admin.Mount(r, db.Clusters(), audit.NewRepo(db), cfg.AdminTokens)
mcpHandler, err := mcpsrv.Handler(&tools.Deps{
Logger: logger.With("component", "mcp"),
ClusterRepo: db.Clusters(),
SparkSubmitTimeout: cfg.SparkSubmitTimeout,
HTTPClient: httpclient.New(httpclient.Config{
Timeout: cfg.HTTPClientTimeout,
MaxResponseBytes: cfg.MaxResponseBytes,
}),
MaxResponseBytes: cfg.MaxResponseBytes,
DataDir: cfg.DataDir,
AnalyzerThresholds: analyzer.Thresholds{
DataSkewRatio: cfg.AnalyzerDataSkewRatio,
GCPressureRatio: cfg.AnalyzerGCPressureRatio,
BottleneckShuffleGB: cfg.AnalyzerBottleneckShuffleGB,
},
})
if err != nil {
return err
}
// Wire AgentAuth for /mcp. We can't use gin.WrapH alone because
// AgentAuth is a gin middleware; compose it manually:
agentAuth := middleware.AgentAuth(cfg.AgentToken)
r.Any("/mcp", func(c *gin.Context) {
agentAuth(c)
if c.IsAborted() {
return
}
mcpHandler.ServeHTTP(c.Writer, c.Request)
})
srv := &http.Server{
Addr: cfg.ListenAddr,
Handler: r,