Agent-Native/runtime/mcp/http.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// httpClient 走 MCP Streamable HTTP 传输：每个 JSON-RPC 请求一次 POST，
// 期待 application/json 的单帧响应（SSE 流式与 server 主动请求不在本客户端范围）。
// initialize 返回的 Mcp-Session-Id 会在后续请求回带。
type httpClient struct {
	cfg       ServerConfig
	client    *http.Client
	sessionID string
	mu        sync.Mutex
	nextID    int64
}

// DialHTTP 建立 HTTP 客户端并完成 initialize 握手。
func DialHTTP(ctx context.Context, cfg ServerConfig) (*httpClient, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}
	if cfg.URL == "" {
		return nil, fmt.Errorf("mcp server %q: http transport requires url", cfg.Name)
	}
	c := &httpClient{cfg: cfg, client: &http.Client{Timeout: requestTimeout(cfg)}}
	if err := c.handshake(ctx); err != nil {
		return nil, err
	}
	return c, nil
}

func requestTimeout(cfg ServerConfig) time.Duration {
	if cfg.Timeout > 0 {
		return cfg.Timeout
	}
	return 30 * time.Second
}

func (c *httpClient) handshake(ctx context.Context) error {
	var res initializeResult
	err := c.roundTrip(ctx, "initialize", initializeParams{
		ProtocolVersion: ProtocolVersion,
		ClientInfo:      clientInfo{Name: "agent-native", Version: "1"},
	}, &res)
	return err
}

func (c *httpClient) Initialize(ctx context.Context) (initializeResult, error) {
	var res initializeResult
	err := c.roundTrip(ctx, "initialize", initializeParams{
		ProtocolVersion: ProtocolVersion,
		ClientInfo:      clientInfo{Name: "agent-native", Version: "1"},
	}, &res)
	return res, err
}

func (c *httpClient) 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 *httpClient) 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 *httpClient) Close() error {
	c.mu.Lock()
	sid := c.sessionID
	url := c.cfg.URL
	client := c.client
	c.mu.Unlock()
	if sid == "" || url == "" || client == nil {
		return nil
	}
	req, err := http.NewRequest(http.MethodDelete, url, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Mcp-Session-Id", sid)
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 2048))
	if resp.StatusCode >= 400 {
		return fmt.Errorf("mcp session terminate: http %d", resp.StatusCode)
	}
	return nil
}

func (c *httpClient) roundTrip(ctx context.Context, method string, params any, out any) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	id := atomic.AddInt64(&c.nextID, 1)
	body, err := json.Marshal(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params})
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.URL, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	for k, v := range c.cfg.Headers {
		req.Header.Set(k, v)
	}
	if c.sessionID != "" {
		req.Header.Set("Mcp-Session-Id", c.sessionID)
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return fmt.Errorf("mcp %s: %w", method, err)
	}
	defer resp.Body.Close()
	if sid := resp.Header.Get("Mcp-Session-Id"); sid != "" {
		c.sessionID = sid
	}
	if resp.StatusCode >= 400 {
		snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return fmt.Errorf("mcp %s: http %d: %s", method, resp.StatusCode, string(snippet))
	}
	// Streamable HTTP 允许 SSE 回包：逐行扫 data: 帧，取第一个 JSON 对象。
	var payload []byte
	ct := resp.Header.Get("Content-Type")
	if bytes.Contains([]byte(ct), []byte("text/event-stream")) {
		payload, err = readSSEFrame(resp.Body)
	} else {
		payload, err = io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	}
	if err != nil {
		return fmt.Errorf("mcp %s: read body: %w", method, err)
	}
	var rpcResp rpcResponse
	if err := json.Unmarshal(payload, &rpcResp); err != nil {
		return fmt.Errorf("mcp %s: bad response frame: %w", method, err)
	}
	if rpcResp.ID == nil || *rpcResp.ID != id {
		return fmt.Errorf("mcp %s: response id mismatch", method)
	}
	if rpcResp.Error != nil {
		return rpcResp.Error
	}
	if out != nil && len(rpcResp.Result) > 0 {
		return json.Unmarshal(rpcResp.Result, out)
	}
	return nil
}

// readSSEFrame 从 SSE 流读取首个 data: JSON 载荷（忽略事件注释行）。
func readSSEFrame(r io.Reader) ([]byte, error) {
	buf := make([]byte, 0, 4096)
	tmp := make([]byte, 2048)
	for {
		n, err := r.Read(tmp)
		buf = append(buf, tmp[:n]...)
		for _, line := range splitLines(buf) {
			if bytes.HasPrefix(line, []byte("data:")) {
				data := bytes.TrimSpace(line[len("data:"):])
				if len(data) > 0 {
					return data, nil
				}
			}
		}
		if err != nil {
			if len(buf) == 0 || err == io.EOF {
				return nil, fmt.Errorf("no data frame in sse response")
			}
			return nil, err
		}
		if len(buf) > 8<<20 {
			return nil, fmt.Errorf("sse response too large")
		}
	}
}

func splitLines(b []byte) [][]byte {
	var out [][]byte
	start := 0
	for i, c := range b {
		if c == '\n' {
			out = append(out, b[start:i])
			start = i + 1
		}
	}
	if start < len(b) {
		out = append(out, b[start:])
	}
	return out
}