AutoZero/mcp/mcp.go

// Package mcp 实现最小 MCP（Model Context Protocol）stdio 客户端。
// 手写 JSON-RPC 2.0（换行分隔帧），不依赖官方 SDK——与项目零重依赖原则一致。
//
// 协议**年代（era）**：2026-07-28 起 MCP 变成无状态，没有 initialize 握手——
//
//   - modern：每个请求在 params._meta 里自带 protocolVersion / clientCapabilities /
//     clientInfo，结果带 resultType；版本不匹配时服务端回 -32022 并列出自己支持的版本。
//   - legacy（2025-11-25 及更早）：initialize → notifications/initialized 握手，
//     版本在握手时协商一次。
//
// 本客户端是 **dual-era**：Start 时先发 server/discover 探测年代，探测不到再回退
// legacy 握手。只认新版的话，线上现存的 server（几乎全是 legacy）会全部不可用。
//
// 本次**没有**实现的部分（都是明确取舍，不是遗漏）：
//   - Streamable HTTP 传输：本包只有 stdio。
//   - subscriptions/listen（替代旧的通知订阅）：我们在启动时列一次工具表，不消费
//     服务端的工具变更通知，所以 server 动态增删工具不会反映出来。
//   - elicitation/sampling/roots：我们声明空的 clientCapabilities，合规服务端因此
//     **不得**向我们发这类请求；真收到 inputRequests 就报可操作的错误。
package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"sort"
	"strings"
	"sync"
	"time"

	"git.gopher.ink/mirrors2/AutoZero/diag"
	"git.gopher.ink/mirrors2/AutoZero/execx"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

const (
	// protocolVersion 是首选（modern）协议版本。
	protocolVersion = "2026-07-28"
	// legacyVersion 是回退握手时声明的版本；legacy 服务端会在 initialize 的响应里
	// 汇报它实际选用的版本。
	legacyVersion = "2025-11-25"
	// errUnsupportedProtocolVersion 是规范定义的错误码：服务端不认识/不支持请求里
	// 声明的版本，并在 data.supported 里列出自己支持的版本。
	errUnsupportedProtocolVersion = -32022

	// discoverTimeout 是年代探测的超时。要短：探测失败就走 legacy 回退，而"服务端
	// 静默"这种最坏情况每个 server 都要等满这个时间（N 个 server 就是 N 倍）。
	discoverTimeout  = 3 * time.Second
	initTimeout      = 15 * time.Second
	defaultCallLimit = 120 * time.Second
	// maxMRTRRetries 限制 input_required 的重试次数。服务端可以反复要求输入，
	// 我们必须有界，否则一次工具调用可能无限空转。
	maxMRTRRetries = 3

	clientName    = "autogo"
	clientVersion = "0.1.0"
)

// _meta 键名（规范要求带前缀，由 MCP 规范方保留）。
const (
	metaProtocolVersion    = "io.modelcontextprotocol/protocolVersion"
	metaClientCapabilities = "io.modelcontextprotocol/clientCapabilities"
	metaClientInfo         = "io.modelcontextprotocol/clientInfo"
	metaServerInfo         = "io.modelcontextprotocol/serverInfo"
)

// 年代取值。
const (
	eraModern = "modern"
	eraLegacy = "legacy"
)

// knownVersions 是我们能讲的协议版本，按新→旧排列（用于协商挑版本）。
var knownVersions = []string{protocolVersion, legacyVersion, "2025-06-18", "2025-03-26"}

// Config 是一个 MCP server 的启动配置（兼容 Claude Code 的 mcpServers 格式）。
//
// 两种传输二选一：
//   - stdio：command/args/env，起本地子进程；
//   - http：url（+ 可选 headers），走 Streamable HTTP（2026-07-28 无会话版；
//     旧版会在 initialize 时下发 Mcp-Session-Id，客户端会自动带上）。
type Config struct {
	Command string            `json:"command"`
	Args    []string          `json:"args"`
	Env     map[string]string `json:"env"`
	URL     string            `json:"url"`
	Headers map[string]string `json:"headers"`
	// Type 只用于识别"明确声明为已废弃的 http+sse（type="sse"）以便给出提示。
	Type string `json:"type"`
}

// Kind 返回配置指定的传输类型：stdio | http | ""（两种都不是）。
func (c Config) Kind() string {
	switch {
	case c.Command != "":
		return "stdio"
	case c.URL != "":
		return "http"
	default:
		return ""
	}
}

