llmgw/internal/mcpclient/mcpclient.go

// Package mcpclient speaks to an MCP server over the streamable HTTP
// transport: the company's own tool servers, which llmgw's /mcp serves on to
// its callers. It initializes once and keeps the session the server gives,
// lists the server's tools, and calls them, taking an answer as JSON or as a
// stream of events, whichever the server sends.
package mcpclient

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
)

// ProtocolVersion is the version asked for; a server may answer an older one.
const ProtocolVersion = "2025-06-18"

// Client is one MCP server.
type Client struct {
	URL string
	// Header gives what each request carries beyond the protocol's own: the
	// server's credentials. Nil sends none.
	Header func(ctx context.Context) (http.Header, error)
	HTTP   *http.Client
	// MaxBytes bounds one answer; zero is 4 MiB.
	MaxBytes int64

	mu      sync.Mutex
	session string
	version string
	ready   bool
	nextID  atomic.Int64
}

// Tool is one tool as the server describes it.
type Tool struct {
	Name        string          `json:"name"`
	Title       string          `json:"title,omitempty"`
	Description string          `json:"description,omitempty"`
	InputSchema json.RawMessage `json:"inputSchema,omitempty"`
	Annotations json.RawMessage `json:"annotations,omitempty"`
}

// Error is the server's own JSON-RPC error.
type Error struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func (e *Error) Error() string {
	return fmt.Sprintf("the MCP server answered %d: %s", e.Code, e.Message)
}

// errSessionGone is the server no longer knowing the session: it is
// initialized again, once.
var errSessionGone = errors.New("the MCP server no longer knows the session")

type rpcMessage struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method,omitempty"`
	Params  json.RawMessage `json:"params,omitempty"`
	Result  json.RawMessage `json:"result,omitempty"`
	Error   *Error          `json:"error,omitempty"`
}

// Tools lists every tool, page by page.
func (c *Client) Tools(ctx context.Context) ([]Tool, error) {
	var out []Tool
	cursor := ""
	for range 100 {
		params := map[string]any{}
		if cursor != "" {
			params["cursor"] = cursor
		}
		raw, err := c.request(ctx, "tools/list", params, nil, nil)
		if err != nil {
			return nil, err
		}
		var page struct {
			Tools      []Tool `json:"tools"`
			NextCursor string `json:"nextCursor"`
		}
		if err := json.Unmarshal(raw, &page); err != nil {
			return nil, fmt.Errorf("the MCP server's tool list did not parse: %w", err)
		}
		out = append(out, page.Tools...)
		if page.NextCursor == "" {
			return out, nil
		}
		cursor = page.NextCursor
	}
	return out, errors.New("the MCP server's tool list never ended")
}

// Call calls a tool and gives its result as the server sent it. extra is
// sent with this request alone (who is calling, on whose behalf); progress
// hears the server's progress notifications while it works, when it sends
// them as a stream.
func (c *Client) Call(ctx context.Context, name string, args json.RawMessage, extra http.Header,
	progress func(params json.RawMessage),
) (json.RawMessage, error) {
	if len(args) == 0 {
		args = json.RawMessage("{}")
	}
	params := map[string]any{"name": name, "arguments": args}
	if progress != nil {
		params["_meta"] = map[string]any{"progressToken": c.nextID.Add(1)}
	}
	return c.request(ctx, "tools/call", params, extra, progress)
}

// request sends one request, initializing first when there is no session,
// and again once when the server has dropped it.
func (c *Client) request(ctx context.Context, method string, params any, extra http.Header,
	progress func(json.RawMessage),
) (json.RawMessage, error) {
	for attempt := 0; ; attempt++ {
		if err := c.ensureReady(ctx); err != nil {
			return nil, err
		}
		raw, err := c.send(ctx, method, params, extra, progress)
		if errors.Is(err, errSessionGone) && attempt == 0 {
			c.mu.Lock()
			c.ready, c.session = false, ""
			c.mu.Unlock()
			continue
		}
		return raw, err
	}
}

func (c *Client) ensureReady(ctx context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.ready {
		return nil
	}
	c.session, c.version = "", ""
	raw, err := c.sendLocked(ctx, "initialize", map[string]any{
		"protocolVersion": ProtocolVersion,
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "llmgw", "version": "1"},
	}, nil, nil, true)
	if err != nil {
		return fmt.Errorf("initializing: %w", err)
	}
	var init struct {
		ProtocolVersion string `json:"protocolVersion"`
	}
	_ = json.Unmarshal(raw, &init)
	c.version = init.ProtocolVersion
	if err := c.notifyLocked(ctx, "notifications/initialized"); err != nil {
		return fmt.Errorf("initializing: %w", err)
	}
	c.ready = true
	return nil
}

func (c *Client) send(ctx context.Context, method string, params any, extra http.Header,
	progress func(json.RawMessage),
) (json.RawMessage, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.sendLocked(ctx, method, params, extra, progress, false)
}

