mcp/internal/agent/agent.go

package agent

import (
	"context"
	"crypto/tls"
	"database/sql"
	"fmt"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"syscall"

	mcpv1 "git.wntrmute.dev/kyle/mcp/gen/mcp/v1"
	"git.wntrmute.dev/kyle/mcp/internal/auth"
	"git.wntrmute.dev/kyle/mcp/internal/config"
	"git.wntrmute.dev/kyle/mcp/internal/monitor"
	"git.wntrmute.dev/kyle/mcp/internal/registry"
	"git.wntrmute.dev/kyle/mcp/internal/runtime"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

// Agent is the MCP node agent. It manages containers, stores the registry,
// monitors for drift, and serves the gRPC API.
type Agent struct {
	mcpv1.UnimplementedMcpAgentServiceServer

	Config    *config.AgentConfig
	DB        *sql.DB
	Runtime   runtime.Runtime
	Monitor   *monitor.Monitor
	Logger    *slog.Logger
	PortAlloc *PortAllocator
	Proxy     *ProxyRouter
}

// Run starts the agent: opens the database, sets up the gRPC server with
// TLS and auth, and blocks until SIGINT/SIGTERM.
func Run(cfg *config.AgentConfig) error {
	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level: parseLogLevel(cfg.Log.Level),
	}))

	db, err := registry.Open(cfg.Database.Path)
	if err != nil {
		return fmt.Errorf("open registry: %w", err)
	}
	defer func() { _ = db.Close() }()

	rt := &runtime.Podman{}

	mon := monitor.New(db, rt, cfg.Monitor, cfg.Agent.NodeName, logger)

	proxy, err := NewProxyRouter(cfg.MCProxy.Socket, cfg.MCProxy.CertDir, logger)
	if err != nil {
		return fmt.Errorf("connect to mc-proxy: %w", err)
	}

	a := &Agent{
		Config:    cfg,
		DB:        db,
		Runtime:   rt,
		Monitor:   mon,
		Logger:    logger,
		PortAlloc: NewPortAllocator(),
		Proxy:     proxy,
	}

	tlsCert, err := tls.LoadX509KeyPair(cfg.Server.TLSCert, cfg.Server.TLSKey)
	if err != nil {
		return fmt.Errorf("load TLS cert: %w", err)
	}
	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		MinVersion:   tls.VersionTLS13,
	}

	validator, err := auth.NewMCIASValidator(cfg.MCIAS.ServerURL, cfg.MCIAS.CACert)
	if err != nil {
		return fmt.Errorf("create MCIAS validator: %w", err)
	}

	server := grpc.NewServer(
		grpc.Creds(credentials.NewTLS(tlsConfig)),
		grpc.ChainUnaryInterceptor(
			auth.AuthInterceptor(validator),
		),
	)
	mcpv1.RegisterMcpAgentServiceServer(server, a)

	lis, err := net.Listen("tcp", cfg.Server.GRPCAddr)
	if err != nil {
		return fmt.Errorf("listen %q: %w", cfg.Server.GRPCAddr, err)
	}

	logger.Info("agent starting",
		"addr", cfg.Server.GRPCAddr,
		"node", cfg.Agent.NodeName,
		"runtime", cfg.Agent.ContainerRuntime,
	)

	mon.Start()

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	errCh := make(chan error, 1)
	go func() {
		errCh <- server.Serve(lis)
	}()

	select {
	case <-ctx.Done():
		logger.Info("shutting down")
		mon.Stop()
		server.GracefulStop()
		_ = proxy.Close()
		return nil
	case err := <-errCh:
		mon.Stop()
		return fmt.Errorf("serve: %w", err)
	}
}

func parseLogLevel(level string) slog.Level {
	switch level {
	case "debug":
		return slog.LevelDebug
	case "warn":
		return slog.LevelWarn
	case "error":
		return slog.LevelError
	default:
		return slog.LevelInfo
	}
}