Maven/internal/mcp/manager.go

package mcp

import (
	"context"
	"errors"
	"fmt"
	"log"
	"sort"
	"strings"
	"sync"
	"time"
)

// Defaults for a server block. Small numbers on purpose — see MaxTools.
const (
	// DefaultTimeout bounds one JSON-RPC call. A tool that takes longer than
	// this is not usable in a spoken turn anyway.
	DefaultTimeout = 15 * time.Second
	// DefaultMaxTools caps how many tools ONE server may contribute. The
	// resident model is a 1.7B with a 4096-token context: a catalogue of forty
	// tool names does not fit in its head, and a name it half-remembers is a
	// wrong act. Twelve per server is already generous.
	DefaultMaxTools = 12
	// DefaultReconnectEvery is how long the manager waits before re-dialing a
	// server whose connection died.
	DefaultReconnectEvery = 30 * time.Second
)

// ErrNoServer — the named server is not configured or not connected.
var ErrNoServer = errors.New("mcp: no such server")

// ServerConfig is one configured MCP server. Off unless present.
//
// Exactly one of Command (a subprocess on this box) or URL (a remote or
// loopback HTTP endpoint) must be set.
type ServerConfig struct {
	// Name is the local handle. It prefixes every tool this server
	// contributes, so it must be short and a valid identifier-ish word.
	Name string `json:"name"`
	// Command + Args + Env + Dir describe a stdio server: a child process of
	// mavend, on this box, under this user. argv, never a shell string.
	Command string   `json:"command,omitempty"`
	Args    []string `json:"args,omitempty"`
	Env     []string `json:"env,omitempty"`
	Dir     string   `json:"dir,omitempty"`
	// URL is a streamable-HTTP endpoint. It goes through internal/webfetch, so
	// it inherits the SSRF guard, the size cap and the per-host rate limit.
	URL string `json:"url,omitempty"`
	// AllowPrivate lets THIS server be a loopback or LAN address
	// (http://localhost:9100/mcp is the Vikunja server on homesrv). It is a
	// per-server hole in the private-address guard and it is not the same trust
	// level as a public endpoint: whatever is behind it is inside the network,
	// so an argument the router got wrong reaches something that matters. Set
	// it only for a server you run.
	AllowPrivate bool `json:"allow_private,omitempty"`
	// AllowTools, when non-empty, is the ONLY set of remote tool names taken
	// from this server. This is the knob for keeping the catalogue small and
	// deliberate rather than "whatever the server grew this week".
	AllowTools []string `json:"allow_tools,omitempty"`
	// MaxTools caps the contribution (0 ⇒ DefaultMaxTools).
	MaxTools int `json:"max_tools,omitempty"`
	// Timeout bounds one call (0 ⇒ DefaultTimeout).
	Timeout time.Duration `json:"-"`
	// Enabled=false keeps a configured server described but dark.
	Enabled bool `json:"enabled"`
}

// PosterFactory builds the HTTP door for one server. It is a factory rather
// than a single shared Poster because allow_private is per server: the fetcher
// that may reach http://localhost:9100/mcp must NOT be the same fetcher another
// server's public URL goes through, or one loopback exemption would quietly
// unlock the LAN for all of them.
type PosterFactory func(cfg ServerConfig) (Poster, error)

// Manager owns the connections. Nothing here starts unless at least one server
// is configured and enabled.
type Manager struct {
	newPoster PosterFactory
	mu        sync.Mutex
	conns     map[string]*conn
	order     []string
}

type conn struct {
	cfg      ServerConfig
	client   *Client
	tools    []Tool
	lastErr  error
	lastTry  time.Time
	dialedAt time.Time
}

// NewManager builds a manager for the enabled servers in cfgs. newPoster is
// the guarded HTTP door factory for url servers; pass nil only when no url
// server is configured (a nil factory with a url server is reported per server
// at dial time rather than fatally, so one bad block never stops the daemon).
//
// Dialing is lazy: NewManager validates and records, Connect dials.
func NewManager(newPoster PosterFactory, cfgs []ServerConfig) (*Manager, error) {
	m := &Manager{newPoster: newPoster, conns: map[string]*conn{}}
	for _, c := range cfgs {
		if !c.Enabled {
			continue
		}
		if err := validate(c); err != nil {
			return nil, err
		}
		if _, dup := m.conns[c.Name]; dup {
			return nil, fmt.Errorf("mcp: duplicate server name %q", c.Name)
		}
		if c.Timeout <= 0 {
			c.Timeout = DefaultTimeout
		}
		if c.MaxTools <= 0 {
			c.MaxTools = DefaultMaxTools
		}
		m.conns[c.Name] = &conn{cfg: c}
		m.order = append(m.order, c.Name)
	}
	sort.Strings(m.order)
	return m, nil
}

