Cyrene_v2/internal/capability/mcp/capabilities.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"

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

// hostFeatures 表示宿主侧配置了哪些客户端能力（决定向服务器宣告什么）。
type hostFeatures struct {
	sampling    bool
	elicitation bool
	roots       bool
}

// hostFeatures 返回宿主侧已配置的客户端能力。
func (c *Client) hostFeatures() hostFeatures {
	return hostFeatures{
		sampling:    c.sampling != nil,
		elicitation: c.elicitation != nil,
		roots:       len(c.roots) > 0,
	}
}

// clientCapabilities 构造客户端能力声明：只宣告宿主真的具备的能力。
// 显式传非 nil 的 Capabilities 会关掉 SDK 默认宣告的 roots——本包默认不宣告任何
// 能力，避免对未声明能力的服务器暴露工作区根目录或代跑 LLM。
func (c *Client) clientCapabilities() *mcpsdk.ClientCapabilities {
	caps := &mcpsdk.ClientCapabilities{}
	if c.sampling != nil {
		caps.Sampling = &mcpsdk.SamplingCapabilities{}
	}
	if c.elicitation != nil {
		caps.Elicitation = &mcpsdk.ElicitationCapabilities{}
	}
	if len(c.roots) > 0 {
		caps.RootsV2 = &mcpsdk.RootCapabilities{}
	}
	return caps
}

// serverInfoFrom 把 SDK 会话的 initialize 结果整理成本包的 ServerInfo，
// 其中 capabilities 原文保留、对客户端可读（design/09 §1.1）。
func serverInfoFrom(session *mcpsdk.ClientSession, host hostFeatures) ServerInfo {
	ir := session.InitializeResult()
	if ir == nil {
		return ServerInfo{Capabilities: map[string]any{}}
	}
	info := ServerInfo{Instructions: ir.Instructions, ProtocolVersion: ir.ProtocolVersion}
	if ir.ServerInfo != nil {
		info.Name = ir.ServerInfo.Name
		info.Version = ir.ServerInfo.Version
	}
	info.Capabilities = capabilitiesMap(ir.Capabilities)
	info.Features = featuresFrom(info.Capabilities, host)
	return info
}

// capabilitiesMap 把 SDK 解析出的服务器能力序列化回原文 map 保留下来。
// 它只做无损转换，不做增删；极端情况下序列化失败则返回空 map（等价服务器未声明）。
func capabilitiesMap(caps *mcpsdk.ServerCapabilities) map[string]any {
	if caps == nil {
		return map[string]any{}
	}
	data, err := json.Marshal(caps)
	if err != nil {
		return map[string]any{}
	}
	out := map[string]any{}
	if err := json.Unmarshal(data, &out); err != nil {
		return map[string]any{}
	}
	return out
}

// featuresFrom 依据服务器 initialize 声明推断可用能力族（design/09 §1.1：
// 声明了才算可用，未声明不得调用）：
//
//   - tools / prompts / resources / logging / completions 直接按服务器是否声明该字段；
//   - 老式服务器完全不发 capabilities（nil 或空）时按 legacy 处理：只认 tools/list
//     是唯一发现入口（Tools=true），其余一律不可用；
//   - sampling / elicitation / roots 属于客户端侧能力，服务器不会声明「自身支持」它们。
//     因此判定为「宿主配置了回调」且「服务器声明过能力族（即确实是会主动交互的对端）」
//     两者同时成立才可用——什么都没声明的服务器不得回调这些能力。
func featuresFrom(caps map[string]any, host hostFeatures) ServerFeatures {
	declared := len(caps) > 0
	f := ServerFeatures{
		Tools:       !declared || hasCap(caps, "tools"),
		Prompts:     hasCap(caps, "prompts"),
		Resources:   hasCap(caps, "resources"),
		Logging:     hasCap(caps, "logging"),
		Completions: hasCap(caps, "completions"),
	}
	f.Sampling = host.sampling && declared
	f.Elicitation = host.elicitation && declared
	f.Roots = host.roots && declared
	return f
}

// hasCap 报告原始 capabilities 里是否声明了名为 name 的能力。
func hasCap(caps map[string]any, name string) bool {
	v, ok := caps[name]
	return ok && v != nil
}

// wrapSampling 包装宿主 sampling 回调：服务器未声明能力族时拒绝，
// 不静默放行（design/09 §1.1）。
func (c *Client) wrapSampling() func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) {
	h := c.sampling
	if h == nil {
		return nil
	}
	return func(ctx context.Context, req *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) {
		if !c.Features().Sampling {
			return nil, fmt.Errorf("mcp sampling %q: %w", c.cfg.Name, ErrCapabilityNotDeclared)
		}
		return h(ctx, req)
	}
}

// wrapElicitation 包装宿主 elicitation 回调，判定同 wrapSampling。
func (c *Client) wrapElicitation() func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) {
	h := c.elicitation
	if h == nil {
		return nil
	}
	return func(ctx context.Context, req *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) {
		if !c.Features().Elicitation {
			return nil, fmt.Errorf("mcp elicitation %q: %w", c.cfg.Name, ErrCapabilityNotDeclared)
		}
		return h(ctx, req)
	}
}