Cyrene_v2/internal/capability/mcp/mcp.go

// Package mcp 用官方 Model Context Protocol Go SDK 接入外部 MCP server，
// 并把服务器暴露的工具自动注册进 capability registry。
//
// 协议层（协议版本协商、stdio / SSE / Streamable HTTP 传输、keepalive ping、
// notifications/tools/list_changed 订阅）全部交给官方 SDK
// （github.com/modelcontextprotocol/go-sdk/mcp）——本包不自造协议，只做四件事：
//
//  1. 配置 → 传输选择（MCPConfig → SDK Transport，含 APIKey 鉴权）；
//  2. 工具自动发现 → 注册进 capability registry（Registry）；
//  3. 保留服务器声明的 capabilities 并推断可用能力族（ServerFeatures，
//     裁决 10.4.6 / design/09 §1.1「声明了才算可用」）；
//  4. tools/list input schema 缓存（SchemaCache）+ 调用前校验与默认值填充。
//
// 典型接法：
//
//	reg := registry.New()
//	mcpReg := mcp.NewRegistry(reg, mcp.MCPConfig{Name: "calc", Command: "mcp-calc"})
//	if err := mcpReg.Register(ctx); err != nil {
//	    log.Warn("mcp server skipped", "error", err)
//	}
//	defer mcpReg.Unregister(ctx)
//
// 注册后的工具就是普通 sdk.Tool：Dispatcher 每轮实时读 registry，
// 因此热接入的 MCP 工具下一轮模型请求即作为 Request.Tools 注入，模型可直接调用。
package mcp

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

	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"

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

// clientName / clientVersion 是 initialize 时上报给 MCP server 的 clientInfo。
const (
	clientName    = "cyrene"
	clientVersion = "0.1.0"
)

// keepAliveFailureThreshold 是 keepalive 关闭会话前容忍的连续 ping 失败次数默认值。
const keepAliveFailureThreshold = 3

// 本包对外哨兵错误，一律用 errors.Is 判定。
var (
	// ErrNotConnected 表示操作要求已建立的连接，但客户端尚未连接或已关闭。
	ErrNotConnected = errors.New("mcp: not connected")
	// ErrClosed 表示在已关闭的客户端上执行操作。
	ErrClosed = errors.New("mcp: closed")
	// ErrNoTransport 表示配置既没有 Command 也没有 URL，选不出传输。
	ErrNoTransport = errors.New("mcp: no transport configured")
	// ErrUnsupportedProtocolVersion 表示客户端与服务器无法就协议版本达成一致。
	ErrUnsupportedProtocolVersion = errors.New("mcp: unsupported protocol version")
	// ErrCapabilityNotDeclared 表示调用了服务器在 initialize 里未声明的能力族
	// （design/09 §1.1「声明了才算可用，未声明不得调用」）。
	ErrCapabilityNotDeclared = errors.New("mcp: capability not declared by server")
)

// MCPConfig 描述一台 MCP server 的连接方式。
// TransportType 非空时按其选择传输；为空则沿用历史推断：
// 有 Command 走 stdio，有 URL 走 SSE（TransportType 可显式改走 Streamable HTTP）。
type MCPConfig struct {
	// Name 是服务器标签，用于日志与工具去重前缀。
	Name string `env:"MCP_NAME"`
	// Command 是 stdio 模式下要拉起的可执行文件。
	Command string `env:"MCP_COMMAND"`
	// Args 是传给 Command 的参数。
	Args []string `env:"MCP_ARGS"`
	// URL 是 SSE 端点（TransportType="sse"）或 Streamable HTTP 端点（"http"/"streamable"）。
	URL string `env:"MCP_URL"`
	// TransportType 显式选择传输："stdio"、"sse"、"http"（等价 "streamable"）。
	TransportType string `env:"MCP_TRANSPORT"`
	// APIKey 是 http/sse 服务器用的 bearer 凭据，作为 Authorization 头发送。
	APIKey string `env:"MCP_API_KEY"`
	// Env 是 stdio 子进程的额外环境变量，叠加在父进程环境之上（同名以这里为准）。
	Env map[string]string `env:"MCP_ENV"`
	// Enabled 表示该服务器是否在启动时接入。
	Enabled bool `env:"MCP_ENABLED"`
}

