spark-mcp/main.go

// Command spark-mcp-go launches the Spark MCP Server.
//
// Phase 5 scope: config + slog + Gin + /healthz + /admin/* + /mcp
// Streamable HTTP. 2 MCP Tools wired (list_clusters, spark_submit).
package main

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"

	"spark-mcp-go/internal/admin"
	"spark-mcp-go/internal/analyzer"
	"spark-mcp-go/internal/audit"
	"spark-mcp-go/internal/config"
	"spark-mcp-go/internal/httpclient"
	"spark-mcp-go/internal/logging"
	mcpsrv "spark-mcp-go/internal/mcp"
	"spark-mcp-go/internal/mcp/tools"
	"spark-mcp-go/internal/middleware"
	"spark-mcp-go/internal/storage"
	"spark-mcp-go/internal/uploads"
)

const startupSweepTimeout = 5 * time.Second

func main() {
	if err := run(); err != nil {
		// run() handles its own logging; this is the last-chance report.
		slog.Error("fatal", "err", err)
		os.Exit(1)
	}
}

func run() error {
	cfg, err := config.Load()
	if err != nil {
		return err
	}

	logger, shutdownLog, err := logging.New(logging.Config{
		Level:    cfg.LogLevel,
		Dir:      cfg.LogDir,
		Format:   cfg.LogFormat,
		LogToStd: true,
	})
	if err != nil {
		return err
	}
	defer shutdownLog()
	slog.SetDefault(logger)

	if err := os.MkdirAll(cfg.LogDir+"/tools", 0o700); err != nil {
		return err
	}

	db, err := storage.Open(cfg.SQLitePath)
	if err != nil {
		return err
	}
	defer db.Close()
	logger.Info("storage.open", "path", cfg.SQLitePath)

	uploadStore, err := uploads.New(filepath.Join(cfg.DataDir, "uploads"))
	if err != nil {
		return err
	}
	sweepCtx, sweepCancel := context.WithTimeout(context.Background(), startupSweepTimeout)
	n, err := uploadStore.Sweep(sweepCtx, cfg.UploadTTL)
	sweepCancel()
	if errors.Is(err, context.DeadlineExceeded) {
		logger.Warn("uploads.sweep.timeout", "deleted", n, "err", err)
	} else if err != nil {
		logger.Error("uploads.sweep", "err", err)
	} else {
		logger.Info("uploads.sweep", "deleted", n)
	}

	backfillCtx, backfillCancel := context.WithTimeout(context.Background(), startupSweepTimeout)
	n, err = uploadStore.Backfill(backfillCtx, db.Uploads())
	backfillCancel()
	if errors.Is(err, context.DeadlineExceeded) {
		logger.Warn("uploads.backfill.timeout", "inserted", n, "err", err)
	} else if err != nil {
		logger.Error("uploads.backfill", "err", err)
	} else {
		logger.Info("uploads.backfill", "inserted", n)
	}

	uploadStore.SetRepo(db.Uploads())

	gin.SetMode(cfg.GinMode)
	r := gin.New()
	r.Use(gin.Recovery())

	r.GET("/healthz", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"ok": true, "version": "0.0.0"})
	})

	admin.Mount(r, db.Clusters(), db.Uploads(), audit.NewRepo(db), &uploadStore, cfg.AdminTokens)

	mcpHandler, err := mcpsrv.Handler(&tools.Deps{
		Logger:             logger.With("component", "mcp"),
		AuditRepo:          audit.NewRepo(db),
		ClusterRepo:        db.Clusters(),
		SparkSubmitTimeout: cfg.SparkSubmitTimeout,
		HTTPClient: httpclient.New(httpclient.Config{
			Timeout:          cfg.HTTPClientTimeout,
			MaxResponseBytes: cfg.MaxResponseBytes,
		}),
		MaxResponseBytes: cfg.MaxResponseBytes,
		DataDir:          cfg.DataDir,
		UploadStore:      uploadStore,
		AnalyzerThresholds: analyzer.Thresholds{
			DataSkewRatio:       cfg.AnalyzerDataSkewRatio,
			GCPressureRatio:     cfg.AnalyzerGCPressureRatio,
			BottleneckShuffleGB: cfg.AnalyzerBottleneckShuffleGB,
		},
	})
	if err != nil {
		return err
	}
	// Wire AgentAuth for /mcp. We can't use gin.WrapH alone because
	// AgentAuth is a gin middleware; compose it manually:
	agentAuth := middleware.AgentAuth(cfg.AgentToken)
	r.Any("/mcp", func(c *gin.Context) {
		agentAuth(c)
		if c.IsAborted() {
			return
		}
		mcpHandler.ServeHTTP(c.Writer, c.Request)
	})

	srv := &http.Server{
		Addr:              cfg.ListenAddr,
		Handler:           r,
		ReadHeaderTimeout: 10 * time.Second,
	}

	errCh := make(chan error, 1)
	go func() {
		logger.Info("server.start", "addr", cfg.ListenAddr, "data_dir", cfg.DataDir)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	select {
	case err := <-errCh:
		return err
	case sig := <-stop:
		logger.Info("server.shutdown", "signal", sig.String())
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	return srv.Shutdown(ctx)
}