Maven/internal/mcp/manager.go

package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"sort"
	"strconv"
	"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. It is the FIRST wait: every consecutive
	// failure doubles it, up to MaxReconnectEvery.
	DefaultReconnectEvery = 30 * time.Second
	// MaxReconnectEvery caps the backoff. Without one, a permanently
	// misconfigured stdio server is exec'd once a minute forever, which is a
	// process spawn per minute in the logs and nothing that ever gets better.
	MaxReconnectEvery = 30 * time.Minute
	// DefaultMaxDescription bounds one tool description. It is written by a
	// server Maven does not control and it lands in two places that cannot
	// absorb an arbitrary blob: the resident model's 4096-token context, and a
	// table cell on /tools.
	DefaultMaxDescription = 400
)

// The two transports, as Status reports them to the web surface.
const (
	transportStdio = "stdio"
	transportHTTP  = "http"
)

var (
	// ErrNoServer — the named server is not configured.
	ErrNoServer = errors.New("mcp: no such server")
	// ErrNotConnected — the server is configured but nothing is dialed. Held
	// apart from ErrNoServer so a caller can say "that tool is not connected"
	// instead of drafting a proposal for a tool that already exists.
	ErrNotConnected = errors.New("mcp: server is not connected")
	// ErrToolGone — the server no longer offers this tool. An enabled row can
	// outlive the tool it names; this is what the act path sees when it does.
	ErrToolGone = errors.New("mcp: server no longer offers this tool")
)

// 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"`
	// Headers are sent verbatim on every request to a url server. This is how
	// a bearer token reaches a real remote server; the Vikunja one on loopback
	// needs none only because it is unauthenticated.
	Headers map[string]string `json:"-"`
	// Timeout bounds one call (0 ⇒ DefaultTimeout).
	Timeout time.Duration `json:"-"`
	// Enabled=false keeps a configured server described but dark.
	Enabled bool `json:"enabled"`
}

// argv is the stdio server's command line. It is argv and never a shell
// string, so the same slice serves both the exec and the /tools display.
func (c ServerConfig) argv() []string {
	return append([]string{c.Command}, c.Args...)
}

// 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
	fails    int // consecutive dial failures, for the backoff
}

// backoff is how long this connection waits before the next dial attempt:
// DefaultReconnectEvery doubled per consecutive failure, capped.
func (c *conn) backoff() time.Duration {
	d := DefaultReconnectEvery
	for i := 1; i < c.fails && d < MaxReconnectEvery; i++ {
		d *= 2
	}
	if d > MaxReconnectEvery {
		d = MaxReconnectEvery
	}
	return d
}

// 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()

	tr, err := m.openTransport(ctx, cfg)
	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.succeed(name, cl, tools)
	log.Printf("mcp: %s connected (%s %s), %d tool(s)", name, cl.Info().Name, cl.Info().Version, len(tools))
	return nil
}

// openTransport builds the door this server is configured for. A url server
// with no factory is one server's problem, reported here, so a bad block never
// stops the daemon.
func (m *Manager) openTransport(ctx context.Context, cfg ServerConfig) (transport, error) {
	if cfg.Command != "" {
		return newStdioTransport(ctx, cfg.argv(), cfg.Env, cfg.Dir)
	}
	if m.newPoster == nil {
		return nil, fmt.Errorf("server %q has a url but no http door was wired", cfg.Name)
	}
	poster, err := m.newPoster(cfg)
	if err != nil {
		return nil, err
	}
	return newHTTPTransport(poster, cfg.URL, cfg.Headers), nil
}

// succeed records a live connection and closes the one it replaces, so a
// re-dial does not leak the previous subprocess.
func (m *Manager) succeed(name string, cl *Client, tools []Tool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	c := m.conns[name]
	if c == nil {
		return
	}
	if c.client != nil {
		_ = c.client.Close()
	}
	c.client, c.tools, c.lastErr, c.fails = cl, tools, nil, 0
	c.dialedAt = time.Now()
}

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
		c.fails++
	}
}

