An MCP client for the company's tool servers

internal/mcpclient speaks streamable HTTP: initializes once and keeps
the session the server gives, beginning it again once if the server
forgot it; lists tools page by page; calls one with headers of its own,
taking the answer as JSON or as events with progress before it; and
bounds what it reads. mcptest is a server for tests.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

internal/mcpclient/mcpclient.go

@@ -0,0 +1,346 @@
+// 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
+}

internal/mcpclient/mcpclient_test.go

@@ -0,0 +1,87 @@
+package mcpclient
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"strings"
+	"testing"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/mcpclient/mcptest"
+)
+
+func TestToolsAreListedPageByPageAndCalledInOneSession(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	srv.PageSize = 2
+	c := &Client{URL: srv.URL, Header: func(context.Context) (http.Header, error) {
+		return http.Header{"Authorization": {"Bearer service"}}, nil
+	}}
+	tools, err := c.Tools(t.Context())
+	if err != nil || len(tools) != 3 || tools[2].Name != "big" {
+		t.Fatalf("tools %+v %v", tools, err)
+	}
+	raw, err := c.Call(t.Context(), "echo", json.RawMessage(`{"text":"hello"}`), http.Header{"X-Llmgw-User": {"ada"}}, nil)
+	if err != nil || !strings.Contains(string(raw), "hello") {
+		t.Fatalf("call %s %v", raw, err)
+	}
+	if srv.Sessions() != 1 {
+		t.Errorf("initialized %d times", srv.Sessions())
+	}
+	h := srv.Headers[0]
+	if h.Get("Authorization") != "Bearer service" || h.Get("X-Llmgw-User") != "ada" || h.Get("MCP-Protocol-Version") != "2025-06-18" {
+		t.Errorf("headers %v", h)
+	}
+}
+
+func TestAStreamedAnswerIsReadAndItsProgressHeard(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	srv.Stream = true
+	c := &Client{URL: srv.URL}
+	var heard []string
+	raw, err := c.Call(t.Context(), "echo", json.RawMessage(`{"text":"streamed"}`), nil, func(p json.RawMessage) {
+		heard = append(heard, string(p))
+	})
+	if err != nil || !strings.Contains(string(raw), "streamed") {
+		t.Fatalf("call %s %v", raw, err)
+	}
+	if len(heard) != 1 || !strings.Contains(heard[0], "halfway") {
+		t.Errorf("progress %v", heard)
+	}
+}
+
+func TestASessionTheServerForgotIsBegunAgainOnce(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	c := &Client{URL: srv.URL}
+	if _, err := c.Tools(t.Context()); err != nil {
+		t.Fatal(err)
+	}
+	srv.Forget()
+	if _, err := c.Call(t.Context(), "echo", json.RawMessage(`{"text":"again"}`), nil, nil); err != nil {
+		t.Fatalf("after the server restarted: %v", err)
+	}
+	if srv.Sessions() != 2 {
+		t.Errorf("sessions %d", srv.Sessions())
+	}
+}
+
+func TestTheServersErrorsAndAnAnswerTooLargeAreSaid(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	c := &Client{URL: srv.URL, MaxBytes: 64 << 10}
+	var rpcErr *Error
+	if _, err := c.Call(t.Context(), "nosuch", nil, nil, nil); !errors.As(err, &rpcErr) || rpcErr.Code != -32602 {
+		t.Errorf("unknown tool: %v", err)
+	}
+	if _, err := c.Call(t.Context(), "big", nil, nil, nil); err == nil || !strings.Contains(err.Error(), "larger than") {
+		t.Errorf("a megabyte through a 64 KiB limit: %v", err)
+	}
+	// A tool's own failure is its result, not an error of the protocol.
+	raw, err := c.Call(t.Context(), "fail", nil, nil, nil)
+	if err != nil || !strings.Contains(string(raw), `"isError":true`) {
+		t.Errorf("fail: %s %v", raw, err)
+	}
+}

internal/mcpclient/mcptest/mcptest.go