// LoadConfigs 合并多个 mcp.json（后者同名覆盖前者）；文件缺失或损坏时跳过。
func LoadConfigs(paths ...string) map[string]Config {
	type fileFormat struct {
		MCPServers map[string]Config `json:"mcpServers"`
	}
	out := map[string]Config{}
	for _, p := range paths {
		data, err := os.ReadFile(p)
		if err != nil {
			continue
		}
		var f fileFormat
		if err := json.Unmarshal(data, &f); err != nil {
			continue
		}
		for name, cfg := range f.MCPServers {
			// 明确声明 type=sse 的是**已废弃的 HTTP+SSE 传输**：即使带 url 也不能当成
			// Streamable HTTP 用——那个传输靠 GET 打开 SSE 流拿到 endpoint，POST 语义
			// 完全不同，硬当成 http 只会在运行时报出看不懂的错。
			if cfg.Type == "sse" {
				log.Printf("[mcp] 跳过 %q：type=sse 是已废弃的 HTTP+SSE 传输，本客户端只支持 stdio 与 Streamable HTTP", name)
				continue
			}
			switch cfg.Kind() {
			case "stdio", "http":
				out[name] = cfg
			default:
				// 以前这里要求 command 非空，url 型条目会被**静默**丢掉（用户的
				// westock-mcp 就是这样消失的）。现在至少说明原因。
				log.Printf("[mcp] 跳过 %q：既没有 command（stdio）也没有 url（http），无法判断传输类型", name)
			}
		}
	}
	return out
}

type rpcError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

// rpcCallError 把 JSON-RPC 错误做成 Go error，并**保留错误码与 data**。
// 年代探测要靠 -32022 及其 data.supported 判断，把错误压成字符串就没法判了。
type rpcCallError struct {
	method string
	code   int
	msg    string
	data   json.RawMessage
}

func (e *rpcCallError) Error() string {
	if len(e.data) > 0 {
		return fmt.Sprintf("%s: %s (code %d, data %s)", e.method, e.msg, e.code, e.data)
	}
	return fmt.Sprintf("%s: %s (code %d)", e.method, e.msg, e.code)
}

type rpcMsg struct {
	JSONRPC string `json:"jsonrpc"`
	// ID 用 RawMessage 原样接收：规范示例里 id 可以是字符串，用 *int64 会让
	// 字符串 id 的消息解析失败。
	ID     json.RawMessage `json:"id"`
	Method string          `json:"method"`
	Result json.RawMessage `json:"result"`
	Error  *rpcError       `json:"error"`
}

// ToolInfo 是 server 声明的一个工具。
type ToolInfo struct {
	Name        string
	Description string
	InputSchema map[string]any
}

// ServerInfo 是服务端自报的身份（仅用于日志/展示，规范明确不可用于安全决策）。
type ServerInfo struct {
	Name    string
	Version string
}

// Client 是与一个 MCP server 进程的连接。
type Client struct {
	tr     transport
	mu     sync.Mutex
	nextID int64
	closed bool

	// 年代与协商结果：进程生命周期内缓存（规范要求：年代是服务端的属性，
	// 不是单次请求的）。Start 返回后不再改写。
	era     string
	version string
	server  ServerInfo
}

// Start 建立与一个 MCP server 的连接并确定协议年代。name 只用于日志。
//
// 传输按配置二选一：command → stdio 子进程；url → Streamable HTTP。
//
// 诊断走 diag（默认 os.Stderr，TUI 会把它换到日志文件）：MCP 子进程的 stderr 直接
// 接终端的话，在 TUI 的 alt-screen 下会打穿界面——而 MCP 的惯例恰恰是往 stderr 打日志。
func Start(name string, cfg Config) (*Client, error) {
	tr, err := openTransport(cfg)
	if err != nil {
		return nil, err
	}
	// 默认按 modern 起步：探测请求自己也要带 _meta，探测失败再降级。
	c := &Client{tr: tr, era: eraModern, version: protocolVersion}
	if err := c.detectEra(); err != nil {
		_ = tr.Close()
		return nil, err
	}
	diag.Printf("[mcp:%s] %s(%s) 协议 %s server=%s %s\n",
		name, tr.Kind(), cfg.Command+cfg.URL, c.versionOrDash(), c.server.Name, c.server.Version)
	return c, nil
}

