Ajax table form with tabledrag only works properly AFTER first submit

enter image description here

So this one has me kind of stumped.

With a simple table with 4 rows, I drag one row to another position, then hit save (submit).
When the ajax returns and replaces the form, the rows snap back to the original position. If I reload the page, I can see the order in which I dragged, did indeed get saved.

I reload the page, starting over, but this time start by pressing save, before dragging anything – and nothing appears to happen. But all form inputs id’s and names now have a random-string, e.g. –XjidQzhH, appended to them. I can now drag a row and hit save again – and it actually works! It stays where I dragged it, also when reloading the page.

Once the form has been replaced the first time, everything seems to work as it should.

Hopefully someone knows what I’m doing wrong here – I’ve been rabbit-holing this for two days now and can’t figure out what I’m missing.

Happens for me on Drupal 8.9.13, both with Bartik and Bootstrap themes.

Here’s some reproducable code:

namespace Drupalsome_module_nameForm;

use DrupalCoreFormFormBase;
use DrupalCoreFormFormStateInterface;
use DrupalCoreAjaxAjaxResponse;
use DrupalCoreAjaxReplaceCommand;


class TestForm extends FormBase {


    public function getFormId() {
        return 'testform';
    } 


    public function buildForm(array $form, FormStateInterface $form_state, $args = NULL) {

        #If form is shown for the first time, get some default values or previous values from session
        $values = $form_state->getValues();
        if (!$values){

            #Load data from session if exists, otherwise we use some temp test data. 
            #Using persistant data to allow for manual reload of page, in order to demonstrate that data and sorting of rows is saved correctly.
            $names = @$_SESSION('test_data_names');
            if (!$names){
                $values =   (
                    "names"=>(
                        ("first_name"=>"ONE",   "last_name"=>"ONE",),
                        ("first_name"=>"TWO",   "last_name"=>"TWO",),
                        ("first_name"=>"THREE", "last_name"=>"THREE",),
                        ("first_name"=>"FOUR",  "last_name"=>"FOUR",),
                    ),
                );
            }else{
                $values = ("names"=>$names,);
            }

            $form_state->setValues($values);
        }

        #Build the simple table for the form
        $form('names') = (
            "#type"=>"table",
            "#tree"=>true,
            '#empty' => t('No names yet.'),
            '#tabledrag' => ((
                'action' => 'order',
                'relationship' => 'sibling',
                'group' => "group-order-weight",
            )),
            '#header' =>(
                ("data"=>t("First name"),),
                ("data"=>t("Last name(s)"),),
                ("data"=>t("Weight"),),
            ),
        );
        
        #Build table rows for the name inputs
        foreach ($values('names') as $key=>$name_data){
            
            $row = (
                '#attributes'=>(
                    'class' => ('draggable'),
                ),

                "first_name"=>(
                    "#type"             => "textfield",
                    "#default_value"    => $name_data('first_name'),
                ),
                
                "last_name"=>(
                    "#type"             => "textfield",
                    "#default_value"    => $name_data('last_name'),
                ),

                "weight"=>(
                    "#type"             => "weight",
                    "#attributes"=>(
                        'class' => ("group-order-weight"),
                    ),
                ),                
            );
            $form('names')() = $row;
            
        }

        $form('save') = (
            "#type"             => "submit",
            "#value"            => t("Save"),
            '#submit'           => ('::ajax_save_submit'),
            '#ajax'             => (
                'callback'  => '::ajax_save_submit_callback',
                'wrapper'   => 'table-wrapper',
            ),
        );

        return $form;

    }
    

    public function validateForm(array &$form, FormStateInterface $form_state) {

        parent::validateForm($form, $form_state);
        // return $form;

    }
    

    public function submitForm(array &$form, FormStateInterface $form_state) {

    }


    public function ajax_save_submit_callback(array &$form, FormStateInterface $form_state){

        $response = new AjaxResponse();
        $response->addCommand(new ReplaceCommand('#testform', $form));
        return $response;

    }


    public function ajax_save_submit(array &$form, FormStateInterface $form_state){

        #Get the names and save 'em in session so they persist if user reloads page
        $names = $form_state->getValue("names");

        $_SESSION('test_data_names') = $names;
    
        $form_state->setValue("names",$names);
        $form_state->setRebuild(TRUE);

    }

}