// Validate checks a set of server blocks without dialling anything, so a typo
// fails at startup rather than at the first turn that needed the tool.
func Validate(cfgs []ServerConfig) error {
	seen := map[string]bool{}
	for _, c := range cfgs {
		if err := validate(c); err != nil {
			return err
		}
		if seen[c.Name] {
			return fmt.Errorf("mcp: duplicate server name %q", c.Name)
		}
		seen[c.Name] = true
	}
	return nil
}

func validate(c ServerConfig) error {
	if strings.TrimSpace(c.Name) == "" {
		return errors.New("mcp: server needs a name")
	}
	if strings.ContainsAny(c.Name, " \t/:") {
		return fmt.Errorf("mcp: server name %q must be one word without spaces, slashes or colons", c.Name)
	}
	hasCmd, hasURL := c.Command != "", c.URL != ""
	if hasCmd == hasURL {
		return fmt.Errorf("mcp: server %q needs exactly one of command or url", c.Name)
	}
	if hasURL && !strings.HasPrefix(c.URL, "http://") && !strings.HasPrefix(c.URL, "https://") {
		return fmt.Errorf("mcp: server %q url must be http or https", c.Name)
	}
	return nil
}

// Servers — the configured, enabled server names, sorted.
func (m *Manager) Servers() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]string(nil), m.order...)
}

// Empty reports whether nothing is configured. The daemon uses it to skip
// wiring entirely.
func (m *Manager) Empty() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.conns) == 0
}

// Connect dials every configured server, handshakes, and discovers tools.
// A server that fails is recorded and retried later by Refresh — one bad
// server never blocks the others, and never blocks boot.
func (m *Manager) Connect(ctx context.Context) {
	for _, name := range m.Servers() {
		if err := m.dial(ctx, name); err != nil {
			log.Printf("mcp: %s: %v", name, err)
		}
	}
}

func (m *Manager) dial(ctx context.Context, name string) error {
	m.mu.Lock()
	c, ok := m.conns[name]
	if !ok {
		m.mu.Unlock()
		return ErrNoServer
	}
	cfg := c.cfg
	c.lastTry = time.Now()
	m.mu.Unlock()

	var tr transport
	var err error
	if cfg.Command != "" {
		tr, err = newStdioTransport(ctx, append([]string{cfg.Command}, cfg.Args...), cfg.Env, cfg.Dir)
	} else if m.newPoster == nil {
		err = fmt.Errorf("server %q has a url but no http door was wired", name)
	} else {
		var poster Poster
		if poster, err = m.newPoster(cfg); err == nil {
			tr = newHTTPTransport(poster, cfg.URL)
		}
	}
	if err != nil {
		m.fail(name, err)
		return err
	}

	cl := newClient(name, tr)
	ictx, cancel := context.WithTimeout(ctx, cfg.Timeout)
	defer cancel()
	if err := cl.Initialize(ictx); err != nil {
		_ = cl.Close()
		m.fail(name, err)
		return err
	}
	tools, err := cl.ListTools(ictx)
	if err != nil {
		// A server with no tools capability is still a usable resource server.
		log.Printf("mcp: %s: list tools: %v", name, err)
		tools = nil
	}
	tools = filterTools(cfg, tools)

	m.mu.Lock()
	if old := m.conns[name].client; old != nil {
		_ = old.Close()
	}
	m.conns[name].client = cl
	m.conns[name].tools = tools
	m.conns[name].lastErr = nil
	m.conns[name].dialedAt = time.Now()
	m.mu.Unlock()
	log.Printf("mcp: %s connected (%s %s), %d tool(s)", name, cl.Info().Name, cl.Info().Version, len(tools))
	return nil
}

func (m *Manager) fail(name string, err error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if c := m.conns[name]; c != nil {
		c.lastErr = err
		c.client = nil
		c.tools = nil
	}
}

// filterTools applies AllowTools and MaxTools, and drops nameless entries.
// Sorted first, so the cap is deterministic rather than "whatever order the
// server felt like".
func filterTools(cfg ServerConfig, in []Tool) []Tool {
	sort.Slice(in, func(i, j int) bool { return in[i].Name < in[j].Name })
	out := make([]Tool, 0, len(in))
	for _, t := range in {
		if len(cfg.AllowTools) > 0 && !contains(cfg.AllowTools, t.Name) {
			continue
		}
		out = append(out, t)
	}
	if cfg.MaxTools > 0 && len(out) > cfg.MaxTools {
		log.Printf("mcp: %s offers %d tools, taking the first %d (raise max_tools or set allow_tools)",
			cfg.Name, len(out), cfg.MaxTools)
		out = out[:cfg.MaxTools]
	}
	return out
}

