Maven/internal/mcp/http.go

package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"
)

// Poster is the HTTP seam: internal/webfetch.Fetcher satisfies it. The
// transport takes it as an interface so a test can serve a fake without a
// listener, and so that the ONLY implementation wired in production is the
// guarded fetcher — an MCP endpoint cannot get a bare http.Client this way.
type Poster interface {
	Post(ctx context.Context, rawURL, contentType string, body []byte, hdr map[string]string) (*PostResponse, error)
}

// PostResponse is the shape webfetch returns, restated here so this package
// does not depend on it structurally.
type PostResponse struct {
	Status      int
	ContentType string
	Body        []byte
	Header      map[string]string
}

// httpTransport speaks streamable HTTP: every request is a POST to one
// endpoint, and the reply is either a JSON object or a text/event-stream frame
// carrying one. Both are accepted — servers pick per response, and the two the
// LAN runs disagree about which.
type httpTransport struct {
	poster Poster
	url    string

	mu      sync.Mutex
	session string // Mcp-Session-Id, echoed back when the server issues one
}

func newHTTPTransport(post Poster, endpoint string) *httpTransport {
	return &httpTransport{poster: post, url: endpoint}
}

func (t *httpTransport) Call(ctx context.Context, req *rpcRequest) (*rpcResponse, error) {
	body, err := t.send(ctx, req)
	if err != nil {
		return nil, err
	}
	frame, err := decodeFrame(body)
	if err != nil {
		return nil, err
	}
	var resp rpcResponse
	if err := json.Unmarshal(frame, &resp); err != nil {
		return nil, fmt.Errorf("mcp: decode response: %w", err)
	}
	return &resp, nil
}

func (t *httpTransport) Notify(ctx context.Context, method string, params any) error {
	_, err := t.send(ctx, &rpcRequest{JSONRPC: "2.0", Method: method, Params: params})
	return err
}

func (t *httpTransport) send(ctx context.Context, req *rpcRequest) ([]byte, error) {
	req.JSONRPC = "2.0"
	raw, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}
	hdr := map[string]string{"Accept": "application/json, text/event-stream"}
	t.mu.Lock()
	if t.session != "" {
		hdr["Mcp-Session-Id"] = t.session
	}
	t.mu.Unlock()

	resp, err := t.poster.Post(ctx, t.url, "application/json", raw, hdr)
	if err != nil {
		return nil, err
	}
	if sid := headerGet(resp.Header, "Mcp-Session-Id"); sid != "" {
		t.mu.Lock()
		t.session = sid
		t.mu.Unlock()
	}
	return resp.Body, nil
}

func (t *httpTransport) Close() error {
	t.mu.Lock()
	t.session = ""
	t.mu.Unlock()
	return nil
}

func headerGet(h map[string]string, key string) string {
	if h == nil {
		return ""
	}
	if v, ok := h[key]; ok {
		return v
	}
	lower := strings.ToLower(key)
	for k, v := range h {
		if strings.ToLower(k) == lower {
			return v
		}
	}
	return ""
}

// decodeFrame pulls the JSON object out of a body that is either raw JSON or
// SSE. For SSE we take the LAST data: payload that parses, which is the
// response — earlier frames on the stream are progress notifications.
func decodeFrame(body []byte) ([]byte, error) {
	trimmed := bytes.TrimSpace(body)
	if len(trimmed) == 0 {
		return nil, errors.New("mcp: empty response body")
	}
	if trimmed[0] == '{' || trimmed[0] == '[' {
		return trimmed, nil
	}
	var last []byte
	sc := bufio.NewScanner(bytes.NewReader(trimmed))
	sc.Buffer(make([]byte, 0, 64<<10), maxLine)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "" {
			continue
		}
		var probe map[string]json.RawMessage
		if json.Unmarshal([]byte(payload), &probe) != nil {
			continue
		}
		if _, isResp := probe["id"]; isResp {
			last = []byte(payload)
		}
	}
	if err := sc.Err(); err != nil {
		return nil, fmt.Errorf("mcp: read event stream: %w", err)
	}
	if last == nil {
		return nil, errors.New("mcp: no JSON-RPC response in event stream")
	}
	return last, nil
}