mcp: factor out reusable functions from TestStreamableServerTransport

The next CL will test stateless and distributable server transport
configurations, using the HTTP testing strategy of
TestStreamableServerTransport.

mcp/server.go

@@ -355,6 +355,8 @@ func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParam
 	if !ok {
 		return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, req.Params.Name)
 	}
+	// TODO: if handler returns nil content, it will serialize as null.
+	// Add a test and fix.
 	return st.handler(ctx, req)
 }
 

mcp/streamable_test.go

@@ -382,30 +382,31 @@ func readNotifications(t *testing.T, ctx context.Context, notifications chan str
 	}
 }
 
-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.
-
 // JSON-RPC message constructors.
-	req := func(id int64, method string, params any) *jsonrpc.Request {
+func req(id int64, method string, params any) *jsonrpc.Request {
 	r := &jsonrpc.Request{
 		Method: method,
-			Params: mustMarshal(t, params),
+		Params: mustMarshal(params),
 	}
 	if id > 0 {
 		r.ID = jsonrpc2.Int64ID(id)
 	}
 	return r
 }
-	resp := func(id int64, result any, err error) *jsonrpc.Response {
+
+func resp(id int64, result any, err error) *jsonrpc.Response {
 	return &jsonrpc.Response{
 		ID:     jsonrpc2.Int64ID(id),
-			Result: mustMarshal(t, result),
+		Result: mustMarshal(result),
 		Error:  err,
 	}
 }
 
+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.
+
 	// Predefined steps, to avoid repetition below.
 	initReq := req(1, methodInitialize, &InitializeParams{})
 	initResp := resp(1, &InitializeResult{
@@ -422,21 +423,23 @@ func TestStreamableServerTransport(t *testing.T) {
 		messages:       []jsonrpc.Message{initReq},
 		wantStatusCode: http.StatusOK,
 		wantMessages:   []jsonrpc.Message{initResp},
+		wantSessionID:  true,
 	}
 	initialized := streamableRequest{
 		method:         "POST",
 		messages:       []jsonrpc.Message{initializedMsg},
 		wantStatusCode: http.StatusAccepted,
+		wantSessionID:  false, // TODO: should this be true?
 	}
 
 	tests := []struct {
 		name     string
 		tool     func(*testing.T, context.Context, *ServerSession)
-		steps []streamableRequest
+		requests []streamableRequest // http requests
 	}{
 		{
 			name: "basic",
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				initialize,
 				initialized,
 				{
@@ -444,12 +447,13 @@ func TestStreamableServerTransport(t *testing.T) {
 					messages:       []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
 					wantStatusCode: http.StatusOK,
 					wantMessages:   []jsonrpc.Message{resp(2, &CallToolResult{}, nil)},
+					wantSessionID:  true,
 				},
 			},
 		},
 		{
 			name: "accept headers",
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				initialize,
 				initialized,
 				// Test various accept headers.
@@ -458,12 +462,14 @@ func TestStreamableServerTransport(t *testing.T) {
 					accept:         []string{"text/plain", "application/*"},
 					messages:       []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})},
 					wantStatusCode: http.StatusBadRequest, // missing text/event-stream
+					wantSessionID:  false,
 				},
 				{
 					method:         "POST",
 					accept:         []string{"text/event-stream"},
 					messages:       []jsonrpc.Message{req(3, "tools/call", &CallToolParams{Name: "tool"})},
 					wantStatusCode: http.StatusBadRequest, // missing application/json
+					wantSessionID:  false,
 				},
 				{
 					method:         "POST",
@@ -471,6 +477,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					messages:       []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})},
 					wantStatusCode: http.StatusOK,
 					wantMessages:   []jsonrpc.Message{resp(4, &CallToolResult{}, nil)},
+					wantSessionID:  true,
 				},
 				{
 					method:         "POST",
@@ -478,6 +485,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					messages:       []jsonrpc.Message{req(4, "tools/call", &CallToolParams{Name: "tool"})},
 					wantStatusCode: http.StatusOK,
 					wantMessages:   []jsonrpc.Message{resp(4, &CallToolResult{}, nil)},
+					wantSessionID:  true,
 				},
 			},
 		},
