mcp: reuse server selection for new streamable requests

mcp/streamable.go

@@ -409,25 +409,26 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 	//
 	// [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header
 	protocolVersion := req.Header.Get(protocolVersionHeader)
-	var supportedVersions []string
-	if server := h.getServer(req); server != nil {
-		supportedVersions = server.protocolVersions
-	} else {
-		supportedVersions = supportedProtocolVersions
-	}
-	if protocolVersion != "" && !slices.Contains(supportedVersions, protocolVersion) && protocolVersion < protocolVersion20260728 {
-		http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedVersions, ",")), http.StatusBadRequest)
-		return
-	}
 	req = req.WithContext(context.WithValue(req.Context(), protocolVersionContextKey{}, protocolVersion))
-
 	if h.opts.Stateless {
 		h.serveStateless(w, req)
 	} else {
 		h.serveStateful(w, req)
 	}
 }
 
+func (h *StreamableHTTPHandler) validateProtocolVersion(w http.ResponseWriter, protocolVersion string, server *Server) bool {
+	supportedVersions := supportedProtocolVersions
+	if server != nil {
+		supportedVersions = server.protocolVersions
+	}
+	if protocolVersion != "" && !slices.Contains(supportedVersions, protocolVersion) && protocolVersion < protocolVersion20260728 {
+		http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedVersions, ",")), http.StatusBadRequest)
+		return false
+	}
+	return true
+}
+
 // serveStateless handles requests for stateless servers.
 // Stateless servers only support POST. Each request creates a temporary
 // session that is closed when the request completes.
@@ -462,6 +463,9 @@ func (h *StreamableHTTPHandler) serveStateless(w http.ResponseWriter, req *http.
 	}
 
 	server := h.getServer(req)
+	if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), server) {
+		return
+	}
 	if server == nil {
 		http.Error(w, "no server available", http.StatusBadRequest)
 		return
@@ -666,6 +670,9 @@ func (h *StreamableHTTPHandler) serveStatefulGET(w http.ResponseWriter, req *htt
 	if !ok {
 		return
 	}
+	if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) {
+		return
+	}
 
 	sessInfo.transport.ServeHTTP(w, req)
 }
@@ -683,6 +690,9 @@ func (h *StreamableHTTPHandler) serveStatefulDELETE(w http.ResponseWriter, req *
 	if !ok {
 		return
 	}
+	if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) {
+		return
+	}
 
 	sessInfo.session.Close()
 	w.WriteHeader(http.StatusNoContent)
@@ -712,6 +722,9 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht
 		if !ok {
 			return
 		}
+		if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), sessInfo.session.server) {
+			return
+		}
 		sessInfo.startPOST()
 		defer sessInfo.endPOST()
 		sessInfo.transport.ServeHTTP(w, req)
@@ -720,6 +733,9 @@ func (h *StreamableHTTPHandler) serveStatefulPOST(w http.ResponseWriter, req *ht
 
 	// No session ID: create a new session.
 	server := h.getServer(req)
+	if !h.validateProtocolVersion(w, protocolVersionFromContext(req.Context()), server) {
+		return
+	}
 	if server == nil {
 		http.Error(w, "no server available", http.StatusBadRequest)
 		return

mcp/streamable_test.go

@@ -4973,3 +4973,65 @@ func TestStreamableSupportedProtocolVersions_Header(t *testing.T) {
 		}
 	})
 }
+
+func TestStreamableGetServerCalledOnceForNewPOST(t *testing.T) {
+	for _, stateless := range []bool{false, true} {
+		t.Run(fmt.Sprintf("stateless=%v", stateless), func(t *testing.T) {
+			server := NewServer(testImpl, nil)
+			var calls atomic.Int32
+			handler := NewStreamableHTTPHandler(func(*http.Request) *Server {
+				calls.Add(1)
+				return server
+			}, &StreamableHTTPOptions{Stateless: stateless})
+
+			req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{}`))
+			req.Header.Set("Content-Type", "application/json")
+			req.Header.Set("Accept", "application/json, text/event-stream")
+			handler.ServeHTTP(httptest.NewRecorder(), req)
+
+			if got := calls.Load(); got != 1 {
+				t.Errorf("getServer called %d times, want 1", got)
+			}
+		})
+	}
+
+	t.Run("existing stateful session", func(t *testing.T) {
+		server := NewServer(testImpl, nil)
+		var calls atomic.Int32
+		handler := NewStreamableHTTPHandler(func(*http.Request) *Server {
+			calls.Add(1)
+			return server
+		}, &StreamableHTTPOptions{JSONResponse: true})
+
+		post := func(body string, sessionID string) *httptest.ResponseRecorder {
+			req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
+			req.Header.Set("Content-Type", "application/json")
+			req.Header.Set("Accept", "application/json, text/event-stream")
+			if sessionID != "" {
+				req.Header.Set(sessionIDHeader, sessionID)
+			}
+			rec := httptest.NewRecorder()
+			handler.ServeHTTP(rec, req)
+			return rec
+		}
+
+		initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}`
+		initRec := post(initBody, "")
+		if initRec.Code != http.StatusOK {
+			t.Fatalf("initialize status = %d, want %d: %s", initRec.Code, http.StatusOK, initRec.Body)
+		}
+		sessionID := initRec.Header().Get(sessionIDHeader)
+		if sessionID == "" {
+			t.Fatal("initialize response missing session ID")
+		}
+
+		pingRec := post(`{"jsonrpc":"2.0","id":2,"method":"ping"}`, sessionID)
+		if pingRec.Code != http.StatusOK {
+			t.Fatalf("ping status = %d, want %d: %s", pingRec.Code, http.StatusOK, pingRec.Body)
+		}
+		if got := calls.Load(); got != 1 {
+			t.Errorf("getServer called %d times, want 1 for new session only", got)
+		}
+		handler.closeAll()
+	})
+}