// filterTools applies AllowTools and MaxTools and truncates descriptions.
//
// Over the cap WITHOUT allow_tools, the whole contribution is dropped. Taking
// the first N of a sorted list was deterministic but it handed the choice of
// which N to the server: a thirteenth tool named "aaa_" would push a tool that
// had already been discovered, proposed and maybe enabled out of the
// catalogue. Determinism was not the property worth buying. With allow_tools
// set, Kami named the tools, so the cap trims a list he chose.
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
		}
		t.Description = truncate(t.Description, DefaultMaxDescription)
		out = append(out, t)
	}
	if cfg.MaxTools > 0 && len(out) > cfg.MaxTools {
		if len(cfg.AllowTools) == 0 {
			log.Printf("mcp: %s offers %d tools, over the cap of %d — taking NONE of them, set allow_tools to choose or raise max_tools",
				cfg.Name, len(out), cfg.MaxTools)
			return nil
		}
		log.Printf("mcp: %s: allow_tools names %d tools, over the cap of %d — taking the first %d",
			cfg.Name, len(out), cfg.MaxTools, cfg.MaxTools)
		out = out[:cfg.MaxTools]
	}
	return out
}

// truncate bounds a server-written string. The ellipsis is there so a reader
// on /tools can tell the text was cut rather than written that way.
func truncate(s string, max int) string {
	r := []rune(s)
	if len(r) <= max {
		return s
	}
	return strings.TrimSpace(string(r[:max])) + "…"
}

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.
// The health check itself is done OUTSIDE m.mu, the way Resources already
// does it. alive() reaches into the transport, and a transport waiting on a
// silent subprocess would otherwise hold m.mu for as long as it waits, which
// blocks Tools, Status and Call for every other server too.
func (m *Manager) Refresh(ctx context.Context) {
	now := time.Now()
	type candidate struct {
		name  string
		cl    *Client
		ready bool
	}
	var cands []candidate
	m.mu.Lock()
	for _, name := range m.order {
		c := m.conns[name]
		cands = append(cands, candidate{name: name, cl: c.client, ready: now.Sub(c.lastTry) >= c.backoff()})
	}
	m.mu.Unlock()
	var stale []string
	for _, c := range cands {
		if !c.ready {
			continue
		}
		if c.cl == nil || !c.cl.alive() {
			stale = append(stale, c.name)
		}
	}
	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
}

