Cyrene_v2/internal/capability/mcp/mcp.go

// Package mcp provides a minimal Model Context Protocol (MCP) client and the
// auto-registration glue that surfaces MCP server tools as capabilities.
//
// An MCP server is an external tool provider. This package speaks JSON-RPC 2.0
// to it over stdio (a spawned subprocess), SSE or Streamable HTTP, discovers the
// tools it exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and
// registers them into the capability registry so any character can invoke them
// (DevDocs/design/09 插件生态 §1).
//
// The client aligns with the MCP Go SDK's mature behaviour: it negotiates the
// protocol version during initialize (requesting the newest version first and
// downgrading to a mutually supported one), subscribes to
// notifications/tools/list_changed via its capabilities so the registry can
// re-sync its tool set, and runs an optional ping keepalive that closes a dead
// session so the host can rebuild it.
//
// It is deliberately dependency-light: only the standard library and the
// project packages (plugins/sdk, capability/registry, internal/logger) are
// used. The wire protocol is implemented directly, so no third-party MCP SDK is
// required.
//
// A caller usually wires it up as:
//
//	reg := registry.New()
//	mcpReg := mcp.New(reg, mcp.MCPConfig{Name: "calc", Command: "mcp-calc"})
//	if err := mcpReg.Register(ctx); err != nil {
//	    log.Fatal(err)
//	}
//	defer mcpReg.Unregister(ctx)
//
// which connects to the server, discovers its tools, and registers them. Call
// sites then invoke the tools by their MCP name through the registry.
package mcp

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"slices"
	"sync"
	"sync/atomic"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// MCP protocol versions this client can speak, newest first. The client advertises
// its preferred version (latestProtocolVersion) during initialize and accepts a
// compatible version from the server, downgrading to a mutually supported version
// when the server rejects or cannot match the request (DevDocs/design/09 §1.1).
const (
	// protocolVersion20251125 is the current preferred protocol version (the
	// MCP 2025-11-25 revision that introduced the Streamable HTTP migration).
	protocolVersion20251125 = "2025-11-25"
	// protocolVersion20241105 is the legacy protocol version, retained for
	// servers that have not yet migrated.
	protocolVersion20241105 = "2024-11-05"
	// latestProtocolVersion is what the client requests first during initialize.
	latestProtocolVersion = protocolVersion20251125
	// codeUnsupportedProtocolVersion is the JSON-RPC error code a server returns
	// during initialize when it does not support the requested protocol version
	// and wishes to advertise the versions it does support.
	codeUnsupportedProtocolVersion = -32022
	// codeMethodNotFound is the JSON-RPC error code a server returns when it does
	// not implement a requested method (e.g. a keepalive ping).
	codeMethodNotFound = -32601
	// maxInitializeAttempts bounds the initialize version-negotiation retries
	// (the first attempt plus one downgrade).
	maxInitializeAttempts = 2
	// notificationToolsListChanged is the server→client notification signalling
	// that the tool list has changed.
	notificationToolsListChanged = "notifications/tools/list_changed"
	// keepAliveFailureThreshold is the number of consecutive ping failures the
	// keepalive loop tolerates before closing the connection for a rebuild.
	keepAliveFailureThreshold = 3
)

// supportedProtocolVersions lists every protocol version this client understands,
// newest first. Version negotiation picks the first entry that the server also
// supports.
var supportedProtocolVersions = []string{protocolVersion20251125, protocolVersion20241105}

// isSupportedVersion reports whether v is a protocol version this client speaks.
func isSupportedVersion(v string) bool {
	return slices.Contains(supportedProtocolVersions, v)
}

// negotiateMutuallySupportedVersion returns the highest protocol version that is
// supported by both this client and the server (which supports supported). It
// returns "" when there is no overlap.
func negotiateMutuallySupportedVersion(supported []string) string {
	for _, v := range supportedProtocolVersions {
		if slices.Contains(supported, v) {
			return v
		}
	}
	return ""
}

// highestCommonVersion returns the highest protocol version this client supports
// that is not newer than v (used to downgrade when the server picks a version the
// client cannot speak). It returns "" when no acceptable version exists.
func highestCommonVersion(v string) string {
	for _, s := range supportedProtocolVersions {
		if s <= v {
			return s
		}
	}
	return ""
}

