afterStateUpdatedJS working differently after upgrade to v5

Group::make()
    ->extraAttributes([
        'x-data' => <<<'JS'
            {
                calculateTotal() {
                    const first = parseFloat($get('first')) || 0;
                    const second = parseFloat($get('second')) || 0;

                    $set('total_amount', first + second);
                }
            }
        JS,
    ])
    ->schema([
        TextInput::make('first')
            ->required()
            ->afterStateUpdatedJs('calculateTotal()');
        TextInput::make('second')
            ->required()
            ->afterStateUpdatedJs('calculateTotal()');

        TextInput::make('total_amount')
            ->readOnly()
            ->required();
    ]),