// openTransport 按配置建立传输：stdio 或 Streamable HTTP。
func openTransport(cfg Config) (transport, error) {
	switch cfg.Kind() {
	case "stdio":
		cmd := exec.Command(cfg.Command, cfg.Args...)
		// 自成进程组 + 有限 Wait：Close 时 kill(-pgid) 才带得走 npx/uvx 包装
		// 后面的真 server，也不会让 Wait 卡在谁手里（见 internal/execx）。
		execx.SeparateGroup(cmd)
		cmd.WaitDelay = execx.WaitGrace
		cmd.Env = os.Environ()
		for k, v := range cfg.Env {
			cmd.Env = append(cmd.Env, k+"="+v)
		}
		// 子进程的 stderr 交给 diag.Indirect：启动这一刻 TUI 还没建好日志文件，
		// 传具体 writer 会把 os.Stderr 钉死（见 diag 包注释）。
		cmd.Stderr = diag.Indirect()
		stdin, err := cmd.StdinPipe()
		if err != nil {
			return nil, err
		}
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			return nil, err
		}
		if err := cmd.Start(); err != nil {
			return nil, err
		}
		return newStdioTransport(cmd, stdin, stdout), nil
	case "http":
		return newHTTPTransport(cfg.URL, cfg.Headers), nil
	default:
		return nil, errors.New("配置既没有 command 也没有 url，无法确定传输类型")
	}
}

// Era 返回探测到的协议年代（modern|legacy）。
func (c *Client) Era() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.era
}

// Version 返回协商后的协议版本。
func (c *Client) Version() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.version
}

func (c *Client) versionOrDash() string {
	if v := c.Version(); v != "" {
		return v
	}
	return "-"
}

// detectEra 按规范确定服务端的协议年代。
//
// stdio 侧的规则（原文）：先用 server/discover 探测，**只有"可识别的现代错误"
// 才算现代**；其它一切（方法不存在 / 超时 / 静默 / 无法解析的输出）都判为 legacy，
// 回退 initialize 握手。
func (c *Client) detectEra() error {
	ctx, cancel := context.WithTimeout(context.Background(), discoverTimeout)
	raw, err := c.rpc(ctx, "server/discover", map[string]any{})
	cancel()

	if err == nil {
		var res struct {
			SupportedVersions []string `json:"supportedVersions"`
			Meta              struct {
				Info ServerInfo `json:"io.modelcontextprotocol/serverInfo"`
			} `json:"_meta"`
		}
		_ = json.Unmarshal(raw, &res)
		c.setServer(res.Meta.Info)
		// 服务端自报支持的版本里挑一个我们能讲的；列表为空视为"接受我声明的版本"。
		v := pickVersion(res.SupportedVersions)
		if v == "" && len(res.SupportedVersions) == 0 {
			v = protocolVersion
		}
		if v != protocolVersion {
			// 只支持 legacy 修订版的服务端：那条路要走握手，而不是"带 _meta 声明旧版本"。
			return c.initLegacy()
		}
		c.setVersion(v)
		return nil
	}

	// 鉴权失败不是年代问题：直接报出去，别去试 legacy 握手。
	var ae *authRequiredError
	if errors.As(err, &ae) {
		return ae
	}
	var ce *rpcCallError
	if errors.As(err, &ce) && ce.code == errUnsupportedProtocolVersion {
		// 现代服务端，只是不支持我们声明的版本：按规范从 data.supported 里挑一个。
		if v := pickVersion(supportedFromError(ce)); v == protocolVersion {
			c.setVersion(v)
			return nil
		}
		return c.initLegacy()
	}
	// 其它一切（method not found / 超时 / 静默 / 垃圾输出）都按 legacy 处理。
	return c.initLegacy()
}

// initLegacy 走握手式协议：initialize → notifications/initialized。
// **用新的 ctx**：探测可能已经把上一个超时耗尽了。
func (c *Client) initLegacy() error {
	c.setEraLegacy()
	ctx, cancel := context.WithTimeout(context.Background(), initTimeout)
	defer cancel()
	raw, err := c.rpc(ctx, "initialize", map[string]any{
		"protocolVersion": legacyVersion,
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": clientName, "version": clientVersion},
	})
	if err != nil {
		return fmt.Errorf("initialize: %w", err)
	}
	var res struct {
		ProtocolVersion string     `json:"protocolVersion"`
		ServerInfo      ServerInfo `json:"serverInfo"`
	}
	_ = json.Unmarshal(raw, &res)
	c.setVersion(res.ProtocolVersion)
	c.setServer(res.ServerInfo)
	_ = c.notify("notifications/initialized")
	return nil
}