// sendLocked posts one request and reads its answer. Held under the lock
// only for what the session needs: the request itself runs unlocked for a
// tool call, so slow calls do not queue behind each other.
func (c *Client) sendLocked(ctx context.Context, method string, params any, extra http.Header,
	progress func(json.RawMessage), initializing bool,
) (json.RawMessage, error) {
	id := c.nextID.Add(1)
	p, err := json.Marshal(params)
	if err != nil {
		return nil, err
	}
	body, _ := json.Marshal(rpcMessage{JSONRPC: "2.0", ID: json.RawMessage(fmt.Sprint(id)), Method: method, Params: p})
	req, err := c.newRequest(ctx, body, extra)
	if err != nil {
		return nil, err
	}
	if !initializing {
		// The request may be long (a tool that searches); the session is
		// read, the lock let go for the wait, and taken again after.
		c.mu.Unlock()
		defer c.mu.Lock()
	}
	resp, err := c.client().Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if s := resp.Header.Get("Mcp-Session-Id"); s != "" && initializing {
		c.session = s
	}
	switch {
	case resp.StatusCode == http.StatusNotFound && req.Header.Get("Mcp-Session-Id") != "":
		return nil, errSessionGone
	case resp.StatusCode >= 300:
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<10))
		return nil, fmt.Errorf("the MCP server answered HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
	}
	limited := io.LimitReader(resp.Body, c.maxBytes()+1)
	var msg rpcMessage
	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
		msg, err = readStream(limited, id, progress, c.maxBytes())
	} else {
		var b []byte
		if b, err = io.ReadAll(limited); err == nil {
			if int64(len(b)) > c.maxBytes() {
				err = fmt.Errorf("the MCP server's answer is larger than %d bytes", c.maxBytes())
			} else {
				err = json.Unmarshal(b, &msg)
			}
		}
	}
	if err != nil {
		return nil, err
	}
	if msg.Error != nil {
		return nil, msg.Error
	}
	return msg.Result, nil
}

// notifyLocked sends a notification, which has no answer.
func (c *Client) notifyLocked(ctx context.Context, method string) error {
	body, _ := json.Marshal(rpcMessage{JSONRPC: "2.0", Method: method})
	req, err := c.newRequest(ctx, body, nil)
	if err != nil {
		return err
	}
	resp, err := c.client().Do(req)
	if err != nil {
		return err
	}
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<16))
	resp.Body.Close()
	if resp.StatusCode >= 300 {
		return fmt.Errorf("the MCP server answered HTTP %d to %s", resp.StatusCode, method)
	}
	return nil
}

func (c *Client) newRequest(ctx context.Context, body []byte, extra http.Header) (*http.Request, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.URL, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	if c.Header != nil {
		h, err := c.Header(ctx)
		if err != nil {
			return nil, err
		}
		for k, vs := range h {
			for _, v := range vs {
				req.Header.Add(k, v)
			}
		}
	}
	for k, vs := range extra {
		for _, v := range vs {
			req.Header.Add(k, v)
		}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	if c.session != "" {
		req.Header.Set("Mcp-Session-Id", c.session)
	}
	if c.version != "" {
		req.Header.Set("MCP-Protocol-Version", c.version)
	}
	return req, nil
}

// readStream reads server-sent events until the answer to id, handing
// progress notifications on.
func readStream(r io.Reader, id int64, progress func(json.RawMessage), limit int64) (rpcMessage, error) {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64<<10), int(limit))
	want := fmt.Sprint(id)
	var data strings.Builder
	flush := func() (rpcMessage, bool) {
		defer data.Reset()
		if data.Len() == 0 {
			return rpcMessage{}, false
		}
		var m rpcMessage
		if json.Unmarshal([]byte(data.String()), &m) != nil {
			return rpcMessage{}, false
		}
		if m.Method == "notifications/progress" && progress != nil {
			progress(m.Params)
			return rpcMessage{}, false
		}
		if string(m.ID) == want && m.Method == "" {
			return m, true
		}
		return rpcMessage{}, false
	}
	for sc.Scan() {
		line := sc.Text()
		switch {
		case line == "":
			if m, ok := flush(); ok {
				return m, nil
			}
		case strings.HasPrefix(line, "data:"):
			if data.Len() > 0 {
				data.WriteByte('\n')
			}
			data.WriteString(strings.TrimPrefix(strings.TrimPrefix(line, "data:"), " "))
		}
	}
	if m, ok := flush(); ok {
		return m, nil
	}
	if err := sc.Err(); err != nil {
		return rpcMessage{}, fmt.Errorf("reading the MCP server's stream: %w", err)
	}
	return rpcMessage{}, errors.New("the MCP server's stream ended without an answer")
}

func (c *Client) client() *http.Client {
	if c.HTTP != nil {
		return c.HTTP
	}
	return http.DefaultClient
}

func (c *Client) maxBytes() int64 {
	if c.MaxBytes > 0 {
		return c.MaxBytes
	}
	return 4 << 20
}