spark-mcp/internal/mcp/tools/fetch_spark_metrics.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"

	"github.com/mark3labs/mcp-go/mcp"

	"spark-mcp-go/internal/analyzer"
	"spark-mcp-go/internal/cluster"
	"spark-mcp-go/internal/httpclient"
)

const FetchSparkMetricsName = "fetch_spark_metrics"

type shsExecutor struct {
	ID        string `json:"id"`
	GCTimeMS  int64  `json:"gc_time_ms"`
	CPUTimeMS int64  `json:"cpu_time_ms"`
}

type shsStage struct {
	StageID           int   `json:"stageId"`
	DurationMS        int64 `json:"duration_ms"`
	ShuffleReadBytes  int64 `json:"shuffle_read_bytes"`
	ShuffleWriteBytes int64 `json:"shuffle_write_bytes"`
	MaxPartitionBytes int64 `json:"max_partition_bytes"`
	MinPartitionBytes int64 `json:"min_partition_bytes"`
}

// NewFetchSparkMetricsTool returns the schema for the fetch_spark_metrics MCP Tool.
func NewFetchSparkMetricsTool() mcp.Tool {
	return mcp.NewTool(FetchSparkMetricsName,
		mcp.WithDescription("Fetch Spark History Server metrics for an application. Returns SHS executor/stage data plus optional heuristic findings."),
		mcp.WithString("cluster_id",
			mcp.Required(),
			mcp.Description("ID of the configured cluster whose SHS is used"),
		),
		mcp.WithString("app_id",
			mcp.Required(),
			mcp.Description("Spark application ID, e.g. application_1234567890_0001"),
		),
		mcp.WithString("format",
			mcp.Description("Output format: raw SHS JSON or summary with heuristic findings"),
			mcp.Enum("summary", "raw"),
			mcp.DefaultString("summary"),
		),
	)
}

// FetchSparkMetricsHandler fetches SHS executor and stage data and optionally runs analysis.
func (d *Deps) FetchSparkMetricsHandler(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	clusterID, err := req.RequireString("cluster_id")
	if err != nil {
		return errResult("fetch_spark_metrics: " + err.Error()), nil
	}
	appID, err := req.RequireString("app_id")
	if err != nil {
		return errResult("fetch_spark_metrics: " + err.Error()), nil
	}

	format := "summary"
	if v, ok := req.GetArguments()["format"].(string); ok && v != "" {
		format = v
	}

	callLog := startToolCall(ctx, d.Logger, FetchSparkMetricsName, map[string]any{
		"cluster_id": clusterID,
		"app_id":     appID,
		"format":     format,
	})
	defer callLog.End()

	cl, err := d.ClusterRepo.Get(ctx, clusterID)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_spark_metrics: cluster " + clusterID + ": " + err.Error()), nil
	}

	base := strings.TrimRight(cl.SHSURL, "/")
	execURL := base + "/api/v1/applications/" + appID + "/executors"
	stageURL := base + "/api/v1/applications/" + appID + "/stages"

	execBody, err := d.fetchInternal(ctx, cl, execURL)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_spark_metrics: executors: " + err.Error()), nil
	}
	stageBody, err := d.fetchInternal(ctx, cl, stageURL)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_spark_metrics: stages: " + err.Error()), nil
	}

	var rawExec, rawStage any
	_ = json.Unmarshal(execBody, &rawExec)
	_ = json.Unmarshal(stageBody, &rawStage)

	if format != "summary" {
		return textResult(encodeJSON(map[string]any{
			"executors": rawExec,
			"stages":    rawStage,
		})), nil
	}

	input, err := d.sparkMetricsInput(ctx, cl, appID)
	if err != nil {
		callLog.WithError(err)
		return errResult("fetch_spark_metrics: analyze: " + err.Error()), nil
	}
	findings := analyzer.Analyze(input, d.AnalyzerThresholds)

	return textResult(encodeJSON(map[string]any{
		"executors": rawExec,
		"stages":    rawStage,
		"findings":  findings,
	})), nil
}

// fetchInternal performs an SSRF-aware, authenticated GET to a cluster URL.
func (d *Deps) fetchInternal(ctx context.Context, cl *cluster.Cluster, rawURL string) ([]byte, error) {
	allowedHosts := buildAllowedHosts(cl)

	targetURL, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("parse url: %w", err)
	}
	targetHost := strings.ToLower(targetURL.Hostname())
	if !hostAllowed(targetHost, allowedHosts) {
		return nil, fmt.Errorf("host %q is not in the cluster allowlist", targetHost)
	}

	if err := httpclient.CheckURL(rawURL, allowedHosts...); err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return nil, fmt.Errorf("build request: %w", err)
	}
	if err := httpclient.ApplyAuth(httpReq, "", cl); err != nil {
		return nil, fmt.Errorf("auth: %w", err)
	}

	resp, err := d.HTTPClient.DoWithRedirect(ctx, httpReq, 5, allowedHosts...)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		return nil, fmt.Errorf("%s returned %d", rawURL, resp.StatusCode)
	}
	return io.ReadAll(resp.Body)
}

// sparkMetricsInput converts SHS executor/stage JSON into analyzer input.
func (d *Deps) sparkMetricsInput(ctx context.Context, cl *cluster.Cluster, appID string) (analyzer.Input, error) {
	base := strings.TrimRight(cl.SHSURL, "/")
	execURL := base + "/api/v1/applications/" + appID + "/executors"
	stageURL := base + "/api/v1/applications/" + appID + "/stages"

	execBody, err := d.fetchInternal(ctx, cl, execURL)
	if err != nil {
		return analyzer.Input{}, err
	}
	stageBody, err := d.fetchInternal(ctx, cl, stageURL)
	if err != nil {
		return analyzer.Input{}, err
	}

	var execs []shsExecutor
	if err := json.Unmarshal(execBody, &execs); err != nil {
		return analyzer.Input{}, fmt.Errorf("parse executors: %w", err)
	}
	var stages []shsStage
	if err := json.Unmarshal(stageBody, &stages); err != nil {
		return analyzer.Input{}, fmt.Errorf("parse stages: %w", err)
	}

	input := analyzer.Input{
		StageMetrics:    make(map[string]analyzer.StageMetric),
		ExecutorMetrics: make(map[string]analyzer.ExecutorMetric),
	}
	for _, e := range execs {
		input.ExecutorMetrics[e.ID] = analyzer.ExecutorMetric{
			GCTimeMS:  e.GCTimeMS,
			CPUTimeMS: e.CPUTimeMS,
		}
	}
	for _, s := range stages {
		key := fmt.Sprintf("stage %d", s.StageID)
		input.StageMetrics[key] = analyzer.StageMetric{
			DurationMS:        s.DurationMS,
			ShuffleReadBytes:  s.ShuffleReadBytes,
			ShuffleWriteBytes: s.ShuffleWriteBytes,
			MaxPartitionBytes: s.MaxPartitionBytes,
			MinPartitionBytes: s.MinPartitionBytes,
		}
	}
	return input, nil
}