diff --git a/internal/mcp/tools/helpers.go b/internal/mcp/tools/helpers.go new file mode 100644 index 0000000..8b92907 --- /dev/null +++ b/internal/mcp/tools/helpers.go @@ -0,0 +1,80 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + + "github.com/mark3labs/mcp-go/mcp" + "spark-mcp-go/internal/logging" +) + +// textResult wraps a string into a text CallToolResult. +func textResult(s string) *mcp.CallToolResult { + return mcp.NewToolResultText(s) +} + +// errResult wraps a string into an error CallToolResult. +// MCP protocol errors are reserved for exceptional conditions; business +// errors are reported inside the tool result with IsError set by the library. +func errResult(s string) *mcp.CallToolResult { + return mcp.NewToolResultError(s) +} + +// encodeJSON marshals v to JSON. On failure it returns a string literal that +// embeds the error so the caller never receives a nil or empty result. +func encodeJSON(v any) string { + b, err := json.Marshal(v) + if err != nil { + return fmt.Sprintf("", err) + } + return string(b) +} + +// startToolCall starts per-tool-call logging. If logging is not initialized +// (e.g. during tests) it falls back to a no-op logger so handlers never panic. +func startToolCall(ctx context.Context, logger *slog.Logger, toolName string, params any) logging.ToolCallLogger { + callLog, err := logging.StartToolCall(ctx, logger, toolName, params) + if err != nil { + return &noopToolCallLogger{} + } + return callLog +} + +type noopToolCallLogger struct{} + +func (n *noopToolCallLogger) WithResult(any) {} +func (n *noopToolCallLogger) WithError(error) {} +func (n *noopToolCallLogger) End() {} + +// truncateMiddle limits s to roughly max bytes by keeping the first and +// last halves, separated by a marker. It preserves UTF-8 rune boundaries +// so the result never starts or ends with a broken multi-byte rune. +func truncateMiddle(s string, max int) string { + if len(s) <= max { + return s + } + half := max / 2 + start := s[:half] + end := s[len(s)-half:] + for len(start) > 0 && start[len(start)-1] >= 0x80 && start[len(start)-1] < 0xC0 { + start = start[:len(start)-1] + } + for len(end) > 0 && end[0]&0xC0 == 0x80 { + end = end[1:] + } + return start + "\n... [truncated middle] ...\n" + end +} + +// truncateTail keeps the last max bytes of s, preserving UTF-8 rune boundaries. +func truncateTail(s string, max int) string { + if len(s) <= max { + return s + } + start := len(s) - max + for start < len(s) && s[start]&0xC0 == 0x80 { + start++ + } + return s[start:] +} diff --git a/internal/mcp/tools/list_clusters.go b/internal/mcp/tools/list_clusters.go new file mode 100644 index 0000000..e4aa200 --- /dev/null +++ b/internal/mcp/tools/list_clusters.go @@ -0,0 +1,39 @@ +package tools + +import ( + "context" + + "github.com/mark3labs/mcp-go/mcp" +) + +const ListClustersName = "list_clusters" + +// NewListClustersTool returns the schema for the list_clusters MCP Tool. +func NewListClustersTool() mcp.Tool { + return mcp.NewTool(ListClustersName, + mcp.WithDescription("List active Spark/YARN clusters configured for the agent. Returns the full Cluster struct (minus auth password) for every active cluster — this is the discovery root: from here the agent knows RM/SHS endpoints, auth flavor, and URL allowlists."), + ) +} + +// ListClustersHandler lists configured clusters and returns only active ones. +func (d *Deps) ListClustersHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + callLog := startToolCall(ctx, d.Logger, ListClustersName, nil) + defer callLog.End() + + all, err := d.ClusterRepo.List(ctx) + if err != nil { + callLog.WithError(err) + return errResult("list_clusters: " + err.Error()), nil + } + + // task_plan.md specifies active clusters only. + active := all[:0] + for _, c := range all { + if c.IsActive { + active = append(active, c) + } + } + + callLog.WithResult(map[string]any{"count": len(active)}) + return textResult(encodeJSON(active)), nil +} diff --git a/internal/mcp/tools/spark_submit.go b/internal/mcp/tools/spark_submit.go new file mode 100644 index 0000000..6666b9b --- /dev/null +++ b/internal/mcp/tools/spark_submit.go @@ -0,0 +1,100 @@ +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 +}