AutoZero/mcp/transport.go

package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"strings"
	"sync"

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

// reqMeta 是发送一条请求时传输层需要的元数据：HTTP 要据此生成 Mcp-Method / Mcp-Name 头，
// stdio 只用到 id（按 id 分发响应）。
type reqMeta struct {
	id      string // 归一化后的请求 id
	method  string // JSON-RPC method
	name    string // tools/call 的 params.name / resources/read 的 params.uri
	version string // 本次请求声明的协议版本（HTTP 放进 MCP-Protocol-Version 头）
}

// transport 是"发一条 JSON-RPC 消息并拿到响应"的抽象。
//
// stdio：一条连接复用多次请求，按 id 分发；HTTP：每个请求一次 POST，响应自带。
type transport interface {
	Call(ctx context.Context, body []byte, meta reqMeta) (rpcMsg, error)
	Notify(body []byte, meta reqMeta) error
	Close() error
	Kind() string // "stdio" | "http"，仅用于日志
}

// legacyServerError 表示"服务端不是现代的（无 initialize 之外的现代特征）"，
// 客户端据此回退到 handshake 式协议。HTTP 侧要靠它区分 400 的两种含义。
type legacyServerError struct{ reason string }

func (e *legacyServerError) Error() string { return e.reason }

// authRequiredError 表示服务端要求鉴权。它**不是年代信号**——两种年代都会返回 401，
// 所以年代探测遇到它应当直接终止，而不是再去试一次 legacy 握手（那样错误里会多出一个
// 误导性的 "initialize:" 前缀，让人以为是握手失败）。
type authRequiredError struct{ msg string }

func (e *authRequiredError) Error() string { return e.msg }

// ---- stdio ----

// stdioTransport 用子进程的 stdin/stdout 通信（JSON-RPC，换行分隔帧）。
type stdioTransport struct {
	cmd     *exec.Cmd
	stdin   io.WriteCloser
	pending map[string]chan rpcMsg
	mu      sync.Mutex
}

func newStdioTransport(cmd *exec.Cmd, stdin io.WriteCloser, stdout io.Reader) *stdioTransport {
	t := &stdioTransport{cmd: cmd, stdin: stdin, pending: map[string]chan rpcMsg{}}
	go t.readLoop(stdout)
	return t
}

func (t *stdioTransport) Kind() string { return "stdio" }

func (t *stdioTransport) Call(ctx context.Context, body []byte, meta reqMeta) (rpcMsg, error) {
	ch := make(chan rpcMsg, 1)
	t.mu.Lock()
	t.pending[meta.id] = ch
	t.mu.Unlock()

	if _, err := t.stdin.Write(append(body, '\n')); err != nil {
		t.drop(meta.id)
		return rpcMsg{}, err
	}
	select {
	case m := <-ch:
		return m, nil
	case <-ctx.Done():
		t.drop(meta.id)
		return rpcMsg{}, ctx.Err()
	}
}

func (t *stdioTransport) Notify(body []byte, _ reqMeta) error {
	_, err := t.stdin.Write(append(body, '\n'))
	return err
}

func (t *stdioTransport) drop(id string) {
	t.mu.Lock()
	delete(t.pending, id)
	t.mu.Unlock()
}

func (t *stdioTransport) Close() error {
	t.mu.Lock()
	for id, ch := range t.pending {
		ch <- rpcMsg{Error: &rpcError{Message: "client closed"}}
		delete(t.pending, id)
	}
	t.mu.Unlock()
	_ = t.stdin.Close()
	if t.cmd != nil {
		// 只 Kill 直接子进程会把真正的 server 留成孤儿（npx/uvx 起的 node/python
		// 就是孙子）：它继续占着 stdout，readLoop goroutine 永不退出。
		_ = execx.KillTree(t.cmd)
		return t.cmd.Wait() // 构造时已设 WaitDelay，这里不会无限阻塞
	}
	return nil
}

// readLoop 按行读 server stdout 并分发给等待者。
//
// **不能用 json.Decoder 直接解流**：stdout 上出现一行非 JSON（很多 server 会打启动
// banner）就会让它报错并停止读取，之后所有请求都等不到回复；而规范要求的恰恰是
// "无法识别的输出 ⇒ 判为 legacy ⇒ 回退 initialize"，按那种写法那次回退必然白等到超时。
// 所以这里跳过解析不了的行，只在 EOF/读错误时唤醒等待者。
func (t *stdioTransport) readLoop(r io.Reader) {
	br := bufio.NewReader(r)
	for {
		line, err := br.ReadBytes('\n')
		if trimmed := bytes.TrimSpace(line); len(trimmed) > 0 {
			var m rpcMsg
			if json.Unmarshal(trimmed, &m) == nil {
				t.dispatch(m)
			} else {
				diag.Printf("[mcp] 忽略无法解析的 stdout 行: %.120s\n", trimmed)
			}
		}
		if err != nil {
			t.failAll(fmt.Errorf("mcp server stopped: %w", err))
			return
		}
	}
}

// dispatch 把带 id 的响应交给等待者。通知与服务端发起的请求都不处理
// （现代协议里服务端不再发起请求；legacy 下我们声明了空 capabilities）。
func (t *stdioTransport) dispatch(m rpcMsg) {
	if len(m.ID) == 0 {
		return
	}
	key := idKey(m.ID)
	t.mu.Lock()
	ch := t.pending[key]
	delete(t.pending, key)
	t.mu.Unlock()
	if ch != nil {
		ch <- m
	}
}

func (t *stdioTransport) failAll(err error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for id, ch := range t.pending {
		ch <- rpcMsg{Error: &rpcError{Message: err.Error()}}
		delete(t.pending, id)
	}
}

// idKey 把 id 归一化成 map 键：去掉引号与空白。有些实现会把数字 id 回成字符串。
func idKey(raw json.RawMessage) string {
	return strings.Trim(strings.TrimSpace(string(raw)), `"`)
}