From e228d78d19de5e14adde52f85b65458a56eadf69 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Wed, 18 Jun 2025 20:39:04 +0000 Subject: [PATCH] internal/mcp: provide a customizable http.Client to client transports Allow customizing the http.Client used for HTTP MCP client transports, by adding client options structs. Change-Id: I2297acb136f8d0f7fa70d58cd244a6a81cc89751 Reviewed-on: https://go-review.googlesource.com/c/tools/+/682756 Reviewed-by: Jonathan Amsterdam Auto-Submit: Robert Findley LUCI-TryBot-Result: Go LUCI --- design/design.md | 21 +++++++++++++++++++-- mcp/sse.go | 23 ++++++++++++++++++++--- mcp/sse_example_test.go | 2 +- mcp/sse_test.go | 2 +- mcp/streamable.go | 25 +++++++++++++++++++++---- mcp/streamable_test.go | 30 +++++++++++++++++++++++++++--- 6 files changed, 89 insertions(+), 14 deletions(-) diff --git a/design/design.md b/design/design.md index af3ca6f..4b39e10 100644 --- a/design/design.md +++ b/design/design.md @@ -187,9 +187,17 @@ The SSE client transport is simpler, and hopefully self-explanatory. ```go type SSEClientTransport struct { /* ... */ } +// SSEClientTransportOptions provides options for the [NewSSEClientTransport] +// constructor. +type SSEClientTransportOptions struct { + // HTTPClient is the client to use for making HTTP requests. If nil, + // http.DefaultClient is used. + HTTPClient *http.Client +} + // NewSSEClientTransport returns a new client transport that connects to the // SSE server at the provided URL. -func NewSSEClientTransport(url string) (*SSEClientTransport, error) { +func NewSSEClientTransport(url string, opts *SSEClientTransportOptions) (*SSEClientTransport, error) // Connect connects through the client endpoint. func (*SSEClientTransport) Connect(ctx context.Context) (Connection, error) @@ -217,7 +225,16 @@ func (*StreamableServerTransport) Connect(context.Context) (Connection, error) // The streamable client handles reconnection transparently to the user. type StreamableClientTransport struct { /* ... */ } -func NewStreamableClientTransport(url string) *StreamableClientTransport { + +// StreamableClientTransportOptions provides options for the +// [NewStreamableClientTransport] constructor. +type StreamableClientTransportOptions struct { + // HTTPClient is the client to use for making HTTP requests. If nil, + // http.DefaultClient is used. + HTTPClient *http.Client +} + +func NewStreamableClientTransport(url string, opts *StreamableClientTransportOptions) *StreamableClientTransport func (*StreamableClientTransport) Connect(context.Context) (Connection, error) ``` diff --git a/mcp/sse.go b/mcp/sse.go index 2ee23d0..bdbb6d0 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -322,20 +322,33 @@ func (s sseServerConn) Close() error { // https://modelcontextprotocol.io/specification/2024-11-05/basic/transports type SSEClientTransport struct { sseEndpoint *url.URL + opts SSEClientTransportOptions +} + +// SSEClientTransportOptions provides options for the [NewSSEClientTransport] +// constructor. +type SSEClientTransportOptions struct { + // HTTPClient is the client to use for making HTTP requests. If nil, + // http.DefaultClient is used. + HTTPClient *http.Client } // NewSSEClientTransport returns a new client transport that connects to the // SSE server at the provided URL. // // NewSSEClientTransport panics if the given URL is invalid. -func NewSSEClientTransport(baseURL string) *SSEClientTransport { +func NewSSEClientTransport(baseURL string, opts *SSEClientTransportOptions) *SSEClientTransport { url, err := url.Parse(baseURL) if err != nil { panic(fmt.Sprintf("invalid base url: %v", err)) } - return &SSEClientTransport{ + t := &SSEClientTransport{ sseEndpoint: url, } + if opts != nil { + t.opts = *opts + } + return t } // Connect connects through the client endpoint. @@ -344,8 +357,12 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) { if err != nil { return nil, err } + httpClient := c.opts.HTTPClient + if httpClient == nil { + httpClient = http.DefaultClient + } req.Header.Set("Accept", "text/event-stream") - resp, err := http.DefaultClient.Do(req) + resp, err := httpClient.Do(req) if err != nil { return nil, err } diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 4602acd..e7a8eb5 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -33,7 +33,7 @@ func ExampleSSEHandler() { defer httpServer.Close() ctx := context.Background() - transport := mcp.NewSSEClientTransport(httpServer.URL) + transport := mcp.NewSSEClientTransport(httpServer.URL, nil) client := mcp.NewClient("test", "v1.0.0", nil) cs, err := client.Connect(ctx, transport) if err != nil { diff --git a/mcp/sse_test.go b/mcp/sse_test.go index aed9982..9a1acd5 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -34,7 +34,7 @@ func TestSSEServer(t *testing.T) { httpServer := httptest.NewServer(sseHandler) defer httpServer.Close() - clientTransport := NewSSEClientTransport(httpServer.URL) + clientTransport := NewSSEClientTransport(httpServer.URL, nil) c := NewClient("testClient", "v1.0.0", nil) cs, err := c.Connect(ctx, clientTransport) diff --git a/mcp/streamable.go b/mcp/streamable.go index 3b860b9..37b0b40 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -571,13 +571,26 @@ func (t *StreamableServerTransport) Close() error { // // TODO(rfindley): support retries and resumption tokens. type StreamableClientTransport struct { - url string + url string + opts StreamableClientTransportOptions +} + +// StreamableClientTransportOptions provides options for the +// [NewStreamableClientTransport] constructor. +type StreamableClientTransportOptions struct { + // HTTPClient is the client to use for making HTTP requests. If nil, + // http.DefaultClient is used. + HTTPClient *http.Client } // NewStreamableClientTransport returns a new client transport that connects to // the streamable HTTP server at the provided URL. -func NewStreamableClientTransport(url string) *StreamableClientTransport { - return &StreamableClientTransport{url: url} +func NewStreamableClientTransport(url string, opts *StreamableClientTransportOptions) *StreamableClientTransport { + t := &StreamableClientTransport{url: url} + if opts != nil { + t.opts = *opts + } + return t } // Connect implements the [Transport] interface. @@ -589,9 +602,13 @@ func NewStreamableClientTransport(url string) *StreamableClientTransport { // When closed, the connection issues a DELETE request to terminate the logical // session. func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, error) { + client := t.opts.HTTPClient + if client == nil { + client = http.DefaultClient + } return &streamableClientConn{ url: t.url, - client: http.DefaultClient, + client: client, incoming: make(chan []byte, 100), done: make(chan struct{}), }, nil diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index c964157..1137d50 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -11,7 +11,9 @@ import ( "fmt" "io" "net/http" + "net/http/cookiejar" "net/http/httptest" + "net/url" "strings" "sync" "sync/atomic" @@ -32,13 +34,35 @@ func TestStreamableTransports(t *testing.T) { server := NewServer("testServer", "v1.0.0", nil) server.AddTools(NewServerTool("greet", "say hi", sayHi)) - // 2. Start an httptest.Server with the StreamableHTTPHandler. + // 2. Start an httptest.Server with the StreamableHTTPHandler, wrapped in a + // cookie-checking middleware. handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) - httpServer := httptest.NewServer(handler) + httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cookie, err := r.Cookie("test-cookie") + if err != nil { + t.Errorf("missing cookie: %v", err) + } else if cookie.Value != "test-value" { + t.Errorf("got cookie %q, want %q", cookie.Value, "test-value") + } + handler.ServeHTTP(w, r) + })) defer httpServer.Close() // 3. Create a client and connect it to the server using our StreamableClientTransport. - transport := NewStreamableClientTransport(httpServer.URL) + // Check that all requests honor a custom client. + jar, err := cookiejar.New(nil) + if err != nil { + t.Fatal(err) + } + u, err := url.Parse(httpServer.URL) + if err != nil { + t.Fatal(err) + } + jar.SetCookies(u, []*http.Cookie{{Name: "test-cookie", Value: "test-value"}}) + httpClient := &http.Client{Jar: jar} + transport := NewStreamableClientTransport(httpServer.URL, &StreamableClientTransportOptions{ + HTTPClient: httpClient, + }) client := NewClient("testClient", "v1.0.0", nil) session, err := client.Connect(ctx, transport) if err != nil {