// Config 是一组 MCP server 配置。
type Config struct {
	// Servers 是按顺序排列的服务器配置。
	Servers []MCPConfig
}

// ServerFeatures 是从服务器 initialize 声明里推断出的能力族可用性。
// 规则见 featuresFrom：声明了才算可用，未声明不得调用（design/09 §1.1）。
type ServerFeatures struct {
	// Tools 表示服务器声明提供工具（tools/list、tools/call 可用）。
	Tools bool
	// Prompts 表示服务器声明提供提示模板。
	Prompts bool
	// Resources 表示服务器声明提供资源。
	Resources bool
	// Logging 表示服务器声明会发日志消息。
	Logging bool
	// Completions 表示服务器声明支持参数自动补全。
	Completions bool
	// Sampling 表示允许服务器回调宿主代跑 LLM：宿主配了回调，
	// 且服务器声明过能力族（即确实是会主动交互的对端）才为真。
	Sampling bool
	// Elicitation 表示允许服务器回调宿主向用户追加提问，判定同 Sampling。
	Elicitation bool
	// Roots 表示允许向服务器暴露工作区根目录，判定同 Sampling。
	Roots bool
}

// ServerInfo 是 initialize 握手的结果，服务器声明的 capabilities 原样保留。
type ServerInfo struct {
	// Name 是服务器自报名称。
	Name string
	// Version 是服务器自报版本。
	Version string
	// Instructions 是服务器给出的使用说明（可注入系统提示）。
	Instructions string
	// ProtocolVersion 是协商后实际使用的 MCP 协议版本。
	ProtocolVersion string
	// Capabilities 是服务器声明的 capabilities 原文（未做删改），对客户端可读。
	Capabilities map[string]any
	// Features 是按 Capabilities 推断出的可用能力族。
	Features ServerFeatures
}

// SamplingHandler 由宿主提供：服务器请求客户端代跑一次 LLM（sampling/createMessage）时调用。
type SamplingHandler func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error)

// ElicitationHandler 由宿主提供：服务器向用户索要补充输入（elicitation/create）时调用。
type ElicitationHandler func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error)

// Client 是一台 MCP server 的客户端。它内部持有一个官方 SDK 会话，
// 并发安全：Connect/Close 互斥，其余调用由 SDK 串行化。
type Client struct {
	cfg MCPConfig
	log *slog.Logger

	mu      sync.Mutex
	state   state
	session *mcpsdk.ClientSession
	info    ServerInfo

	// tr 是注入或按 cfg 选出的传输；schemas 是 tools/list input schema 缓存。
	tr      mcpsdk.Transport
	schemas *SchemaCache

	// preferredProtocolVersion 非空时固定 initialize 请求的协议版本（测试/排障用）。
	preferredProtocolVersion string
	// keepAlive 是 keepalive ping 间隔，零值关闭；keepAliveFailures 是容忍失败数。
	keepAlive         time.Duration
	keepAliveFailures int

	sampling    SamplingHandler
	elicitation ElicitationHandler
	roots       []*mcpsdk.Root

	// onToolListChanged 在服务器发 notifications/tools/list_changed 时被调用
	// （新 goroutine 里，已被 refreshPending 合并），用于增量刷新 registry。
	onToolListChanged func()
	refreshPending    atomic.Bool
}

// state 是客户端生命周期。
type state int

const (
	stateNew state = iota
	stateConnected
	stateClosed
)

// Option 定制 New 构建的 Client。
type Option func(*Client)

// WithLogger 覆盖客户端日志器；给定的日志器会被标上 source=capability/mcp。
func WithLogger(l *slog.Logger) Option {
	return func(c *Client) {
		if l != nil {
			c.log = logger.WithSource(l, "capability/mcp")
		}
	}
}