@@ -0,0 +1,158 @@
+// Package mcptest is an MCP server for tests: a few tools, a session it
+// gives at initialize and can forget, its tool list in pages, and answers
+// as JSON or, when asked, as a stream with progress before the answer.
+package mcptest
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+)
+
+// Server is the fake server and what it saw.
+type Server struct {
+	*httptest.Server
+
+	// Stream answers tools/call as server-sent events, with a progress
+	// notification first.
+	Stream bool
+	// PageSize splits the tool list into pages; zero sends it whole.
+	PageSize int
+
+	mu       sync.Mutex
+	session  string
+	sessions int
+	// Calls are the tools called, and Headers the headers of each call.
+	Calls   []string
+	Args    []json.RawMessage
+	Headers []http.Header
+}
+
+// Tools are what it offers: echo returns its text, fail answers with a
+// tool error, big returns more than most limits allow.
+var Tools = []map[string]any{
+	{"name": "echo", "description": "Says the text back.", "inputSchema": map[string]any{
+		"type": "object", "properties": map[string]any{"text": map[string]any{"type": "string"}}}},
+	{"name": "fail", "description": "Always fails.", "inputSchema": map[string]any{"type": "object"}},
+	{"name": "big", "description": "Answers with a megabyte.", "inputSchema": map[string]any{"type": "object"}},
+}
+
+// New starts one.
+func New() *Server {
+	s := &Server{}
+	s.Server = httptest.NewServer(http.HandlerFunc(s.serve))
+	return s
+}
+
+// Forget drops the session, as a restarted server would.
+func (s *Server) Forget() {
+	s.mu.Lock()
+	s.session = ""
+	s.mu.Unlock()
+}
+
+// Sessions is how many times it was initialized.
+func (s *Server) Sessions() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.sessions
+}
+
+func (s *Server) serve(w http.ResponseWriter, r *http.Request) {
+	var msg struct {
+		ID     json.RawMessage `json:"id"`
+		Method string          `json:"method"`
+		Params json.RawMessage `json:"params"`
+	}
+	if json.NewDecoder(r.Body).Decode(&msg) != nil {
+		http.Error(w, "bad", http.StatusBadRequest)
+		return
+	}
+	s.mu.Lock()
+	if msg.Method == "initialize" {
+		s.sessions++
+		s.session = fmt.Sprintf("s%d", s.sessions)
+		w.Header().Set("Mcp-Session-Id", s.session)
+	} else if r.Header.Get("Mcp-Session-Id") != s.session || s.session == "" {
+		s.mu.Unlock()
+		http.Error(w, "no such session", http.StatusNotFound)
+		return
+	}
+	s.mu.Unlock()
+	if len(msg.ID) == 0 {
+		w.WriteHeader(http.StatusAccepted)
+		return
+	}
+	answer := func(result any) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": msg.ID, "result": result})
+	}
+	switch msg.Method {
+	case "initialize":
+		answer(map[string]any{"protocolVersion": "2025-06-18", "capabilities": map[string]any{"tools": map[string]any{}},
+			"serverInfo": map[string]any{"name": "test", "version": "1"}})
+	case "tools/list":
+		var p struct {
+			Cursor string `json:"cursor"`
+		}
+		_ = json.Unmarshal(msg.Params, &p)
+		start := 0
+		_, _ = fmt.Sscan(p.Cursor, &start)
+		end := len(Tools)
+		if s.PageSize > 0 && start+s.PageSize < end {
+			end = start + s.PageSize
+		}
+		res := map[string]any{"tools": Tools[start:end]}
+		if end < len(Tools) {
+			res["nextCursor"] = fmt.Sprint(end)
+		}
+		answer(res)
+	case "tools/call":
+		var p struct {
+			Name      string          `json:"name"`
+			Arguments json.RawMessage `json:"arguments"`
+		}
+		_ = json.Unmarshal(msg.Params, &p)
+		s.mu.Lock()
+		s.Calls, s.Args, s.Headers = append(s.Calls, p.Name), append(s.Args, p.Arguments), append(s.Headers, r.Header.Clone())
+		s.mu.Unlock()
+		var result map[string]any
+		switch p.Name {
+		case "echo":
+			var a struct {
+				Text string `json:"text"`
+			}
+			_ = json.Unmarshal(p.Arguments, &a)
+			result = map[string]any{"content": []any{map[string]any{"type": "text", "text": a.Text}}}
+		case "fail":
+			result = map[string]any{"content": []any{map[string]any{"type": "text", "text": "it failed"}}, "isError": true}
+		case "big":
+			result = map[string]any{"content": []any{map[string]any{"type": "text", "text": strings.Repeat("x", 1<<20)}}}
+		default:
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": msg.ID,
+				"error": map[string]any{"code": -32602, "message": "unknown tool " + p.Name}})
+			return
+		}
+		if s.Stream {
+			w.Header().Set("Content-Type", "text/event-stream")
+			fmt.Fprintf(w, "data: %s\n\n", mustJSON(map[string]any{"jsonrpc": "2.0", "method": "notifications/progress",
+				"params": map[string]any{"progress": 1, "total": 2, "message": "halfway"}}))
+			fmt.Fprintf(w, "event: message\ndata: %s\n\n", mustJSON(map[string]any{"jsonrpc": "2.0", "id": msg.ID, "result": result}))
+			return
+		}
+		answer(result)
+	default:
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": msg.ID,
+			"error": map[string]any{"code": -32601, "message": "method not found"}})
+	}
+}
+
+func mustJSON(v any) string {
+	b, _ := json.Marshal(v)
+	return string(b)
+}