From 765b6f5603b494a5e97aa5f781672f4ef628964a Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Tue, 20 Oct 2026 10:43:35 +0200 Subject: [PATCH] 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) --- internal/mcpclient/mcpclient.go | 346 ++++++++++++++++++++++++++ internal/mcpclient/mcpclient_test.go | 87 +++++++ internal/mcpclient/mcptest/mcptest.go | 158 ++++++++++++ 3 files changed, 591 insertions(+) create mode 100644 internal/mcpclient/mcpclient.go create mode 100644 internal/mcpclient/mcpclient_test.go create mode 100644 internal/mcpclient/mcptest/mcptest.go diff --git a/internal/mcpclient/mcpclient.go b/internal/mcpclient/mcpclient.go new file mode 100644 index 0000000..850670a --- /dev/null +++ b/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 +} diff --git a/internal/mcpclient/mcpclient_test.go b/internal/mcpclient/mcpclient_test.go new file mode 100644 index 0000000..2c215fb --- /dev/null +++ b/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) + } +} diff --git a/internal/mcpclient/mcptest/mcptest.go b/internal/mcpclient/mcptest/mcptest.go new file mode 100644 index 0000000..c90331f --- /dev/null +++ b/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) +}