// WithTransport 注入传输（官方 SDK 的 Transport），主要用于测试。
func WithTransport(t mcpsdk.Transport) Option {
	return func(c *Client) {
		c.tr = t
	}
}

// WithKeepAlive 打开 keepalive：每 interval 发一次 ping，连续失败 keepAliveFailures
// 次即关闭会话（由 SDK 执行）。间隔为零则关闭 keepalive；失败次数缺省用包默认值。
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 注册工具清单变更回调，由 registry 用于增量刷新。
func WithToolListChangedHandler(fn func()) Option {
	return func(c *Client) {
		c.onToolListChanged = fn
	}
}

// WithProtocolVersion 固定 initialize 请求的协议版本；不设置则用 SDK 最新版本，
// 服务器拒绝时由 SDK 自动回退到双方都支持的版本。
func WithProtocolVersion(v string) Option {
	return func(c *Client) {
		c.preferredProtocolVersion = v
	}
}

// WithSchemaCache 复用外部 schema 缓存；不设置则每个客户端自建一份。
func WithSchemaCache(cache *SchemaCache) Option {
	return func(c *Client) {
		if cache != nil {
			c.schemas = cache
		}
	}
}

// WithSamplingHandler 注册 sampling 回调。设了才向服务器宣告 sampling 能力；
// 服务器未声明能力族时回调会被拒绝（ErrCapabilityNotDeclared）。
func WithSamplingHandler(h SamplingHandler) Option {
	return func(c *Client) {
		c.sampling = h
	}
}

// WithElicitationHandler 注册 elicitation 回调。设了才向服务器宣告 elicitation 能力。
func WithElicitationHandler(h ElicitationHandler) Option {
	return func(c *Client) {
		c.elicitation = h
	}
}

// WithRoots 注册工作区根目录。设了才向服务器宣告 roots 能力。
func WithRoots(roots ...*mcpsdk.Root) Option {
	return func(c *Client) {
		c.roots = append([]*mcpsdk.Root(nil), roots...)
	}
}

// New 构建 cfg 对应的 Client，默认用丢弃日志器并自动选传输；
// 用 WithLogger / WithTransport 覆盖。
func New(cfg MCPConfig, opts ...Option) *Client {
	c := &Client{
		cfg:     cfg,
		log:     defaultLogger(),
		state:   stateNew,
		schemas: NewSchemaCache(),
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// defaultLogger 返回带 capability/mcp 来源标签的丢弃日志器；宿主应传 WithLogger。
func defaultLogger() *slog.Logger {
	base := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelWarn}))
	return logger.WithSource(base, "capability/mcp")
}

// Name 返回服务器标签。
func (c *Client) Name() string { return c.cfg.Name }

// ServerInfo 返回 initialize 握手结果；未连接时为零值。
func (c *Client) ServerInfo() ServerInfo {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.info
}

// Features 返回按服务器声明推断出的能力族可用性；未连接时全为 false。
func (c *Client) Features() ServerFeatures {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.info.Features
}

// SchemaCache 返回该客户端使用的 input schema 缓存。
func (c *Client) SchemaCache() *SchemaCache { return c.schemas }

// Session 返回底层官方 SDK 会话（未连接时为 nil），供需要 SDK 原生能力的调用方使用。
func (c *Client) Session() *mcpsdk.ClientSession {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.state != stateConnected {
		return nil
	}
	return c.session
}

// connected 返回已建立的 SDK 会话，未连接时返回 nil。
func (c *Client) connected() *mcpsdk.ClientSession {
	return c.Session()
}

// begin 做 Connect 前的状态检查：已连接返回 same=true（幂等），已关闭返回错误。
func (c *Client) begin() (same bool, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	switch c.state {
	case stateClosed:
		return false, fmt.Errorf("mcp connect %q: %w", c.cfg.Name, ErrClosed)
	case stateConnected:
		return true, nil
	}
	return false, nil
}