// supportedFromError 从 -32022 的 data.supported 里取支持的版本列表。
func supportedFromError(e *rpcCallError) []string {
	var data struct {
		Supported []string `json:"supported"`
	}
	if len(e.data) == 0 || json.Unmarshal(e.data, &data) != nil {
		return nil
	}
	return data.Supported
}

// pickVersion 从服务端声明的版本里挑一个我们也能讲的（优先新版本）。
func pickVersion(server []string) string {
	for _, ours := range knownVersions {
		for _, theirs := range server {
			if theirs == ours {
				return ours
			}
		}
	}
	return ""
}

func (c *Client) setVersion(v string) {
	c.mu.Lock()
	c.version = v
	c.mu.Unlock()
}

func (c *Client) setServer(s ServerInfo) {
	c.mu.Lock()
	c.server = s
	c.mu.Unlock()
}

// setEraLegacy 把年代降级为 legacy（此后请求不再带 _meta，改走握手语义）。
func (c *Client) setEraLegacy() {
	c.mu.Lock()
	c.era = eraLegacy
	c.mu.Unlock()
}

// isModern 报告当前是否按 modern 协议说话。
func (c *Client) isModern() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.era == eraModern
}

// ListTools 拉取 server 的工具列表。
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
	raw, err := c.rpc(ctx, "tools/list", map[string]any{})
	if err != nil {
		return nil, err
	}
	var res struct {
		Tools []struct {
			Name        string         `json:"name"`
			Description string         `json:"description"`
			InputSchema map[string]any `json:"inputSchema"`
		} `json:"tools"`
	}
	if err := json.Unmarshal(raw, &res); err != nil {
		return nil, err
	}
	out := make([]ToolInfo, 0, len(res.Tools))
	for _, t := range res.Tools {
		out = append(out, ToolInfo{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
	}
	return out, nil
}

// CallTool 调用 server 工具，返回拼接后的文本内容；isError 结果转为 error。
func (c *Client) CallTool(ctx context.Context, name, argsJSON string) (string, error) {
	if !json.Valid([]byte(argsJSON)) {
		return "", errors.New("invalid arguments JSON")
	}
	if strings.TrimSpace(argsJSON) == "" {
		argsJSON = "{}"
	}
	raw, err := c.rpc(ctx, "tools/call", map[string]any{
		"name": name, "arguments": json.RawMessage(argsJSON),
	})
	if err != nil {
		return "", err
	}
	var res struct {
		Content []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		} `json:"content"`
		IsError bool `json:"isError"`
	}
	if err := json.Unmarshal(raw, &res); err != nil {
		return "", err
	}
	var texts []string
	for _, c := range res.Content {
		if c.Text != "" {
			texts = append(texts, c.Text)
		}
	}
	joined := strings.Join(texts, "\n")
	if res.IsError {
		return "", errors.New(joined)
	}
	return joined, nil
}

// AsTool 把一个 server 工具包装为 Registry 工具（命名 mcp__<server>__<tool>，
// 与 Claude Code 约定一致）。语义未知 → fail-closed（视为副作用、不可并发）。
func (c *Client) AsTool(server, name, description string, schema map[string]any) *tools.Tool {
	full := "mcp__" + server + "__" + name
	if description != "" {
		description = "[mcp:" + server + "] " + description
	} else {
		description = "[mcp:" + server + "] tool " + name
	}
	if schema == nil {
		schema = map[string]any{"type": "object"}
	}
	return &tools.Tool{
		Name:        full,
		Description: description,
		Parameters:  schema,
		SideEffect:  true,
		Run: func(ctx context.Context, input string) (tools.Result, error) {
			cctx, cancel := context.WithTimeout(ctx, defaultCallLimit)
			defer cancel()
			out, err := c.CallTool(cctx, name, input)
			return tools.TextResult(out), err
		},
	}
}

// Close 关闭传输（stdio 会终止子进程；HTTP 无连接可关）。
func (c *Client) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.mu.Unlock()
	return c.tr.Close()
}

