上一个 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>
101 lines
3.1 KiB
Go
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
|
|
}
|