refactor: delay initialization state updates and enforce protocol version for testing

Author: guglielmo-san

mcp/client_test.go

@@ -814,7 +814,7 @@ func TestClientConnectDiscover_RequestContents(t *testing.T) {
 			return nil, nil
 		},
 	})
-	cs, err := c.Connect(ctx, ct, nil)
+	cs, err := c.Connect(ctx, ct, &ClientSessionOptions{protocolVersion: protocolVersion20260630})
 	if err != nil {
 		t.Fatalf("Connect: %v", err)
 	}

mcp/server.go

@@ -1475,12 +1475,6 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
 		return nil, perRequestErr
 	}
 
-	if !initialized && validatedMeta.usesNewProtocol && validatedMeta.initializeParams != nil {
-		ss.updateState(func(state *ServerSessionState) {
-			state.InitializeParams = validatedMeta.initializeParams
-		})
-	}
-
 	switch req.Method {
 	case methodInitialize, methodPing, notificationInitialized:
 		if validatedMeta.usesNewProtocol {
@@ -1490,11 +1484,20 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
 				Message: fmt.Sprintf("%q is not supported in the new protocol", req.Method),
 			}
 		}
+	case methodDiscover:
+		// In case of methodDiscover call the state.initializeParams is populated
+		// within the discover handle function to make sure the method is supported
+		// when the user is probing a pre-2026-06-30 server.
 	default:
 		if !initialized && !validatedMeta.usesNewProtocol {
 			ss.server.opts.Logger.Error("method invalid during initialization", "method", req.Method)
 			return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method)
 		}
+		if !initialized && validatedMeta.usesNewProtocol && validatedMeta.initializeParams != nil {
+			ss.updateState(func(state *ServerSessionState) {
+				state.InitializeParams = validatedMeta.initializeParams
+			})
+		}
 	}
 
 	// modelcontextprotocol/go-sdk#26: handle calls asynchronously, and

mcp/streamable_client_test.go

@@ -1290,7 +1290,7 @@ func TestStreamableClientConnect_DiscoverSuccess(t *testing.T) {
 
 	transport := &StreamableClientTransport{Endpoint: httpServer.URL}
 	client := NewClient(testImpl, nil)
-	session, err := client.Connect(ctx, transport, nil)
+	session, err := client.Connect(ctx, transport, &ClientSessionOptions{protocolVersion: protocolVersion20260630})
 	if err != nil {
 		t.Fatalf("Connect: %v", err)
 	}
@@ -1513,7 +1513,7 @@ func TestStreamableClientConnect_DiscoverPropagatesOtherErrors(t *testing.T) {
 
 	transport := &StreamableClientTransport{Endpoint: httpServer.URL}
 	client := NewClient(testImpl, nil)
-	session, err := client.Connect(ctx, transport, nil)
+	session, err := client.Connect(ctx, transport, &ClientSessionOptions{protocolVersion: protocolVersion20260630})
 	if err == nil {
 		_ = session.Close()
 		t.Fatal("Connect succeeded; want propagated error")