// Connected — the names of servers that are dialed right now. A caller that
// wants to act on a tool's ABSENCE needs this: a tool missing from Tools()
// because its server is down is not a tool the server withdrew.
func (m *Manager) Connected() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []string
	for _, name := range m.order {
		if m.conns[name].client != nil {
			out = append(out, name)
		}
	}
	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 = transportStdio, strings.Join(c.cfg.argv(), " ")
		} else {
			s.Transport, s.Target = transportHTTP, 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) {
	cl, cfg, _, configured, known := m.lookup(server, tool)
	if !configured {
		return "", fmt.Errorf("%w: %s", ErrNoServer, server)
	}
	if cl == nil {
		return "", fmt.Errorf("%w: %s", ErrNotConnected, 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("%w: %s/%s", ErrToolGone, server, tool)
	}
	cctx, cancel := context.WithTimeout(ctx, cfg.Timeout)
	defer cancel()
	return cl.CallTool(cctx, tool, args)
}

// lookup reads one server's live state under the lock. Every call path needs
// the same four answers, and reading them in one critical section keeps a
// server that goes down mid-check from answering half yes.
func (m *Manager) lookup(server, tool string) (cl *Client, cfg ServerConfig, found Tool, configured, known bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	c := m.conns[server]
	if c == nil {
		return nil, ServerConfig{}, Tool{}, false, false
	}
	for _, t := range c.tools {
		if t.Name == tool {
			found, known = t, true
			break
		}
	}
	return c.client, c.cfg, found, true, known
}

// 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) {
	cl, cfg, _, configured, _ := m.lookup(server, "")
	if !configured {
		return "", fmt.Errorf("%w: %s", ErrNoServer, server)
	}
	// A configured server that is merely down is not an unknown server. Call
	// already keeps the two apart; reporting ErrNoServer here tells a caller
	// the resource can never exist, when the truth is "not right now".
	if cl == nil {
		return "", fmt.Errorf("%w: %s", ErrNotConnected, server)
	}
	cctx, cancel := context.WithTimeout(ctx, cfg.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
}

// ErrNeedsArgs — the tool requires arguments that a voice verb cannot supply.
var ErrNeedsArgs = errors.New("mcp: tool needs named arguments")

// CallPositional is the voice path's way in. The router gives an act a verb and
// a tail of positional words; an MCP tool wants a named-argument object. There
// is no general mapping between those two, and inventing one is exactly the
// improvisation this codebase refuses, so the rule is deliberately narrow:
//
//   - a tool with no required properties runs with no arguments (a spare tail
//     is ignored — "покажи проекты пожалуйста" should still list projects);
//   - a READ-ONLY tool NAMED IN allow_tools, with exactly one required
//     property, of type string or integer/number, gets the tail bound to it;
//   - anything else is refused with ErrNeedsArgs. Such a tool is still callable
//     with explicit arguments from the authed surface, where a human types
//     them.
//
// The refusal is the point, and the read-only condition on it was learned the
// hard way while testing against the Vikunja server: `update_task` requires
// only `task_id` and takes every other field as optional, so calling it with
// one guessed argument and no others BLANKED the fields it did not receive. A
// mutating tool therefore never gets a guessed argument — the one thing a
// partially-filled write can do is destroy what it did not mention. A mutating
// tool with nothing required is still fine: nothing was guessed, and it still
// goes through the confirm turn.
//
// The allow_tools condition is the second half, and it is there because
// readOnlyHint is the SERVER's claim about itself. It already buys one
// exemption (destructive=false, so no confirm turn); letting it buy argument
// binding as well means one lie converts a spoken utterance into an
// unconfirmed, argument-carrying write. A server advertising delete_project
// with readOnlyHint true and the description "show a project and its tasks"
// would be enough. So the binding half rests on something local instead: a
// name Kami typed into mavend.json. The tool name is not a defence — the
// router picks tools by name similarity and the description a human reads is
// server-written too.
func (m *Manager) CallPositional(ctx context.Context, server, tool string, args []string) (string, error) {
	cl, cfg, t, configured, known := m.lookup(server, tool)
	if !known {
		if !configured {
			return "", fmt.Errorf("%w: %s", ErrNoServer, server)
		}
		if cl == nil {
			return "", fmt.Errorf("%w: %s", ErrNotConnected, server)
		}
		return "", fmt.Errorf("%w: %s/%s", ErrToolGone, server, tool)
	}
	bindable := t.ReadOnly && contains(cfg.AllowTools, tool)
	named, err := bindPositional(t.InputSchema, args, bindable)
	if err != nil {
		return "", err
	}
	return m.Call(ctx, server, tool, named)
}

// bindPositional implements the rule documented on CallPositional. bind is the
// caller's verdict on whether a guessed argument is allowed at all: read-only
// AND named in allow_tools.
func bindPositional(schema json.RawMessage, args []string, bind bool) (map[string]any, error) {
	var s struct {
		Required   []string `json:"required"`
		Properties map[string]struct {
			Type string `json:"type"`
		} `json:"properties"`
	}
	if len(schema) > 0 {
		if err := json.Unmarshal(schema, &s); err != nil {
			return nil, fmt.Errorf("mcp: unreadable input schema: %w", err)
		}
	}
	switch len(s.Required) {
	case 0:
		return map[string]any{}, nil
	case 1:
		name := s.Required[0]
		prop, described := s.Properties[name]
		if !described {
			// required names it, properties does not describe it. The zero
			// value would make it a string, which is a guess about a guess.
			return nil, fmt.Errorf("%w: %q, which the schema never describes", ErrNeedsArgs, name)
		}
		if !bind {
			return nil, fmt.Errorf("%w: %q, and a guessed argument goes only to a read-only tool named in allow_tools", ErrNeedsArgs, name)
		}
		return bindOne(name, prop.Type, args)
	default:
		return nil, fmt.Errorf("%w: %s", ErrNeedsArgs, strings.Join(s.Required, ", "))
	}
}

// bindOne puts the whole positional tail in the one required property. The
// tail is spoken words, so only a scalar can hold it and anything else is
// refused rather than coerced.
func bindOne(name, typ string, args []string) (map[string]any, error) {
	tail := strings.TrimSpace(strings.Join(args, " "))
	if tail == "" {
		return nil, fmt.Errorf("%w: %q", ErrNeedsArgs, name)
	}
	switch typ {
	case "string", "":
		return map[string]any{name: tail}, nil
	case "integer", "number":
		n, err := strconv.ParseFloat(tail, 64)
		if err != nil {
			return nil, fmt.Errorf("%w: %q wants a number, got %q", ErrNeedsArgs, name, tail)
		}
		return map[string]any{name: n}, nil
	default:
		return nil, fmt.Errorf("%w: %q is a %s", ErrNeedsArgs, name, typ)
	}
}