llmgw/internal/control/adminapi/mcp.go

package adminapi

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/mcpclient"
	"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// MCPStore keeps the company's MCP servers and the secrets their headers
// carry.
type MCPStore interface {
	MCPServers(ctx context.Context) ([]store.MCPServer, error)
	MCPServer(ctx context.Context, name string) (store.MCPServer, error)
	PutMCPServer(ctx context.Context, m store.MCPServer) error
	DeleteMCPServer(ctx context.Context, name string) error
	SetClusterSecret(ctx context.Context, name, value string) error
}

// SetMCP enables the endpoints for the company's MCP servers.
func (a *API) SetMCP(st MCPStore) { a.mcp = st }

// MCPServerView is a server as the API shows it: header values only for the
// headers sent as written, and for the others the reference read.
type MCPServerView struct {
	Name           string            `json:"name"`
	URL            string            `json:"url"`
	Description    string            `json:"description,omitempty"`
	Headers        map[string]string `json:"headers,omitempty"`
	HeaderRefs     map[string]string `json:"header_refs,omitempty"`
	Enabled        bool              `json:"enabled"`
	TimeoutSeconds int               `json:"timeout_seconds,omitempty"`
	RatePerMinute  int               `json:"rate_per_minute,omitempty"`
	CreatedAt      time.Time         `json:"created_at"`
	CreatedBy      string            `json:"created_by,omitempty"`
}

// MCPServerRequest adds a server or changes one. Secrets are header values
// to keep in the store, sealed, each then read by reference; HeaderRefs name
// secrets kept elsewhere (a file, env:, vault:).
type MCPServerRequest struct {
	URL            string            `json:"url"`
	Description    string            `json:"description,omitempty"`
	Headers        map[string]string `json:"headers,omitempty"`
	HeaderRefs     map[string]string `json:"header_refs,omitempty"`
	Secrets        map[string]string `json:"secrets,omitempty"`
	Enabled        *bool             `json:"enabled,omitempty"`
	TimeoutSeconds int               `json:"timeout_seconds,omitempty"`
	// RatePerMinute bounds each caller's calls to its tools; a negative
	// number takes the bound away.
	RatePerMinute int `json:"rate_per_minute,omitempty"`
}

// MCPTestView is what the server said when asked: its tools, or why not.
type MCPTestView struct {
	OK    bool     `json:"ok"`
	Tools []string `json:"tools,omitempty"`
	Error string   `json:"error,omitempty"`
}

func mcpServerView(m store.MCPServer) MCPServerView {
	return MCPServerView{Name: m.Name, URL: m.URL, Description: m.Description, Headers: m.Headers,
		HeaderRefs: m.HeaderRefs, Enabled: m.Enabled, TimeoutSeconds: int(m.Timeout / time.Second), RatePerMinute: m.RatePerMinute,
		CreatedAt: m.CreatedAt, CreatedBy: m.CreatedBy}
}

func (a *API) registerMCP(mux Mux) {
	need := func(w http.ResponseWriter) bool {
		if a.mcp == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store to keep MCP servers in")
			return false
		}
		return true
	}
	mux.HandleFunc("POST /api/v1/mcp", a.MCP)
	mux.HandleFunc("GET /api/v1/mcp/servers", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		list, err := a.mcp.MCPServers(r.Context())
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		out := make([]MCPServerView, len(list))
		for i, m := range list {
			out[i] = mcpServerView(m)
		}
		writeJSON(w, http.StatusOK, out)
	})
	mux.HandleFunc("PUT /api/v1/mcp/servers/{name}", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		name := r.PathValue("name")
		var req MCPServerRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		m, err := a.mcpServerFrom(r.Context(), name, req, actorOf(r))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		if err := a.mcp.PutMCPServer(r.Context(), m); err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, mcpServerView(m))
	})
	mux.HandleFunc("DELETE /api/v1/mcp/servers/{name}", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		switch err := a.mcp.DeleteMCPServer(r.Context(), r.PathValue("name")); {
		case errors.Is(err, store.ErrNotFound):
			writeError(w, http.StatusNotFound, "no MCP server "+r.PathValue("name"))
		case err != nil:
			writeError(w, http.StatusInternalServerError, err.Error())
		default:
			writeJSON(w, http.StatusOK, map[string]string{"removed": r.PathValue("name")})
		}
	})
	mux.HandleFunc("POST /api/v1/mcp/servers/{name}/test", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		m, err := a.mcp.MCPServer(r.Context(), r.PathValue("name"))
		if errors.Is(err, store.ErrNotFound) {
			writeError(w, http.StatusNotFound, "no MCP server "+r.PathValue("name"))
			return
		}
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, a.testMCPServer(r.Context(), m))
	})
}

