Cyrene_v2/internal/capability/mcp/client.go

package mcp

import (
	"context"
	"errors"
	"fmt"

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

	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// Connect 选择传输（未注入时）、建立官方 SDK 会话并完成 initialize 握手。
// 幂等：已连接时直接返回 nil；已关闭时返回 ErrClosed。
func (c *Client) Connect(ctx context.Context) error {
	same, err := c.begin()
	if err != nil {
		return err
	}
	if same {
		return nil
	}

	tr := c.tr
	if tr == nil {
		if tr, err = newTransport(c.cfg); err != nil {
			return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
		}
	}

	session, err := c.newSDKClient().Connect(ctx, tr, c.sessionOptions())
	if err != nil {
		return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, wrapConnectErr(err))
	}
	info := serverInfoFrom(session, c.hostFeatures())
	c.mu.Lock()
	c.session = session
	c.tr = tr
	c.state = stateConnected
	c.info = info
	c.mu.Unlock()

	c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name,
		"version", info.Version, "protocol", info.ProtocolVersion,
		"tools_declared", info.Features.Tools)
	return nil
}

// Close 关闭 SDK 会话并标记客户端已关闭。幂等。
func (c *Client) Close() error {
	c.mu.Lock()
	if c.state == stateClosed {
		c.mu.Unlock()
		return nil
	}
	session := c.session
	c.session = nil
	c.state = stateClosed
	c.mu.Unlock()

	if session != nil {
		if err := session.Close(); err != nil {
			return fmt.Errorf("mcp close %q: %w", c.cfg.Name, upstreamError(err))
		}
	}
	return nil
}

// ListTools 返回服务器暴露的工具（SDK 自动翻页）。
// 服务器在 initialize 里没有声明 tools 能力族时拒绝调用、不发出 RPC：
// design/09 §1.1「声明了才算可用，未声明不得调用」。
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
	session := c.connected()
	if session == nil {
		return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrNotConnected)
	}
	if !c.Features().Tools {
		return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrCapabilityNotDeclared)
	}
	out := make([]ToolInfo, 0, 8)
	for tool, err := range session.Tools(ctx, nil) {
		if err != nil {
			return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, upstreamError(err))
		}
		out = append(out, toolInfoFrom(tool))
	}
	return out, nil
}

// CallTool 调用服务器上的工具。args 必须是 JSON 对象（map[string]any）。
// 调用前用 schema 缓存里的 input schema 校验参数并补默认值。
func (c *Client) CallTool(ctx context.Context, name string, args any) (ToolResult, error) {
	arguments, err := normalizeArgs(args)
	if err != nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
	}
	return c.callTool(ctx, name, arguments, nil)
}

// callTool 是 CallTool 的内部实现，fallback 是缓存未命中时可用的 input schema
// （来自 tools/list 的发现结果）。
func (c *Client) callTool(ctx context.Context, name string, args map[string]any, fallback map[string]any) (ToolResult, error) {
	session := c.connected()
	if session == nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, ErrNotConnected)
	}
	prepared, err := c.schemas.Prepare(c.cfg.Name, name, fallback, args)
	if err != nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
	}
	res, err := session.CallTool(ctx, &mcpsdk.CallToolParams{Name: name, Arguments: prepared})
	if err != nil {
		return ToolResult{}, fmt.Errorf("mcp call tool %q on %q: %w", name, c.cfg.Name, upstreamError(err))
	}
	return ToolResult{
		Content:    toContentBlocks(res.Content),
		Structured: res.StructuredContent,
		IsError:    res.IsError,
	}, nil
}

// DiscoverTools 列出服务器工具，把每把工具包装成 sdk.Tool 并把 input schema
// 放进缓存，使 MCP server 看起来就是一个普通工具提供者。
func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) {
	tools, err := c.ListTools(ctx)
	if err != nil {
		return nil, err
	}
	out := make([]sdk.Tool, 0, len(tools))
	for _, t := range tools {
		c.cacheSchema(t)
		out = append(out, &mcpTool{client: c, info: t})
	}
	return out, nil
}

// cacheSchema 把一次发现到的 input schema 解析并放进缓存；解析失败只降级为
// 「调用时不校验」，不影响工具接入。
func (c *Client) cacheSchema(t ToolInfo) {
	if err := c.schemas.Put(c.cfg.Name, t.Name, t.InputSchema); err != nil {
		c.log.Warn("mcp tool input schema not resolvable; validation skipped",
			"name", c.cfg.Name, "tool", t.Name, "err", err)
	}
}

// handleToolListChanged 是 SDK 的 notifications/tools/list_changed 回调：
// 先让 schema 缓存失效（design/09 §1.1），再合并触发宿主的增量刷新。
// 它可能运行在 SDK 的读路径上，因此真正的刷新放到新 goroutine 里。
func (c *Client) handleToolListChanged(context.Context, *mcpsdk.ToolListChangedRequest) {
	c.schemas.Invalidate(c.cfg.Name)
	if c.onToolListChanged == nil {
		return
	}
	if !c.refreshPending.CompareAndSwap(false, true) {
		return
	}
	go func() {
		defer c.refreshPending.Store(false)
		c.onToolListChanged()
	}()
}

// newSDKClient 按配置构建官方 SDK 客户端：只宣告宿主真正具备的客户端能力。
func (c *Client) newSDKClient() *mcpsdk.Client {
	client := mcpsdk.NewClient(&mcpsdk.Implementation{Name: clientName, Version: clientVersion}, &mcpsdk.ClientOptions{
		Logger:                    c.log,
		Capabilities:              c.clientCapabilities(),
		CreateMessageHandler:      c.wrapSampling(),
		ElicitationHandler:        c.wrapElicitation(),
		ToolListChangedHandler:    c.handleToolListChanged,
		KeepAlive:                 c.keepAlive,
		KeepAliveFailureThreshold: c.keepAliveFailures,
	})
	if len(c.roots) > 0 {
		client.AddRoots(c.roots...)
	}
	return client
}

// sessionOptions 返回 SDK 会话选项；宿主固定了协议版本时才返回非 nil。
func (c *Client) sessionOptions() *mcpsdk.ClientSessionOptions {
	if c.preferredProtocolVersion == "" {
		return nil
	}
	return &mcpsdk.ClientSessionOptions{ProtocolVersion: c.preferredProtocolVersion}
}

// wrapConnectErr 把 SDK 报的「协议版本不支持」映射到本包哨兵错误，便于调用方判定。
//
// 它同时是握手失败的**脱敏点**：原始 SDK 错误（第三方 server 的自由文本）只经
// upstreamError 留下 code 与长度，绝不被 %w 带进文案 —— 而哨兵映射仍然成立，因为
// upstreamFailure 的 cause 保留了整条链，errors.Is 不受影响。
func wrapConnectErr(err error) error {
	safe := upstreamError(err)
	var wire *sdkjsonrpc.Error
	if errors.As(err, &wire) && wire.Code == mcpsdk.CodeUnsupportedProtocolVersion {
		return fmt.Errorf("%w: %w", ErrUnsupportedProtocolVersion, safe)
	}
	return safe
}

// normalizeArgs 把 agent 传来的参数值规整成 MCP arguments 对象（必须是 JSON 对象）。
func normalizeArgs(args any) (map[string]any, error) {
	if args == nil {
		return map[string]any{}, nil
	}
	if m, ok := args.(map[string]any); ok {
		return m, nil
	}
	return nil, fmt.Errorf("arguments must be a JSON object, got %T", args)
}