@@ -489,7 +497,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					t.Errorf("Notify failed: %v", err)
 				}
 			},
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				initialize,
 				initialized,
 				{
@@ -502,6 +510,7 @@ func TestStreamableServerTransport(t *testing.T) {
 						req(0, "notifications/progress", &ProgressNotificationParams{}),
 						resp(2, &CallToolResult{}, nil),
 					},
+					wantSessionID: true,
 				},
 			},
 		},
@@ -513,7 +522,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					t.Errorf("Call failed: %v", err)
 				}
 			},
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				initialize,
 				initialized,
 				{
@@ -523,6 +532,7 @@ func TestStreamableServerTransport(t *testing.T) {
 						resp(1, &ListRootsResult{}, nil),
 					},
 					wantStatusCode: http.StatusAccepted,
+					wantSessionID:  false,
 				},
 				{
 					method: "POST",
@@ -534,6 +544,7 @@ func TestStreamableServerTransport(t *testing.T) {
 						req(1, "roots/list", &ListRootsParams{}),
 						resp(2, &CallToolResult{}, nil),
 					},
+					wantSessionID: true,
 				},
 			},
 		},
@@ -554,7 +565,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					t.Errorf("Notify failed: %v", err)
 				}
 			},
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				initialize,
 				initialized,
 				{
@@ -564,6 +575,7 @@ func TestStreamableServerTransport(t *testing.T) {
 						resp(1, &ListRootsResult{}, nil),
 					},
 					wantStatusCode: http.StatusAccepted,
+					wantSessionID:  false,
 				},
 				{
 					method:         "GET",
@@ -574,6 +586,7 @@ func TestStreamableServerTransport(t *testing.T) {
 						req(0, "notifications/progress", &ProgressNotificationParams{}),
 						req(1, "roots/list", &ListRootsParams{}),
 					},
+					wantSessionID: true,
 				},
 				{
 					method: "POST",
@@ -584,12 +597,13 @@ func TestStreamableServerTransport(t *testing.T) {
 					wantMessages: []jsonrpc.Message{
 						resp(2, &CallToolResult{}, nil),
 					},
+					wantSessionID: true,
 				},
 			},
 		},
 		{
 			name: "errors",
-			steps: []streamableRequest{
+			requests: []streamableRequest{
 				{
 					method:         "PUT",
 					wantStatusCode: http.StatusMethodNotAllowed,
@@ -615,6 +629,7 @@ func TestStreamableServerTransport(t *testing.T) {
 					wantMessages: []jsonrpc.Message{resp(2, nil, &jsonrpc2.WireError{
 						Message: `method "tools/call" is invalid during session initialization`,
 					})},
+					wantSessionID: true, // TODO: this is probably wrong; we don't have a valid session
 				},
 			},
 		},
@@ -636,6 +651,12 @@ func TestStreamableServerTransport(t *testing.T) {
 			handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
 			defer handler.closeAll()
 
+			testStreamableHandler(t, handler, test.requests)
+		})
+	}
+}
+
+func testStreamableHandler(t *testing.T, handler http.Handler, requests []streamableRequest) {
 	httpServer := httptest.NewServer(handler)
 	defer httpServer.Close()
 
@@ -645,9 +666,9 @@ func TestStreamableServerTransport(t *testing.T) {
 	// 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 {
+	for _, req := range requests {
-				if step.onRequest > 0 {
+		if req.onRequest > 0 {
-					blocks[step.onRequest] = make(chan struct{})
+			blocks[req.onRequest] = make(chan struct{})
 		}
 	}
 
@@ -661,16 +682,16 @@ func TestStreamableServerTransport(t *testing.T) {
 	sessionID.Store("")
 
 	// doStep executes a single step.
-			doStep := func(t *testing.T, step streamableRequest) {
+	doStep := func(t *testing.T, i int, request streamableRequest) {
-				if step.onRequest > 0 {
+		if request.onRequest > 0 {
 			// Block the step until we've received the server->client request.
 			mu.Lock()
-					block := blocks[step.onRequest]
+			block := blocks[request.onRequest]
 			mu.Unlock()
 			select {
 			case <-block:
 			case <-syncRequestsDone:
-						t.Errorf("after all sync requests are complete, request still blocked on %d", step.onRequest)
+				t.Errorf("after all sync requests are complete, request still blocked on %d", request.onRequest)
 				return
 			}
 		}
@@ -702,13 +723,13 @@ func TestStreamableServerTransport(t *testing.T) {
 					mu.Unlock()
 				}
 				got = append(got, m)
-						if step.closeAfter > 0 && len(got) == step.closeAfter {
+				if request.closeAfter > 0 && len(got) == request.closeAfter {
 					cancel()
 				}
 			}
 		}()
 
-				gotSessionID, gotStatusCode, err := step.do(ctx, httpServer.URL, sessionID.Load().(string), out)
+		gotSessionID, gotStatusCode, err := request.do(ctx, httpServer.URL, sessionID.Load().(string), out)
 
 		// Don't fail on cancelled requests: error (if any) is handled
 		// elsewhere.
@@ -716,28 +737,33 @@ func TestStreamableServerTransport(t *testing.T) {
 			t.Fatal(err)
 		}
 
-				if gotStatusCode != step.wantStatusCode {
+		if gotStatusCode != request.wantStatusCode {
-					t.Errorf("got status %d, want %d", gotStatusCode, step.wantStatusCode)
+			t.Errorf("request #%d: got status %d, want %d", i, gotStatusCode, request.wantStatusCode)
+		}
+		if got := gotSessionID != ""; got != request.wantSessionID {
+			t.Errorf("request #%d: got session id: %t, want %t", i, got, request.wantSessionID)
 		}
 		wg.Wait()
 
+		if !request.ignoreResponse {
 			transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() })
-				if diff := cmp.Diff(step.wantMessages, got, transform); diff != "" {
+			if diff := cmp.Diff(request.wantMessages, 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 {
+	for i, request := range requests {
-				if step.async || step.onRequest > 0 {
+		if request.async || request.onRequest > 0 {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
-						doStep(t, step)
+				doStep(t, i, request)
 			}()
 		} else {
-					doStep(t, step)
+			doStep(t, i, request)
 		}
 	}
 
@@ -746,8 +772,6 @@ func TestStreamableServerTransport(t *testing.T) {
 	close(syncRequestsDone)
 
 	wg.Wait()
-		})
-	}
 }
 
 // A streamableRequest describes a single streamable HTTP request, consisting
@@ -768,13 +792,15 @@ type streamableRequest struct {
 	async bool
 
 	// Request attributes
-	method   string            // HTTP request method
+	method   string            // HTTP request method (required)
 	accept   []string          // if non-empty, the Accept header to use; otherwise the default header is used
 	messages []jsonrpc.Message // messages to send
 
 	closeAfter     int               // if nonzero, close after receiving this many messages
 	wantStatusCode int               // expected status code
+	ignoreResponse bool              // if set, don't check the response messages
 	wantMessages   []jsonrpc.Message // expected messages to receive
+	wantSessionID  bool              // whether or not a session ID is expected in the response
 }
 
 // streamingRequest makes a request to the given streamable server with the
@@ -840,7 +866,8 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string,
 
 	newSessionID := resp.Header.Get("Mcp-Session-Id")
 
-	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
+	contentType := resp.Header.Get("Content-Type")
+	if strings.HasPrefix(contentType, "text/event-stream") {
 		for evt, err := range scanEvents(resp.Body) {
 			if err != nil {
 				return newSessionID, resp.StatusCode, fmt.Errorf("reading events: %v", err)
@@ -853,7 +880,7 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string,
 			}
 			out <- msg
 		}
-	} else if strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") {
+	} else if strings.HasPrefix(contentType, "application/json") {
 		data, err := io.ReadAll(resp.Body)
 		if err != nil {
 			return newSessionID, resp.StatusCode, fmt.Errorf("reading json body: %w", err)
@@ -868,14 +895,13 @@ func (s streamableRequest) do(ctx context.Context, serverURL, sessionID string,
 	return newSessionID, resp.StatusCode, nil
 }
 
-func mustMarshal(t *testing.T, v any) json.RawMessage {
+func mustMarshal(v any) json.RawMessage {
 	if v == nil {
 		return nil
 	}
-	t.Helper()
 	data, err := json.Marshal(v)
 	if err != nil {
-		t.Fatal(err)
+		panic(err)
 	}
 	return data
 }
@@ -886,7 +912,7 @@ func TestStreamableClientTransportApplicationJSON(t *testing.T) {
 	resp := func(id int64, result any, err error) *jsonrpc.Response {
 		return &jsonrpc.Response{
 			ID:     jsonrpc2.Int64ID(id),
-			Result: mustMarshal(t, result),
+			Result: mustMarshal(result),
 			Error:  err,
 		}
 	}
@@ -970,13 +996,10 @@ func TestEventID(t *testing.T) {
 }
 
 func TestStreamableStateless(t *testing.T) {
-	// Test stateless mode behavior
-	ctx := context.Background()
-
 	// This version of sayHi doesn't make a ping request (we can't respond to
 	// that request from our client).
-	sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) {
+	sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResult, error) {
-		return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
+		return &CallToolResult{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
 	}
 	server := NewServer(testImpl, nil)
 	AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
@@ -985,57 +1008,27 @@ func TestStreamableStateless(t *testing.T) {
 	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
 		GetSessionID: func() string { return "" },
 	})
-	httpServer := httptest.NewServer(handler)
-	defer httpServer.Close()
 
-	checkRequest := func(body string) {
+	requests := []streamableRequest{
-		// Verify we can call tools/list directly without initialization in stateless mode
+		{
-		req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, strings.NewReader(body))
+			method:         "POST",
-		if err != nil {
+			wantStatusCode: http.StatusOK,
-			t.Fatal(err)
+			messages:       []jsonrpc.Message{req(1, "tools/list", struct{}{})},
-		}
+			ignoreResponse: true,
-		req.Header.Set("Content-Type", "application/json")
+			wantSessionID:  false,
-		req.Header.Set("Accept", "application/json, text/event-stream")
+		},
-
+		{
-		resp, err := http.DefaultClient.Do(req)
+			method:         "POST",
-		if err != nil {
+			wantStatusCode: http.StatusOK,
-			t.Fatal(err)
+			messages: []jsonrpc.Message{
-		}
+				req(2, "tools/call", &CallToolParams{Name: "greet", Arguments: hiParams{Name: "World"}}),
-		defer resp.Body.Close()
+			},
-
+			wantMessages: []jsonrpc.Message{
-		// Verify that no session ID header is returned in stateless mode
+				resp(2, &CallToolResult{Content: []Content{&TextContent{Text: "hi World"}}}, nil),
-		if sessionID := resp.Header.Get(sessionIDHeader); sessionID != "" {
+			},
-			t.Errorf("%s = %s, want no session ID header", sessionIDHeader, sessionID)
+			wantSessionID: false,
-		}
+		},
-		if resp.StatusCode != http.StatusOK {
-			t.Errorf("Status code = %d; want successful response", resp.StatusCode)
 	}
 
-		var events []Event
+	testStreamableHandler(t, handler, requests)
-		for event, err := range scanEvents(resp.Body) {
-			if err != nil {
-				t.Fatal(err)
-			}
-			events = append(events, event)
-		}
-		if len(events) != 1 {
-			t.Fatalf("got %d SSE events, want 1; events: %v", len(events), events)
-		}
-		msg, err := jsonrpc.DecodeMessage(events[0].Data)
-		if err != nil {
-			t.Fatal(err)
-		}
-		jsonResp, ok := msg.(*jsonrpc.Response)
-		if !ok {
-			t.Errorf("event is %T, want response", jsonResp)
-		}
-		if jsonResp.Error != nil {
-			t.Errorf("request failed: %v", jsonResp.Error)
-		}
-	}
-
-	checkRequest(`{"jsonrpc":"2.0","method":"tools/list","id":1,"params":{}}`)
-
-	// Verify we can make another request without session ID
-	checkRequest(`{"jsonrpc":"2.0","method":"tools/call","id":2,"params":{"name":"greet","arguments":{"name":"World"}}}`)
 }