How do you prevent a re-created React component from losing state?

The example in this sandbox is a contrived example, but it illustrates the point.

Clicking the "Add column" button is supposed to add a new column. It works the first time, but doesn’t work after that. You’ll notice from the log that this issue has to do with the fact that columns is always in its original state. Therefore, a column is always being added to this original state, not the current state.

I imagine this issue is related to the fact that the column header is being re-created on each call to renderHeader, but I’m unsure about how to pass the state to the newly created header component.

>Solution :

There is 2 ways to update your state:

• The first way:
  setColumns(newColumns)
• The second way:
  setColumns(columns => newColumns)

You should use the second way because the new state depend on the current state.

onClick={() => {
  setColumns((columns) => {
    console.log(columns.map((column) => column.field));
    const newColumnName = `Column ${columns.length}`;
    const newColumns = [...columns];
    newColumns.splice(
      columns.length - 1,
      0,
      createColumn(newColumnName)
    );
    return newColumns;
  });
}}

Your component:

import React, { useState } from "react";
import { DataGrid } from "@mui/x-data-grid";
import Button from "@mui/material/Button";
import "./styles.css";

export default function App() {
  const [rows] = useState([
    { id: 1, "Column 1": 1, "Column 2": 2 },
    { id: 2, "Column 1": 3, "Column 2": 3 },
    { id: 3, "Column 1": 4, "Column 2": 5 }
  ]);
  const createColumn = (name) => {
    return {
      field: name,
      align: "center",
      editable: true,
      sortable: false
    };
  };
  const [columns, setColumns] = useState([
    createColumn("Column 1"),
    createColumn("Column 2"),
    {
      field: "Add a split",
      width: 150,
      sortable: false,
      renderHeader: (params) => {
        return (
          <Button
            variant="contained"
            onClick={() => {
              setColumns((columns) => { // <== use callback
                console.log(columns.map((column) => column.field));
                const newColumnName = `Column ${columns.length}`;
                const newColumns = [...columns];
                newColumns.splice(
                  columns.length - 1,
                  0,
                  createColumn(newColumnName)
                );
                return newColumns;
              });
            }}
          >
            Add column
          </Button>
        );
      }
    }
  ]);

  return (
    <div className="App">
      <DataGrid
        className="App-data-grid"
        rows={rows}
        columns={columns}
        disableSelectionOnClick
        disableColumnMenu
      />
    </div>
  );
}

https://codesandbox.io/s/mui-sandbox-forked-1im0p?file=/src/App.js:0-1562