Files
spark-mcp/internal/mcp/tools/spark_submit.go
T
tao.chenandClaude be4ca460c9 Phase 5 Batch 1 (补): mcp/server + 2 Tool (Batch 1 commit 时漏 stage)
上一个 commit (7056657) 漏 add 的 mcp 包文件:
- internal/mcp/server.go (Streamable HTTP Handler)
- internal/mcp/tools/list_clusters.go (发现入口)
- internal/mcp/tools/spark_submit.go (本地 exec 提交)
- internal/mcp/tools/helpers.go (textResult/errResult/encodeJSON)

deps.go / rm.go 已在 Phase 5 Batch 2 一起 commit, 此处不重复。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-10 17:24:10 +08:00

101 lines
3.1 KiB
Go

package tools
import (
"context"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"spark-mcp-go/internal/executor"
)
const SparkSubmitName = "spark_submit"
// NewSparkSubmitTool returns the schema for the spark_submit MCP Tool.
func NewSparkSubmitTool() mcp.Tool {
return mcp.NewTool(SparkSubmitName,
mcp.WithDescription("Submit a Spark application to a configured cluster using the cluster's local spark-submit binary. Returns {app_id, exit_code, stdout_tail, stderr_tail, duration_ms}. The cluster's default_submit_args are prepended to your args. Binary is invoked as a child process — no shell, no command injection."),
mcp.WithString("cluster_id",
mcp.Required(),
mcp.Description("ID of the configured cluster (from list_clusters)"),
),
// mcp-go v0.56.0 uses PropertyOption for array item schema.
mcp.WithArray("args",
mcp.Description("Spark-submit CLI args (after default_submit_args). Each element becomes one argv entry — no shell interpretation."),
mcp.Items(map[string]any{"type": "string"}),
),
)
}
// SparkSubmitHandler runs spark-submit against the requested cluster.
func (d *Deps) SparkSubmitHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
clusterID, err := req.RequireString("cluster_id")
if err != nil {
return errResult("spark_submit: " + err.Error()), nil
}
args := req.GetArguments()
var userArgs []string
if rawArgs, ok := args["args"]; ok && rawArgs != nil {
arr, ok := rawArgs.([]any)
if !ok {
return errResult("spark_submit: args is not an array"), nil
}
for i, item := range arr {
s, ok := item.(string)
if !ok {
return errResult(fmt.Sprintf("spark_submit: args[%d] is not a string", i)), nil
}
userArgs = append(userArgs, s)
}
}
callLog := startToolCall(ctx, d.Logger, SparkSubmitName, map[string]any{
"cluster_id": clusterID,
"args": userArgs,
})
defer callLog.End()
cl, err := d.ClusterRepo.Get(ctx, clusterID)
if err != nil {
callLog.WithError(err)
return errResult("spark_submit: cluster " + clusterID + ": " + err.Error()), nil
}
fullArgs := append([]string{}, cl.DefaultSubmitArgs...)
fullArgs = append(fullArgs, userArgs...)
result, err := executor.Run(ctx, executor.SparkSubmitOpts{
Binary: cl.SparkSubmitExecuteBin,
Args: fullArgs,
Timeout: d.SparkSubmitTimeout,
})
if err != nil {
callLog.WithError(err)
// Even on error, result carries ExitCode/Stderr for the LLM.
if result != nil {
callLog.WithResult(map[string]any{
"app_id": result.AppID,
"exit_code": result.ExitCode,
"duration_ms": result.DurationMS,
})
return textResult(encodeJSON(map[string]any{
"app_id": result.AppID,
"exit_code": result.ExitCode,
"stdout_tail": result.StdoutTail,
"stderr_tail": result.StderrTail,
"duration_ms": result.DurationMS,
"error": err.Error(),
})), nil
}
return errResult("spark_submit: " + err.Error()), nil
}
callLog.WithResult(map[string]any{
"app_id": result.AppID,
"exit_code": result.ExitCode,
"duration_ms": result.DurationMS,
})
return textResult(encodeJSON(result)), nil
}