Agent-Native/runtime/mcp/stdio.go

package mcp

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"sync"
	"sync/atomic"
)

// stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。
// 同一时刻只允许一个在途请求（串行锁），匹配简单且足以支撑智能体工具调用。
type stdioClient struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	stdout *bufio.Reader
	mu     sync.Mutex
	nextID int64
	closed bool
}

// DialStdio 启动 server 子进程并完成 initialize 握手。
func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	if cfg.Command == "" {
		return nil, fmt.Errorf("mcp server %q: stdio transport requires command", cfg.Name)
	}
	cmd := exec.Command(cfg.Command, cfg.Args...)
	cmd.Env = cfg.Env
	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, fmt.Errorf("mcp server %q: start %s: %w", cfg.Name, cfg.Command, err)
	}
	c := &stdioClient{cmd: cmd, stdin: stdin, stdout: bufio.NewReaderSize(stdout, 1<<20)}
	if err := c.handshake(ctx); err != nil {
		_ = c.Close()
		return nil, err
	}
	return c, nil
}

func (c *stdioClient) handshake(ctx context.Context) error {
	var res initializeResult
	if err := c.roundTrip(ctx, "initialize", initializeParams{
		ProtocolVersion: ProtocolVersion,
		ClientInfo:      clientInfo{Name: "agent-native", Version: "1"},
	}, &res); err != nil {
		return err
	}
	// initialized 通知无需响应。
	return c.notify("notifications/initialized")
}

func (c *stdioClient) Initialize(ctx context.Context) (initializeResult, error) {
	// Initialize 已在 Dial 时完成；这里补一次协议层调用以支持探测场景。
	var res initializeResult
	err := c.roundTrip(ctx, "initialize", initializeParams{
		ProtocolVersion: ProtocolVersion,
		ClientInfo:      clientInfo{Name: "agent-native", Version: "1"},
	}, &res)
	return res, err
}

func (c *stdioClient) ListTools(ctx context.Context) ([]toolInfo, error) {
	var out []toolInfo
	cursor := ""
	for {
		params := map[string]any{}
		if cursor != "" {
			params["cursor"] = cursor
		}
		var res toolsListResult
		if err := c.roundTrip(ctx, "tools/list", params, &res); err != nil {
			return nil, err
		}
		out = append(out, res.Tools...)
		if res.NextCursor == "" || len(out) > maxToolsPerServer*2 {
			return out, nil
		}
		cursor = res.NextCursor
	}
}

func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) {
	var res callToolResult
	if len(args) == 0 {
		args = json.RawMessage(`{}`)
	}
	err := c.roundTrip(ctx, "tools/call", callToolParams{Name: name, Arguments: args}, &res)
	return res, err
}

func (c *stdioClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return nil
	}
	c.closed = true
	_ = c.stdin.Close()
	if c.cmd.Process != nil {
		_ = c.cmd.Process.Kill()
	}
	return c.cmd.Wait()
}

func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return fmt.Errorf("mcp client closed")
	}
	id := atomic.AddInt64(&c.nextID, 1)
	req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}
	body, err := json.Marshal(req)
	if err != nil {
		return err
	}
	body = append(body, '\n')
	if _, err := c.stdin.Write(body); err != nil {
		return fmt.Errorf("mcp write %s: %w", method, err)
	}
	line, err := c.readResponseLine(ctx)
	if err != nil {
		return fmt.Errorf("mcp read %s: %w", method, err)
	}
	var resp rpcResponse
	if err := json.Unmarshal(line, &resp); err != nil {
		return fmt.Errorf("mcp %s: bad response frame: %w", method, err)
	}
	if resp.ID == nil || *resp.ID != id {
		return fmt.Errorf("mcp %s: response id mismatch", method)
	}
	if resp.Error != nil {
		return resp.Error
	}
	if out != nil && len(resp.Result) > 0 {
		return json.Unmarshal(resp.Result, out)
	}
	return nil
}

// readResponseLine 逐行读取，跳过通知帧（无 id），带 ctx 取消感知。
func (c *stdioClient) readResponseLine(ctx context.Context) ([]byte, error) {
	type result struct {
		line []byte
		err  error
	}
	done := make(chan result, 1)
	go func() {
		for {
			line, err := c.stdout.ReadBytes('\n')
			if err != nil {
				done <- result{nil, err}
				return
			}
			var probe rpcResponse
			if json.Unmarshal(line, &probe) == nil && probe.ID != nil {
				done <- result{line, nil}
				return
			}
			// 通知或日志行：忽略。
		}
	}()
	select {
	case r := <-done:
		return r.line, r.err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (c *stdioClient) notify(method string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return fmt.Errorf("mcp client closed")
	}
	body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method})
	if err != nil {
		return err
	}
	body = append(body, '\n')
	_, err = c.stdin.Write(body)
	return err
}