internal/mcp: implement streamable HTTP transport

Add an (almost) complete implementation of the streamable HTTP transport, including: - A StreamableHTTPHandler for serving sessions. - A StreamableHTTPServerTransport that supports session resumption. - A minimal StreamableHTTPClientTransport. Additionally, include the following minor changes: - Refactor SSE event reading to use an iterator. - Start to to aliases for JSON-RPC types (incomplete) - Fix handling for null tool arguments (again) Still to do: - The streamable client does not support retry/resumption, or the side-channel GET requests. - Support configurable session and event retention (garbage collection). - Expose the HTTP client used for HTTP requests (for auth). Change-Id: Ieba216ef901169b4660d601386fa0c7c7bb4c4b2 Reviewed-on: https://go-review.googlesource.com/c/tools/+/667515 Reviewed-by: Jonathan Amsterdam <jba@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
2025-04-23 10:52:44 +00:00 · 2025-04-23 10:52:44 +00:00 · d2f2f87bbe
commit d2f2f87bbe
parent 11b17f5858
1 changed files with 1570 additions and 149 deletions
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@ -0,0 +1,516 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package mcp
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
+)
+
+func TestStreamableTransports(t *testing.T) {
+	// This test checks that the streamable server and client transports can
+	// communicate.
+
+	ctx := context.Background()
+
+	// 1. Create a server with a simple "greet" tool.
+	server := NewServer("testServer", "v1.0.0", nil)
+	server.AddTools(NewTool("greet", "say hi", sayHi))
+
+	// 2. Start an httptest.Server with the StreamableHTTPHandler.
+	handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
+	httpServer := httptest.NewServer(handler)
+	defer httpServer.Close()
+
+	// 3. Create a client and connect it to the server using our StreamableClientTransport.
+	transport := NewStreamableClientTransport(httpServer.URL)
+	client := NewClient("testClient", "v1.0.0", nil)
+	session, err := client.Connect(ctx, transport)
+	if err != nil {
+		t.Fatalf("client.Connect() failed: %v", err)
+	}
+	defer session.Close()
+
+	// 4. The client calls the "greet" tool.
+	params := &CallToolParams{
+		Name:      "greet",
+		Arguments: map[string]any{"name": "streamy"},
+	}
+	got, err := session.CallTool(ctx, params)
+	if err != nil {
+		t.Fatalf("CallTool() failed: %v", err)
+	}
+
+	// 5. Verify that the correct response is received.
+	want := &CallToolResult{
+		Content: []*Content{{Type: "text", Text: "hi streamy"}},
+	}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("CallTool() returned unexpected content (-want +got):\n%s", diff)
+	}
+}
+
+func TestStreamableServerTransport(t *testing.T) {
+	// This test checks detailed behavior of the streamable server transport, by
+	// faking the behavior of a streamable client using a sequence of HTTP
+	// requests.
+
+	// A step is a single step in the tests below, consisting of a request payload
+	// and expected response.
+	type step struct {
+		// If OnRequest is > 0, this step only executes after a request with the
+		// given ID is received.
+		//
+		// All OnRequest steps must occur before the step that creates the request.
+		//
+		// To avoid tests hanging when there's a bug, it's expected that this
+		// request is received in the course of a *synchronous* request to the
+		// server (otherwise, we wouldn't be able to terminate the test without
+		// analyzing a dependency graph).
+		OnRequest int64
+		// If set, Async causes the step to run asynchronously to other steps.
+		// Redundant with OnRequest: all OnRequest steps are asynchronous.
+		Async bool
+
+		Method     string           // HTTP request method
+		Send       []JSONRPCMessage // messages to send
+		CloseAfter int              // if nonzero, close after receiving this many messages
+		StatusCode int              // expected status code
+		Recv       []JSONRPCMessage // expected messages to receive
+	}
+
+	// JSON-RPC message constructors.
+	req := func(id int64, method string, params any) *JSONRPCRequest {
+		r := &JSONRPCRequest{
+			Method: method,
+			Params: mustMarshal(t, params),
+		}
+		if id > 0 {
+			r.ID = jsonrpc2.Int64ID(id)
+		}
+		return r
+	}
+	resp := func(id int64, result any, err error) *JSONRPCResponse {
+		return &JSONRPCResponse{
+			ID:     jsonrpc2.Int64ID(id),
+			Result: mustMarshal(t, result),
+			Error:  err,
+		}
+	}
+
+	// Predefined steps, to avoid repetition below.
+	initReq := req(1, "initialize", &InitializeParams{})
+	initResp := resp(1, &InitializeResult{
+		Capabilities: &serverCapabilities{
+			Logging:   &loggingCapabilities{},
+			Prompts:   &promptCapabilities{ListChanged: true},
+			Resources: &resourceCapabilities{ListChanged: true},
+			Tools:     &toolCapabilities{ListChanged: true},
+		},
+		ProtocolVersion: "2024-11-05",
+		ServerInfo:      &implementation{Name: "testServer", Version: "v1.0.0"},
+	}, nil)
+	initializedMsg := req(0, "initialized", &InitializedParams{})
+	initialize := step{
+		Method:     "POST",
+		Send:       []JSONRPCMessage{initReq},
+		StatusCode: http.StatusOK,
+		Recv:       []JSONRPCMessage{initResp},
+	}
+	initialized := step{
+		Method:     "POST",
+		Send:       []JSONRPCMessage{initializedMsg},
+		StatusCode: http.StatusAccepted,
+	}
+
+	tests := []struct {
+		name  string
+		tool  func(*testing.T, context.Context, *ServerSession)
+		steps []step
+	}{
+		{
+			name: "basic",
+			steps: []step{
+				initialize,
+				initialized,
+				{
+					Method:     "POST",
+					Send:       []JSONRPCMessage{req(2, "tools/call", &CallToolParams{Name: "tool"})},
+					StatusCode: http.StatusOK,
+					Recv:       []JSONRPCMessage{resp(2, &CallToolResult{}, nil)},
+				},
+			},
+		},
+		{
+			name: "tool notification",
+			tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
+				// Send an arbitrary notification.
+				if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
+					t.Errorf("Notify failed: %v", err)
+				}
+			},
+			steps: []step{
+				initialize,
+				initialized,
+				{
+					Method: "POST",
+					Send: []JSONRPCMessage{
+						req(2, "tools/call", &CallToolParams{Name: "tool"}),
+					},
+					StatusCode: http.StatusOK,
+					Recv: []JSONRPCMessage{
+						req(0, "notifications/progress", &ProgressNotificationParams{}),
+						resp(2, &CallToolResult{}, nil),
+					},
+				},
+			},
+		},
+		{
+			name: "tool upcall",
+			tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
+				// Make an arbitrary call.
+				if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil {
+					t.Errorf("Call failed: %v", err)
+				}
+			},
+			steps: []step{
+				initialize,
+				initialized,
+				{
+					Method:    "POST",
+					OnRequest: 1,
+					Send: []JSONRPCMessage{
+						resp(1, &ListRootsResult{}, nil),
+					},
+					StatusCode: http.StatusAccepted,
+				},
+				{
+					Method: "POST",
+					Send: []JSONRPCMessage{
+						req(2, "tools/call", &CallToolParams{Name: "tool"}),
+					},
+					StatusCode: http.StatusOK,
+					Recv: []JSONRPCMessage{
+						req(1, "roots/list", &ListRootsParams{}),
+						resp(2, &CallToolResult{}, nil),
+					},
+				},
+			},
+		},
+		{
+			name: "background",
+			tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
+				// Perform operations on a background context, and ensure the client
+				// receives it.
+				ctx = context.Background()
+				if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
+					t.Errorf("Notify failed: %v", err)
+				}
+				// TODO(rfindley): finish implementing logging.
+				// if err := ss.LoggingMessage(ctx, &LoggingMessageParams{}); err != nil {
+				// 	t.Errorf("Logging failed: %v", err)
+				// }
+				if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil {
+					t.Errorf("Notify failed: %v", err)
+				}
+			},
+			steps: []step{
+				initialize,
+				initialized,
+				{
+					Method:    "POST",
+					OnRequest: 1,
+					Send: []JSONRPCMessage{
+						resp(1, &ListRootsResult{}, nil),
+					},
+					StatusCode: http.StatusAccepted,
+				},
+				{
+					Method:     "GET",
+					Async:      true,
+					StatusCode: http.StatusOK,
+					CloseAfter: 2,
+					Recv: []JSONRPCMessage{
+						req(0, "notifications/progress", &ProgressNotificationParams{}),
+						req(1, "roots/list", &ListRootsParams{}),
+					},
+				},
+				{
+					Method: "POST",
+					Send: []JSONRPCMessage{
+						req(2, "tools/call", &CallToolParams{Name: "tool"}),
+					},
+					StatusCode: http.StatusOK,
+					Recv: []JSONRPCMessage{
+						resp(2, &CallToolResult{}, nil),
+					},
+				},
+			},
+		},
+		{
+			name: "errors",
+			steps: []step{
+				{
+					Method:     "PUT",
+					StatusCode: http.StatusMethodNotAllowed,
+				},
+				{
+					Method:     "DELETE",
+					StatusCode: http.StatusBadRequest,
+				},
+				{
+					Method:     "POST",
+					Send:       []JSONRPCMessage{req(2, "tools/call", &CallToolParams{Name: "tool"})},
+					StatusCode: http.StatusOK,
+					Recv: []JSONRPCMessage{resp(2, nil, &jsonrpc2.WireError{
+						Message: `method "tools/call" is invalid during session initialization`,
+					})},
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			// Create a server containing a single tool, which runs the test tool
+			// behavior, if any.
+			server := NewServer("testServer", "v1.0.0", nil)
+			tool := NewTool("tool", "test tool", func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) {
+				if test.tool != nil {
+					test.tool(t, ctx, ss)
+				}
+				return &CallToolResultFor[any]{}, nil
+			})
+			server.AddTools(tool)
+
+			// Start the streamable handler.
+			handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
+			defer handler.closeAll()
+
+			httpServer := httptest.NewServer(handler)
+			defer httpServer.Close()
+
+			// blocks records request blocks by JSONRPC ID.
+			//
+			// When an OnRequest step is encountered, it waits on the corresponding
+			// block. When a request with that ID is received, the block is closed.
+			var mu sync.Mutex
+			blocks := make(map[int64]chan struct{})
+			for _, step := range test.steps {
+				if step.OnRequest > 0 {
+					blocks[step.OnRequest] = make(chan struct{})
+				}
+			}
+
+			// signal when all synchronous requests have executed, so we can fail
+			// async requests that are blocked.
+			syncRequestsDone := make(chan struct{})
+
+			// To avoid complicated accounting for session ID, just set the first
+			// non-empty session ID from a response.
+			var sessionID atomic.Value
+			sessionID.Store("")
+
+			// doStep executes a single step.
+			doStep := func(t *testing.T, step step) {
+				if step.OnRequest > 0 {
+					// Block the step until we've received the server->client request.
+					mu.Lock()
+					block := blocks[step.OnRequest]
+					mu.Unlock()
+					select {
+					case <-block:
+					case <-syncRequestsDone:
+						t.Errorf("after all sync requests are complete, request still blocked on %d", step.OnRequest)
+						return
+					}
+				}
+
+				// Collect messages received during this request, unblock other steps
+				// when requests are received.
+				var got []JSONRPCMessage
+				out := make(chan JSONRPCMessage)
+				// Cancel the step if we encounter a request that isn't going to be
+				// handled.
+				ctx, cancel := context.WithCancel(context.Background())
+
+				var wg sync.WaitGroup
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+
+					for m := range out {
+						if req, ok := m.(*JSONRPCRequest); ok && req.ID.IsValid() {
+							// Encountered a server->client request. We should have a
+							// response queued. Otherwise, we may deadlock.
+							mu.Lock()
+							if block, ok := blocks[req.ID.Raw().(int64)]; ok {
+								close(block)
+							} else {
+								t.Errorf("no queued response for %v", req.ID)
+								cancel()
+							}
+							mu.Unlock()
+						}
+						got = append(got, m)
+						if step.CloseAfter > 0 && len(got) == step.CloseAfter {
+							cancel()
+						}
+					}
+				}()
+
+				gotSessionID, gotStatusCode, err := streamingRequest(ctx,
+					httpServer.URL, sessionID.Load().(string), step.Method, step.Send, out)
+
+				// Don't fail on cancelled requests: error (if any) is handled
+				// elsewhere.
+				if err != nil && ctx.Err() == nil {
+					t.Fatal(err)
+				}
+
+				if gotStatusCode != step.StatusCode {
+					t.Errorf("got status %d, want %d", gotStatusCode, step.StatusCode)
+				}
+				wg.Wait()
+
+				transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id JSONRPCID) any { return id.Raw() })
+				if diff := cmp.Diff(step.Recv, got, transform); diff != "" {
+					t.Errorf("received unexpected messages (-want +got):\n%s", diff)
+				}
+				sessionID.CompareAndSwap("", gotSessionID)
+			}
+
+			var wg sync.WaitGroup
+			for _, step := range test.steps {
+				if step.Async || step.OnRequest > 0 {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						doStep(t, step)
+					}()
+				} else {
+					doStep(t, step)
+				}
+			}
+
+			// Fail any blocked responses if they weren't needed by a synchronous
+			// request.
+			close(syncRequestsDone)
+
+			wg.Wait()
+		})
+	}
+}
+
+// streamingRequest makes a request to the given streamable server with the
+// given url, sessionID, and method.
+//
+// If provided, the in messages are encoded in the request body. A single
+// message is encoded as a JSON object. Multiple messages are batched as a JSON
+// array.
+//
+// Any received messages are sent to the out channel, which is closed when the
+// request completes.
+//
+// Returns the sessionID and http status code from the response. If an error is
+// returned, sessionID and status code may still be set if the error occurs
+// after the response headers have been received.
+func streamingRequest(ctx context.Context, serverURL, sessionID, method string, in []JSONRPCMessage, out chan<- JSONRPCMessage) (string, int, error) {
+	defer close(out)
+
+	var body []byte
+	if len(in) == 1 {
+		data, err := jsonrpc2.EncodeMessage(in[0])
+		if err != nil {
+			return "", 0, fmt.Errorf("encoding message: %w", err)
+		}
+		body = data
+	} else {
+		var rawMsgs []json.RawMessage
+		for _, msg := range in {
+			data, err := jsonrpc2.EncodeMessage(msg)
+			if err != nil {
+				return "", 0, fmt.Errorf("encoding message: %w", err)
+			}
+			rawMsgs = append(rawMsgs, data)
+		}
+		data, err := json.Marshal(rawMsgs)
+		if err != nil {
+			return "", 0, fmt.Errorf("marshaling batch: %w", err)
+		}
+		body = data
+	}
+
+	req, err := http.NewRequestWithContext(ctx, method, serverURL, bytes.NewReader(body))
+	if err != nil {
+		return "", 0, fmt.Errorf("creating request: %w", err)
+	}
+	if sessionID != "" {
+		req.Header.Set("Mcp-Session-Id", sessionID)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Add("Accept", "text/plain") // ensure multiple accept headers are allowed
+	req.Header.Add("Accept", "application/json, text/event-stream")
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", 0, fmt.Errorf("request failed: %v", err)
+	}
+	defer resp.Body.Close()
+
+	newSessionID := resp.Header.Get("Mcp-Session-Id")
+
+	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
+		for evt, err := range scanEvents(resp.Body) {
+			if err != nil {
+				return newSessionID, resp.StatusCode, fmt.Errorf("reading events: %v", err)
+			}
+			// TODO(rfindley): do we need to check evt.name?
+			// Does the MCP spec say anything about this?
+			msg, err := jsonrpc2.DecodeMessage(evt.data)
+			if err != nil {
+				return newSessionID, resp.StatusCode, fmt.Errorf("decoding message: %w", err)
+			}
+			out <- msg
+		}
+	} else if strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") {
+		data, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return newSessionID, resp.StatusCode, fmt.Errorf("reading json body: %w", err)
+		}
+		msg, err := jsonrpc2.DecodeMessage(data)
+		if err != nil {
+			return newSessionID, resp.StatusCode, fmt.Errorf("decoding message: %w", err)
+		}
+		out <- msg
+	}
+
+	return newSessionID, resp.StatusCode, nil
+}
+
+func mustMarshal(t *testing.T, v any) json.RawMessage {
+	if v == nil {
+		return nil
+	}
+	t.Helper()
+	data, err := json.Marshal(v)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return data
+}