func contains(hay []string, needle string) bool {
	for _, h := range hay {
		if h == needle {
			return true
		}
	}
	return false
}

// Refresh re-dials any server that is down, if enough time has passed since the
// last attempt. Call it from the daemon's periodic tick — it is cheap when
// everything is up.
func (m *Manager) Refresh(ctx context.Context) {
	now := time.Now()
	var stale []string
	m.mu.Lock()
	for _, name := range m.order {
		c := m.conns[name]
		down := c.client == nil || !c.client.alive()
		if down && now.Sub(c.lastTry) >= DefaultReconnectEvery {
			stale = append(stale, name)
		}
	}
	m.mu.Unlock()
	for _, name := range stale {
		if err := m.dial(ctx, name); err != nil {
			log.Printf("mcp: %s: reconnect: %v", name, err)
		}
	}
}

// Tools — every discovered tool across connected servers, sorted by
// server then name.
func (m *Manager) Tools() []Tool {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []Tool
	for _, name := range m.order {
		out = append(out, m.conns[name].tools...)
	}
	return out
}

// Status is one server's health, for the web surface.
type Status struct {
	Name      string
	Transport string // "stdio" or "http"
	Target    string // command or url
	Connected bool
	Server    string // the server's own name+version
	Tools     int
	Err       string
}

// Status reports every configured server.
func (m *Manager) Status() []Status {
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make([]Status, 0, len(m.order))
	for _, name := range m.order {
		c := m.conns[name]
		s := Status{Name: name, Tools: len(c.tools)}
		if c.cfg.Command != "" {
			s.Transport, s.Target = "stdio", strings.Join(append([]string{c.cfg.Command}, c.cfg.Args...), " ")
		} else {
			s.Transport, s.Target = "http", c.cfg.URL
		}
		if c.client != nil {
			s.Connected = true
			s.Server = strings.TrimSpace(c.client.Info().Name + " " + c.client.Info().Version)
		}
		if c.lastErr != nil {
			s.Err = c.lastErr.Error()
		}
		out = append(out, s)
	}
	return out
}

// Call runs server's tool with args. Args come from the router and nothing
// else; there is no path here through which a note or a fact could travel.
func (m *Manager) Call(ctx context.Context, server, tool string, args map[string]any) (string, error) {
	m.mu.Lock()
	c := m.conns[server]
	m.mu.Unlock()
	if c == nil {
		return "", fmt.Errorf("%w: %s", ErrNoServer, server)
	}
	m.mu.Lock()
	cl, timeout, known := c.client, c.cfg.Timeout, false
	for _, t := range c.tools {
		if t.Name == tool {
			known = true
			break
		}
	}
	m.mu.Unlock()
	if cl == nil {
		return "", fmt.Errorf("mcp: %s is not connected", server)
	}
	// The discovered-and-filtered set is the second allowlist: even an enabled
	// store row cannot reach a tool the server stopped offering, or one
	// allow_tools excludes.
	if !known {
		return "", fmt.Errorf("mcp: %s offers no tool %q", server, tool)
	}
	cctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	return cl.CallTool(cctx, tool, args)
}

// Resources lists resources across connected servers.
func (m *Manager) Resources(ctx context.Context) []Resource {
	m.mu.Lock()
	clients := make([]*Client, 0, len(m.order))
	for _, name := range m.order {
		if cl := m.conns[name].client; cl != nil {
			clients = append(clients, cl)
		}
	}
	m.mu.Unlock()
	var out []Resource
	for _, cl := range clients {
		rs, err := cl.ListResources(ctx)
		if err != nil {
			continue // no resources capability; not an error worth logging per tick
		}
		out = append(out, rs...)
	}
	return out
}

// ReadResource reads one resource from one server.
func (m *Manager) ReadResource(ctx context.Context, server, uri string) (string, error) {
	m.mu.Lock()
	c := m.conns[server]
	var cl *Client
	var timeout time.Duration
	if c != nil {
		cl, timeout = c.client, c.cfg.Timeout
	}
	m.mu.Unlock()
	if cl == nil {
		return "", fmt.Errorf("%w: %s", ErrNoServer, server)
	}
	cctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	return cl.ReadResource(cctx, uri)
}

// Close shuts every connection down.
func (m *Manager) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, name := range m.order {
		if cl := m.conns[name].client; cl != nil {
			_ = cl.Close()
			m.conns[name].client = nil
		}
	}
	return nil
}