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 通信。
// 一个长驻读循环按 id 把响应分发给各请求；写 stdin 仍串行。
// 取消会摘掉等待项，迟到的回复直接丢弃。读循环退出时，在等的请求全部失败，客户端标记为已死。
type stdioClient struct {
	cmd        *exec.Cmd
	stdin      io.WriteCloser
	stdout     *bufio.Reader
	writeMu    sync.Mutex
	mu         sync.Mutex
	pending    map[int64]chan stdioResult
	nextID     int64
	readers    atomic.Int32
	closed     bool
	readerOnce sync.Once
}

type stdioResult struct {
	resp rpcResponse
	err  error
}

// 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),
		pending: map[int64]chan stdioResult{},
	}
	c.startReader()
	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()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.mu.Unlock()
	_ = 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.startReader()
	id := atomic.AddInt64(&c.nextID, 1)
	ch := make(chan stdioResult, 1)
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return fmt.Errorf("mcp client closed")
	}
	c.pending[id] = ch
	c.mu.Unlock()

	req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}
	body, err := json.Marshal(req)
	if err != nil {
		c.dropPending(id)
		return err
	}
	body = append(body, '\n')
	c.writeMu.Lock()
	_, err = c.stdin.Write(body)
	c.writeMu.Unlock()
	if err != nil {
		c.dropPending(id)
		return fmt.Errorf("mcp write %s: %w", method, err)
	}

	var got stdioResult
	select {
	case got = <-ch:
	case <-ctx.Done():
		c.dropPending(id)
		return fmt.Errorf("mcp read %s: %w", method, ctx.Err())
	}
	if got.err != nil {
		return fmt.Errorf("mcp read %s: %w", method, got.err)
	}
	resp := got.resp
	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
}

func (c *stdioClient) dropPending(id int64) {
	c.mu.Lock()
	delete(c.pending, id)
	c.mu.Unlock()
}

func (c *stdioClient) startReader() {
	c.readerOnce.Do(func() {
		go c.readLoop()
	})
}

func (c *stdioClient) readLoop() {
	c.readers.Add(1)
	defer c.readers.Add(-1)
	var fatal error
	for {
		line, err := c.stdout.ReadBytes('\n')
		if err != nil {
			fatal = err
			break
		}
		var resp rpcResponse
		if json.Unmarshal(line, &resp) != nil || resp.ID == nil {
			continue
		}
		c.mu.Lock()
		ch := c.pending[*resp.ID]
		if ch != nil {
			delete(c.pending, *resp.ID)
		}
		c.mu.Unlock()
		if ch == nil {
			continue
		}
		ch <- stdioResult{resp: resp}
	}
	c.mu.Lock()
	c.closed = true
	pending := c.pending
	c.pending = map[int64]chan stdioResult{}
	c.mu.Unlock()
	for _, ch := range pending {
		ch <- stdioResult{err: fatal}
	}
}

func (c *stdioClient) Readers() int { return int(c.readers.Load()) }

func (c *stdioClient) notify(method string) error {
	c.mu.Lock()
	closed := c.closed
	c.mu.Unlock()
	if 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')
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	_, err = c.stdin.Write(body)
	return err
}