// Sentinel errors returned by this package. Test with errors.Is.
var (
	// ErrNotConnected is returned when an operation requires an established
	// connection but the client is not yet connected or is closed.
	ErrNotConnected = errors.New("mcp: not connected")
	// ErrClosed is returned when operating on a client or transport that has
	// been closed.
	ErrClosed = errors.New("mcp: closed")
	// ErrNoTransport is returned when the configuration selects no transport
	// (neither Command nor URL is set).
	ErrNoTransport = errors.New("mcp: no transport configured")
	// ErrRPC is returned when the MCP server replies with a JSON-RPC error
	// object. The concrete server error is wrapped and accessible via
	// errors.Is/errors.As on an *rpcError.
	ErrRPC = errors.New("mcp: rpc error")
	// ErrUnsupportedProtocolVersion is returned when the client and server
	// cannot agree on a mutually supported protocol version after the
	// negotiation retries are exhausted.
	ErrUnsupportedProtocolVersion = errors.New("mcp: unsupported protocol version")
)

// MCPConfig describes a single MCP server connection. The transport is selected
// by TransportType when set, otherwise inferred from the fields: stdio when
// Command is set, SSE when URL is set. A blank TransportType keeps that legacy
// inference so existing configuration keeps working unchanged.
type MCPConfig struct {
	// Name is the human-readable label of the server, used in logs and as the
	// prefix for registering its tools uniquely.
	Name string `env:"MCP_NAME"`
	// Command is the executable path to spawn for a stdio MCP server.
	Command string `env:"MCP_COMMAND"`
	// Args are the arguments passed to Command.
	Args []string `env:"MCP_ARGS"`
	// URL is the endpoint of an MCP server (SSE stream endpoint, or the
	// Streamable HTTP endpoint when TransportType is "http").
	URL string `env:"MCP_URL"`
	// TransportType selects the transport explicitly: "stdio", "sse" or
	// "http" (Streamable HTTP). Empty means infer from Command/URL.
	TransportType string `env:"MCP_TRANSPORT"`
	// Enabled reports whether this server should be activated at startup.
	Enabled bool `env:"MCP_ENABLED"`
}

// Config holds the set of configured MCP servers.
type Config struct {
	// Servers is the ordered list of MCP server configurations.
	Servers []MCPConfig
}

// ServerInfo describes the MCP server reported by the initialize handshake.
type ServerInfo struct {
	// Name is the server-provided name.
	Name string
	// Version is the server-provided version string.
	Version string
}

// ToolInfo is a tool exposed by an MCP server.
type ToolInfo struct {
	// Name is the tool identifier used to invoke it.
	Name string
	// Description explains what the tool does.
	Description string
	// InputSchema is the tool's JSON Schema, typically
	// {"type":"object","properties":{...}}.
	InputSchema map[string]any
}

// ContentBlock is a single item in a tool call result (MCP content block).
type ContentBlock struct {
	// Type is the block type, most commonly "text".
	Type string `json:"type"`
	// Text is the textual payload.
	Text string `json:"text"`
}

// ToolResult is the structured result of a tools/call.
type ToolResult struct {
	// Content holds the response content blocks.
	Content []ContentBlock
	// IsError reports whether the server flagged the call as an error.
	IsError bool
}

// Client is a minimal MCP client bound to a single server described by cfg. It
// is safe for concurrent use; all calls are serialized internally.
type Client struct {
	cfg   MCPConfig
	log   *slog.Logger
	tr    Transport
	mu    sync.Mutex
	state state
	info  ServerInfo

	// negotiatedVersion is the protocol version agreed during initialize; it is
	// sent back to the server in the Mcp-Protocol-Version header by the HTTP
	// transport.
	negotiatedVersion string

	// preferredProtocolVersion, when non-empty, overrides latestProtocolVersion
	// as the version requested first during initialize. Tests use it to drive the
	// negotiation path deterministically.
	preferredProtocolVersion string

	// keepAlive is the ping interval for the connection keepalive; zero disables
	// it. keepAliveFailures is the tolerated consecutive-failure count.
	keepAlive         time.Duration
	keepAliveFailures int
	keepaliveCancel   context.CancelFunc

	// onToolListChanged is invoked (in a fresh goroutine) when the server sends
	// notifications/tools/list_changed, so the host can refresh its tool set.
	onToolListChanged func()
	refreshPending    atomic.Bool
}

// protocolVersion returns the protocol version negotiated by initialize, or ""
// before a successful handshake.
func (c *Client) protocolVersion() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.negotiatedVersion
}

// setProtocolVersion records the version negotiated by initialize.
func (c *Client) setProtocolVersion(v string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.negotiatedVersion = v
}

// capabilities returns the client capability declaration for the given protocol
// version. It advertises tools.listChanged so the server notifies the client when
// the tool list changes.
func (c *Client) capabilities(version string) map[string]any {
	caps := map[string]any{}
	if c.onToolListChanged != nil {
		caps["tools"] = map[string]any{"listChanged": true}
	}
	return caps
}