// rpc 发送请求并等待响应，内部处理 MRTR（input_required）的有界重试。
//
// MRTR 放在这里而不是暴露给调用方：requestState 是**服务端在调用过程中**给的，
// 上层（ListTools/CallTool/AsTool）根本无从提供，只能是本层的重试循环。
func (c *Client) rpc(ctx context.Context, method string, params map[string]any) (json.RawMessage, error) {
	var requestState string
	for attempt := 1; ; attempt++ {
		raw, err := c.rpcOnce(ctx, method, params, requestState)
		if err != nil {
			return nil, err
		}
		var env struct {
			ResultType    string                     `json:"resultType"`
			RequestState  string                     `json:"requestState"`
			InputRequests map[string]json.RawMessage `json:"inputRequests"`
		}
		if json.Unmarshal(raw, &env) != nil {
			return raw, nil // 不是对象（例如个别实现的裸值）：原样返回
		}
		// 缺 resultType 按 complete 处理——兼容 2026-07-28 之前的服务端。
		if env.ResultType != "input_required" {
			return raw, nil
		}
		if len(env.InputRequests) > 0 {
			// 我们声明的是空 clientCapabilities，合规服务端不得发 elicitation/sampling/
			// roots 请求。真收到就说清楚我们满足不了、它要什么。
			keys := make([]string, 0, len(env.InputRequests))
			for k := range env.InputRequests {
				keys = append(keys, k)
			}
			sort.Strings(keys)
			return nil, fmt.Errorf("%s: 服务端要求补充输入，但本客户端未声明相应能力（elicitation/sampling/roots），无法满足：%s",
				method, strings.Join(keys, ", "))
		}
		if env.RequestState == "" {
			// 既没 inputRequests 也没 requestState：重试也没有新信息，报错防空转。
			return nil, fmt.Errorf("%s: 服务端返回 input_required 但既无 inputRequests 也无 requestState", method)
		}
		if attempt >= maxMRTRRetries {
			return nil, fmt.Errorf("%s: 服务端连续 %d 次要求补充输入，放弃", method, attempt)
		}
		requestState = env.RequestState // 下一轮原样回显
	}
}

// rpcOnce 发一次请求并等待同 id 响应。requestState 非空时作为 MRTR 重试回显。
func (c *Client) rpcOnce(ctx context.Context, method string, params map[string]any, requestState string) (json.RawMessage, error) {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil, errors.New("mcp client closed")
	}
	c.nextID++
	id := c.nextID
	key := fmt.Sprintf("%d", id)
	modern := c.era == eraModern
	version := c.version
	c.mu.Unlock()

	msg := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
	p := make(map[string]any, len(params)+2)
	for k, v := range params {
		p[k] = v
	}
	if modern {
		// _meta 放在 params 顶层。**绝不并进 arguments**：那是工具入参，混进去会污染
		// 用户数据。
		p["_meta"] = map[string]any{
			metaProtocolVersion:    version,
			metaClientCapabilities: map[string]any{},
			metaClientInfo:         map[string]any{"name": clientName, "version": clientVersion},
		}
	}
	if requestState != "" {
		p["requestState"] = requestState
	}
	if len(p) > 0 {
		msg["params"] = p
	}

	body, err := json.Marshal(msg)
	if err != nil {
		return nil, err
	}
	// Mcp-Name：tools/call 用 params.name（stdio 不用头，传了也无害。服务端不认识就忽略。
	meta := reqMeta{id: key, method: method, name: nameOf(params), version: version}
	m, err := c.tr.Call(ctx, body, meta)
	if err != nil {
		return nil, err
	}
	if m.Error != nil {
		return nil, &rpcCallError{method: method, code: m.Error.Code, msg: m.Error.Message, data: m.Error.Data}
	}
	return m.Result, nil
}

// nameOf 取 params 里的 name / uri——HTTP 的 Mcp-Name 头需要它。
func nameOf(params map[string]any) string {
	for _, k := range []string{"name", "uri"} {
		if v, ok := params[k]; ok {
			if s, ok := v.(string); ok {
				return s
			}
		}
	}
	return ""
}

// notify 发送无 id 通知（仅 legacy 握手的 notifications/initialized 用）。
func (c *Client) notify(method string) error {
	body, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "method": method})
	if err != nil {
		return err
	}
	return c.tr.Notify(body, reqMeta{method: method, version: c.Version()})
}