Phase 5 Batch 1: executor + list_clusters + spark_submit + /mcp 路由
- 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:
@@ -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
|
||||
Reference in New Issue
Block a user