// state tracks the client lifecycle.
type state int

const (
	stateNew state = iota
	stateConnected
	stateClosed
)

// Transport is the JSON-RPC message exchange with an MCP server over a concrete
// transport (stdio subprocess, SSE or Streamable HTTP). Implementations must be
// safe for concurrent use. This is exported so tests and hosts can inject a fake
// transport.
type Transport interface {
	// Start makes the transport ready to exchange messages.
	Start(ctx context.Context) error
	// Call sends a JSON-RPC request and decodes the matching response into
	// result. Params is marshaled into the request body; result is unmarshaled
	// from the response's result field.
	Call(ctx context.Context, method string, params any, result any) error
	// Close releases every underlying resource (subprocess, streams, sockets).
	Close() error
}

// notifier is implemented by transports that can send a JSON-RPC notification
// (a request with no id and no response).
type notifier interface {
	notify(ctx context.Context, method string, params any) error
}

// serverNotifier is implemented by transports that can deliver server-initiated
// notifications (messages with no id, e.g. notifications/tools/list_changed) back
// to the client. The handler is invoked from the transport's read path, so it must
// not block or re-enter the same transport synchronously.
type serverNotifier interface {
	setOnNotification(fn func(method string))
}

// attachNotificationSink wires a server→client notification handler into tr, if
// tr supports one.
func attachNotificationSink(tr Transport, fn func(method string)) {
	if on, ok := tr.(serverNotifier); ok {
		on.setOnNotification(fn)
	}
}

// protocolProvider is implemented by transports that include the negotiated MCP
// protocol version in requests (the Mcp-Protocol-Version header of the HTTP
// transport).
type protocolProvider interface {
	setProtocolProvider(fn func() string)
}

// attachProtocolProvider wires the negotiated-version getter into tr, if tr
// supports one.
func attachProtocolProvider(tr Transport, fn func() string) {
	if p, ok := tr.(protocolProvider); ok {
		p.setProtocolProvider(fn)
	}
}

// Option customizes a Client built by New.
type Option func(*Client)

// WithLogger overrides the client's logger. The given logger is tagged with the
// source capability/mcp.
func WithLogger(l *slog.Logger) Option {
	return func(c *Client) {
		if l != nil {
			c.log = logger.WithSource(l, "capability/mcp")
		}
	}
}

// WithTransport overrides the transport used by the client. This is primarily
// for tests that inject a fake transport instead of spawning a subprocess or
// dialing SSE.
func WithTransport(t Transport) Option {
	return func(c *Client) {
		c.tr = t
	}
}

// WithKeepAlive enables the connection keepalive, sending an MCP ping every
// interval and closing the session after keepAliveFailures consecutive missed
// pings. A zero interval leaves keepalive disabled. When keepAliveFailures is
// omitted, the package default (keepAliveFailureThreshold) is used.
func WithKeepAlive(interval time.Duration, keepAliveFailures ...int) Option {
	failures := keepAliveFailureThreshold
	if len(keepAliveFailures) > 0 && keepAliveFailures[0] > 0 {
		failures = keepAliveFailures[0]
	}
	return func(c *Client) {
		c.keepAlive = interval
		c.keepAliveFailures = failures
	}
}

// WithToolListChangedHandler registers a handler that is invoked (in a fresh
// goroutine) whenever the server sends notifications/tools/list_changed, letting
// the host refresh its registered tool set.
func WithToolListChangedHandler(fn func()) Option {
	return func(c *Client) {
		c.onToolListChanged = fn
	}
}

// WithProtocolVersion overrides the protocol version requested first during
// initialize. It is primarily for tests that drive the negotiation path
// deterministically; callers should normally leave it unset to use
// latestProtocolVersion.
func WithProtocolVersion(v string) Option {
	return func(c *Client) {
		c.preferredProtocolVersion = v
	}
}

// New builds a Client for cfg. By default it uses a discard logger and selects a
// stdio or SSE transport from cfg; pass WithLogger and/or WithTransport to
// override.
func New(cfg MCPConfig, opts ...Option) *Client {
	c := &Client{
		cfg:   cfg,
		log:   defaultLogger(),
		state: stateNew,
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// defaultLogger returns a discard logger tagged with the mcp source. Hosts
// should prefer WithLogger so records are observable.
func defaultLogger() *slog.Logger {
	base := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelWarn}))
	return logger.WithSource(base, "capability/mcp")
}

// transport returns the established transport, or nil when not connected.
func (c *Client) transport() Transport {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.state != stateConnected || c.tr == nil {
		return nil
	}
	return c.tr
}