diff --git a/go.mod b/go.mod index e62ef92..1698bff 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index f861ada..586df15 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/internal/executor/spark_submit.go b/internal/executor/spark_submit.go new file mode 100644 index 0000000..b4fd4eb --- /dev/null +++ b/internal/executor/spark_submit.go @@ -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 diff --git a/internal/executor/spark_submit_test.go b/internal/executor/spark_submit_test.go new file mode 100644 index 0000000..79a718d --- /dev/null +++ b/internal/executor/spark_submit_test.go @@ -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") + } +} diff --git a/main.go b/main.go index e8b21ba..9236e7c 100644 --- a/main.go +++ b/main.go @@ -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,