// mcpServerFrom is the server a request makes, over the one of that name if
// there is one: its secrets kept in the store, and named by reference.
func (a *API) mcpServerFrom(ctx context.Context, name string, req MCPServerRequest, by string) (store.MCPServer, error) {
	if !store.ValidMCPServerName(name) {
		return store.MCPServer{}, errors.New("a server's name is lowercase letters, digits, - and _, starting with a letter: " +
			"its tools are NAME.tool")
	}
	m, err := a.mcp.MCPServer(ctx, name)
	switch {
	case errors.Is(err, store.ErrNotFound):
		m = store.MCPServer{Name: name, Enabled: true, CreatedBy: by}
	case err != nil:
		return m, err
	}
	if req.URL != "" {
		m.URL = req.URL
	}
	if u, err := url.Parse(m.URL); m.URL != AdminMCPAddress && (err != nil || (u.Scheme != "https" && u.Scheme != "http") || u.Host == "") {
		return m, errors.New("the server's url is an http:// or https:// address, or " + AdminMCPAddress + " for this cluster's own")
	}
	if req.Description != "" {
		m.Description = req.Description
	}
	if req.Headers != nil {
		m.Headers = req.Headers
	}
	if req.HeaderRefs != nil {
		m.HeaderRefs = req.HeaderRefs
	}
	if m.HeaderRefs == nil {
		m.HeaderRefs = map[string]string{}
	}
	for header, value := range req.Secrets {
		if strings.TrimSpace(header) == "" || value == "" {
			return m, errors.New("a secret header needs a name and a value")
		}
		key := "mcp/" + name + "/" + header
		if err := a.mcp.SetClusterSecret(ctx, key, value); err != nil {
			return m, err
		}
		m.HeaderRefs[header] = "store:" + key
		delete(m.Headers, header)
	}
	if req.Enabled != nil {
		m.Enabled = *req.Enabled
	}
	if req.TimeoutSeconds > 0 {
		m.Timeout = time.Duration(req.TimeoutSeconds) * time.Second
	}
	switch {
	case req.RatePerMinute > 0:
		m.RatePerMinute = req.RatePerMinute
	case req.RatePerMinute < 0:
		m.RatePerMinute = 0
	}
	return m, nil
}

// testMCPServer asks a server for its tools, with its credentials, as /mcp
// would.
func (a *API) testMCPServer(ctx context.Context, m store.MCPServer) MCPTestView {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	c := &mcpclient.Client{URL: m.URL, Header: func(ctx context.Context) (http.Header, error) {
		h := http.Header{}
		for k, v := range m.Headers {
			h.Set(k, v)
		}
		for k, ref := range m.HeaderRefs {
			v, err := secrets.ReadString(ctx, ref)
			if err != nil {
				return nil, err
			}
			h.Set(k, v)
		}
		return h, nil
	}}
	if m.URL == AdminMCPAddress {
		c.URL, c.HTTP = InProcessURL, &http.Client{Transport: InProcess(http.HandlerFunc(a.MCP))}
	}
	tools, err := c.Tools(ctx)
	if err != nil {
		return MCPTestView{Error: err.Error()}
	}
	v := MCPTestView{OK: true}
	for _, t := range tools {
		v.Tools = append(v.Tools, m.Name+"."+t.Name)
	}
	sort.Strings(v.Tools)
	return v
}