registry_mcp/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strings"

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

	"fn-registry/registry"
)

const version = "0.1.0"

type config struct {
	httpAddr     string
	bind         string
	enableRun    bool
	enableWrite  bool
	registryRoot string
	logLevel     string
}

func main() {
	var cfg config
	flag.StringVar(&cfg.httpAddr, "http", "", "Listen on HTTP+SSE address (e.g. :7733). Empty = stdio.")
	flag.StringVar(&cfg.bind, "bind", "127.0.0.1", "HTTP bind address. Use 0.0.0.0 only with REGISTRY_API_TOKEN set.")
	flag.BoolVar(&cfg.enableRun, "enable-run", false, "Enable fn_run tool (executes registry functions/pipelines).")
	flag.BoolVar(&cfg.enableWrite, "enable-write", false, "Enable fn_create_function tool (writes files + runs fn index).")
	flag.StringVar(&cfg.registryRoot, "registry-root", "", "Override FN_REGISTRY_ROOT.")
	flag.StringVar(&cfg.logLevel, "log-level", "info", "Log level: debug, info, warn, error.")
	flag.Parse()

	// Slog → stderr (stdio JSON-RPC owns stdout).
	lvl := parseLevel(cfg.logLevel)
	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: lvl})))

	root, err := resolveRoot(cfg.registryRoot)
	if err != nil {
		slog.Error("resolve registry root", "err", err)
		os.Exit(1)
	}
	dbPath := filepath.Join(root, "registry.db")
	if _, err := os.Stat(dbPath); err != nil {
		slog.Error("registry.db not found", "path", dbPath, "hint", "set FN_REGISTRY_ROOT")
		os.Exit(1)
	}

	db, err := registry.Open(dbPath)
	if err != nil {
		slog.Error("open registry.db", "err", err)
		os.Exit(1)
	}
	defer db.Close()

	srv := server.NewMCPServer(
		"registry_mcp",
		version,
		server.WithToolCapabilities(true),
	)

	registerTools(srv, db, root, cfg)

	slog.Info("starting registry_mcp",
		"version", version,
		"root", root,
		"transport", transportLabel(cfg),
		"enable_run", cfg.enableRun,
		"enable_write", cfg.enableWrite,
	)

	if cfg.httpAddr == "" {
		if err := server.ServeStdio(srv); err != nil {
			slog.Error("stdio server", "err", err)
			os.Exit(1)
		}
		return
	}

	if err := serveHTTP(srv, cfg); err != nil {
		slog.Error("http server", "err", err)
		os.Exit(1)
	}
}

func transportLabel(cfg config) string {
	if cfg.httpAddr == "" {
		return "stdio"
	}
	return fmt.Sprintf("http %s%s", cfg.bind, cfg.httpAddr)
}

func parseLevel(s string) slog.Level {
	switch strings.ToLower(s) {
	case "debug":
		return slog.LevelDebug
	case "warn":
		return slog.LevelWarn
	case "error":
		return slog.LevelError
	default:
		return slog.LevelInfo
	}
}

func resolveRoot(override string) (string, error) {
	if override != "" {
		return filepath.Abs(override)
	}
	if env := os.Getenv("FN_REGISTRY_ROOT"); env != "" {
		return filepath.Abs(env)
	}
	cwd, err := os.Getwd()
	if err != nil {
		return "", err
	}
	dir := cwd
	for {
		if _, err := os.Stat(filepath.Join(dir, "registry.db")); err == nil {
			return dir, nil
		}
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			// Found a go.mod but no registry.db here: keep going up.
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			break
		}
		dir = parent
	}
	return "", fmt.Errorf("registry.db not found upward from %s", cwd)
}

func registerTools(s *server.MCPServer, db *registry.DB, root string, cfg config) {
	deps := &deps{db: db, root: root, cfg: cfg}

	// Read-only tools — always on.
	s.AddTool(searchTool(), mcp.NewTypedToolHandler(deps.handleSearch))
	s.AddTool(showTool(), mcp.NewTypedToolHandler(deps.handleShow))
	s.AddTool(codeTool(), mcp.NewTypedToolHandler(deps.handleCode))
	s.AddTool(listDomainsTool(), mcp.NewTypedToolHandler(deps.handleListDomains))
	s.AddTool(usesTool(), mcp.NewTypedToolHandler(deps.handleUses))
	s.AddTool(doctorTool(), mcp.NewTypedToolHandler(deps.handleDoctor))
	s.AddTool(proposalTool(), mcp.NewTypedToolHandler(deps.handleProposal))

	// Mutating tools — opt-in.
	if cfg.enableRun {
		s.AddTool(runTool(), mcp.NewTypedToolHandler(deps.handleRun))
	}
	if cfg.enableWrite {
		s.AddTool(createFunctionTool(), mcp.NewTypedToolHandler(deps.handleCreateFunction))
	}
}

// deps carries state into tool handlers.
type deps struct {
	db   *registry.DB
	root string
	cfg  config
}

// serveHTTP hosts the MCP server over Streamable HTTP with optional bearer auth.
func serveHTTP(s *server.MCPServer, cfg config) error {
	addr := cfg.bind + cfg.httpAddr

	httpSrv := server.NewStreamableHTTPServer(s)

	token := os.Getenv("REGISTRY_API_TOKEN")
	if cfg.bind == "0.0.0.0" && token == "" {
		return fmt.Errorf("--bind 0.0.0.0 requires REGISTRY_API_TOKEN")
	}

	mux := http.NewServeMux()
	if token != "" {
		mux.Handle("/", authMiddleware(token, httpSrv))
	} else {
		mux.Handle("/", httpSrv)
	}

	slog.Info("listening http", "addr", addr)
	return http.ListenAndServe(addr, mux)
}

func authMiddleware(token string, next http.Handler) http.Handler {
	expected := "Bearer " + token
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != expected {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// silence unused-import for context if no tool uses it directly here
var _ = context.Background