From 08fd9213e4ec34bec76550e503b70fee9cd4bca5 Mon Sep 17 00:00:00 2001
From: Adam Obuchowicz <adam.obuchowicz@enso.org>
Date: Fri, 25 Oct 2024 14:26:21 +0200
Subject: [PATCH] Table Input Widget: Add column with plus (#11388)

Fixes #10863

(The `Column #3` in a screenshot below is just created and actually existing).

![image](https://github.com/user-attachments/assets/995d1a85-eb04-4b8d-bf23-b47ea61185e4)
---
 CHANGELOG.md                                  |   3 +
 app/gui/e2e/project-view/rightPanel.spec.ts   |  11 +-
 .../project-view/tableVisualisation.spec.ts   |   1 +
 app/gui/e2e/project-view/widgets.spec.ts      |  29 ++--
 .../GraphEditor/widgets/WidgetTableEditor.vue |  18 ++-
 .../widgets/WidgetTableEditor/TableHeader.vue |  75 +++++++---
 .../__tests__/tableNewArgument.test.ts        | 132 +++++++++---------
 .../WidgetTableEditor/tableNewArgument.ts     |  64 +++++----
 .../components/TooltipTrigger.vue             |   1 -
 app/gui/src/project-view/providers/index.ts   |   5 +-
 10 files changed, 203 insertions(+), 136 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b397ec7fa7..cbfc042aba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,12 +9,15 @@
 - [Copying and pasting in Table Editor Widget now works properly][11332]
 - [Fix invisible selection in Table Input Widget][11358]
 - [Enable cloud file browser in local projects][11383]
+- [Changed the way of adding new column in Table Input Widget][11388]. The
+  "virtual column" is replaced with an explicit (+) button.
 
 [11151]: https://github.com/enso-org/enso/pull/11151
 [11271]: https://github.com/enso-org/enso/pull/11271
 [11332]: https://github.com/enso-org/enso/pull/11332
 [11358]: https://github.com/enso-org/enso/pull/11358
 [11383]: https://github.com/enso-org/enso/pull/11383
+[11388]: https://github.com/enso-org/enso/pull/11388
 
 #### Enso Standard Library
 
diff --git a/app/gui/e2e/project-view/rightPanel.spec.ts b/app/gui/e2e/project-view/rightPanel.spec.ts
index d8cc94aa46..474231e5ec 100644
--- a/app/gui/e2e/project-view/rightPanel.spec.ts
+++ b/app/gui/e2e/project-view/rightPanel.spec.ts
@@ -26,15 +26,16 @@ test('Doc panel focus (regression #10471)', async ({ page }) => {
   await page.keyboard.press(`${CONTROL_KEY}+D`)
   await page.keyboard.press(`${CONTROL_KEY}+\``)
   await expect(locate.rightDock(page)).toBeVisible()
-  await expect(locate.bottomDock(page)).toBeVisible()
+  const codeEditor = page.locator('.CodeEditor')
+  await expect(codeEditor).toBeVisible()
 
   // Focus code editor.
-  await locate.bottomDock(page).click()
+  await codeEditor.click()
 
   await page.evaluate(() => {
-    const codeEditor = (window as any).__codeEditorApi
-    const docStart = codeEditor.indexOf('The main method')
-    codeEditor.placeCursor(docStart + 8)
+    const codeEditorApi = (window as any).__codeEditorApi
+    const docStart = codeEditorApi.indexOf('The main method')
+    codeEditorApi.placeCursor(docStart + 8)
   })
   await page.keyboard.press('Space')
   await page.keyboard.press('T')
diff --git a/app/gui/e2e/project-view/tableVisualisation.spec.ts b/app/gui/e2e/project-view/tableVisualisation.spec.ts
index 168ff37cc9..73fd06f10f 100644
--- a/app/gui/e2e/project-view/tableVisualisation.spec.ts
+++ b/app/gui/e2e/project-view/tableVisualisation.spec.ts
@@ -58,6 +58,7 @@ test('Copy from Table Visualization', async ({ page, context }) => {
   const node = await actions.createTableNode(page)
   const widget = node.locator('.WidgetTableEditor')
   await expect(widget).toBeVisible()
+  await widget.getByRole('button', { name: 'Add new column' }).click()
   await widget.locator('.ag-cell', { hasNotText: /0/ }).first().click()
   await page.keyboard.press('Control+V')
   await expect(widget.locator('.ag-cell')).toHaveText([
diff --git a/app/gui/e2e/project-view/widgets.spec.ts b/app/gui/e2e/project-view/widgets.spec.ts
index 8c33e473f7..21a1b2697e 100644
--- a/app/gui/e2e/project-view/widgets.spec.ts
+++ b/app/gui/e2e/project-view/widgets.spec.ts
@@ -544,27 +544,32 @@ test('Table widget', async ({ page }) => {
   const node = await actions.createTableNode(page)
   const widget = node.locator('.WidgetTableEditor')
   await expect(widget).toBeVisible()
-  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'New Column'])
-  await expect(widget.locator('.ag-header-cell-text', { hasText: 'New Column' })).toHaveClass(
-    /(?<=^| )virtualColumn(?=$| )/,
-  )
-  // There are two cells, one with row number, second allowing creating first row and column
-  await expect(widget.locator('.ag-cell')).toHaveCount(2)
+  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#'])
+  await expect(widget.getByRole('button', { name: 'Add new column' })).toExist()
+  await expect(widget.locator('.ag-cell')).toHaveText(['0', ''])
+
+  // Create first column
+  await widget.getByRole('button', { name: 'Add new column' }).click()
+  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'Column #1'])
+  await expect(widget.locator('.ag-cell')).toHaveText(['0', '', ''])
 
   // Putting first value
-  await widget.locator('.ag-cell', { hasNotText: '0' }).click()
+  await widget.locator('.ag-cell', { hasNotText: '0' }).first().click()
   await page.keyboard.type('Value')
   await page.keyboard.press('Enter')
-  // There will be new blank column and new blank row allowing adding new columns and rows
-  // (so 4 cells in total)
-  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'Column #0', 'New Column'])
+  // There will be new blank row allowing adding new rows.
   await expect(widget.locator('.ag-cell')).toHaveText(['0', 'Value', '', '1', '', ''])
 
   // Renaming column
-  await widget.locator('.ag-header-cell-text', { hasText: 'Column #0' }).first().click()
+  await widget.locator('.ag-header-cell-text', { hasText: 'Column #1' }).first().click()
   await page.keyboard.type('Header')
   await page.keyboard.press('Enter')
-  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'Header', 'New Column'])
+  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'Header'])
+
+  // Adding next column
+  await widget.getByRole('button', { name: 'Add new column' }).click()
+  await expect(widget.locator('.ag-header-cell-text')).toHaveText(['#', 'Header', 'Column #2'])
+  await expect(widget.locator('.ag-cell')).toHaveText(['0', 'Value', '', '', '1', '', '', ''])
 
   // Switching edit between cells and headers - check we will never edit two things at once.
   await expect(widget.locator('.ag-text-field-input')).toHaveCount(0)
diff --git a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor.vue b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor.vue
index 78dcedc52c..da5a17f3f2 100644
--- a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor.vue
+++ b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor.vue
@@ -9,6 +9,7 @@ import {
 import ResizeHandles from '@/components/ResizeHandles.vue'
 import AgGridTableView from '@/components/shared/AgGridTableView.vue'
 import { injectGraphNavigator } from '@/providers/graphNavigator'
+import { useTooltipRegistry } from '@/providers/tooltipState'
 import { Score, defineWidget, widgetProps } from '@/providers/widgetRegistry'
 import { WidgetEditHandler } from '@/providers/widgetRegistry/editHandler'
 import { useGraphStore } from '@/stores/graph'
@@ -20,12 +21,13 @@ import '@ag-grid-community/styles/ag-theme-alpine.css'
 import type {
   CellEditingStartedEvent,
   CellEditingStoppedEvent,
+  ColDef,
   Column,
   ColumnMovedEvent,
   ProcessDataFromClipboardParams,
   RowDragEndEvent,
 } from 'ag-grid-enterprise'
-import { computed, ref } from 'vue'
+import { computed, markRaw, ref } from 'vue'
 import type { ComponentExposed } from 'vue-component-type-helpers'
 
 const props = defineProps(widgetProps(widgetDefinition))
@@ -178,14 +180,22 @@ function processDataFromClipboard({ data, api }: ProcessDataFromClipboardParams<
 
 // === Column Default Definition ===
 
-const defaultColDef = {
+const tooltipRegistry = useTooltipRegistry()
+const defaultColDef: ColDef<RowData> = {
   editable: true,
   resizable: true,
   sortable: false,
   lockPinned: true,
+  menuTabs: ['generalMenuTab'],
   headerComponentParams: {
-    onHeaderEditingStarted: headerEditHandler.headerEditedInGrid.bind(headerEditHandler),
-    onHeaderEditingStopped: headerEditHandler.headerEditingStoppedInGrid.bind(headerEditHandler),
+    // TODO[ao]: we mark raw, because otherwise any change _inside_ tooltipRegistry causes the grid
+    //  to be refreshed. Technically, shallowReactive should work here, but it does not,
+    //  I don't know why
+    tooltipRegistry: markRaw(tooltipRegistry),
+    editHandlers: {
+      onHeaderEditingStarted: headerEditHandler.headerEditedInGrid.bind(headerEditHandler),
+      onHeaderEditingStopped: headerEditHandler.headerEditingStoppedInGrid.bind(headerEditHandler),
+    },
   },
 }
 </script>
diff --git a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/TableHeader.vue b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/TableHeader.vue
index 8ffe338fef..710a48b8d9 100644
--- a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/TableHeader.vue
+++ b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/TableHeader.vue
@@ -1,22 +1,39 @@
 <script lang="ts">
+import SvgButton from '@/components/SvgButton.vue'
+import { provideTooltipRegistry, TooltipRegistry } from '@/providers/tooltipState'
 import type { IHeaderParams } from 'ag-grid-community'
-import { ref, watch } from 'vue'
+import { computed, ref, watch } from 'vue'
+
+export interface HeaderEditHandlers {
+  /** Setter called when column name is changed by the user. */
+  nameSetter: (newName: string) => void
+  onHeaderEditingStarted?: (stop: (cancel: boolean) => void) => void
+  onHeaderEditingStopped?: () => void
+}
+
+/**
+ * A subset of {@link HeaderParams} which is meant to be specified for columns separately
+ * (not in defaultColumnDef).
+ */
+export type ColumnSpecificHeaderParams =
+  | {
+      type: 'astColumn'
+      editHandlers: HeaderEditHandlers
+    }
+  | { type: 'newColumn'; newColumnRequested: () => void }
+  | { type: 'rowIndexColumn' }
 
 /**
  * Parameters recognized by this header component.
  *
  * They are set through `headerComponentParams` option in AGGrid column definition.
  */
-export interface HeaderParams {
-  /** Setter called when column name is changed by the user. */
-  nameSetter?: (newName: string) => void
+export type HeaderParams = ColumnSpecificHeaderParams & {
   /**
-   * Column is virtual if it is not represented in the AST. Such column might be used
-   * to create new one.
+   * AgGrid mounts header components as separate "App", so we don't have access to any context.
+   * Threfore the tooltip registry must be provided by props.
    */
-  virtualColumn?: boolean
-  onHeaderEditingStarted?: (stop: (cancel: boolean) => void) => void
-  onHeaderEditingStopped?: () => void
+  tooltipRegistry: TooltipRegistry
 }
 </script>
 
@@ -25,17 +42,23 @@ const props = defineProps<{
   params: IHeaderParams & HeaderParams
 }>()
 
+/** Re-provide tooltipRegistry. See `tooltipRegistry` docs in {@link HeaderParams} */
+provideTooltipRegistry.provideConstructed(props.params.tooltipRegistry)
+
 const editing = ref(false)
 const inputElement = ref<HTMLInputElement>()
+const editHandlers = computed(() =>
+  props.params.type === 'astColumn' ? props.params.editHandlers : undefined,
+)
 
 watch(editing, (newVal) => {
   if (newVal) {
-    props.params.onHeaderEditingStarted?.((cancel: boolean) => {
+    editHandlers.value?.onHeaderEditingStarted?.((cancel: boolean) => {
       if (cancel) editing.value = false
       else acceptNewName()
     })
   } else {
-    props.params.onHeaderEditingStopped?.()
+    editHandlers.value?.onHeaderEditingStopped?.()
   }
 })
 
@@ -48,33 +71,47 @@ watch(inputElement, (newVal, oldVal) => {
 })
 
 function acceptNewName() {
+  if (editHandlers.value == null) {
+    console.error("Tried to accept header new name where it's not editable!")
+    return
+  }
   if (inputElement.value == null) {
     console.error('Tried to accept header new name without input element!')
     return
   }
-  props.params.nameSetter?.(inputElement.value.value)
+  editHandlers.value.nameSetter(inputElement.value.value)
   editing.value = false
 }
 
-function onMouseClick() {
-  if (!editing.value && props.params.nameSetter != null) {
+function onMouseClick(event: MouseEvent) {
+  if (!editing.value && props.params.type === 'astColumn') {
     editing.value = true
+    event.stopPropagation()
   }
 }
 
 function onMouseRightClick(event: MouseEvent) {
   if (!editing.value) {
     props.params.showColumnMenuAfterMouseClick(event)
+    event.preventDefault()
+    event.stopPropagation()
   }
 }
 </script>
 
 <template>
+  <SvgButton
+    v-if="params.type === 'newColumn'"
+    class="addColumnButton"
+    name="add"
+    title="Add new column"
+    @click.stop="params.newColumnRequested()"
+  />
   <div
+    v-else
     class="ag-cell-label-container"
     role="presentation"
     @pointerdown.stop
-    @click.stop
     @click="onMouseClick"
     @click.right="onMouseRightClick"
   >
@@ -93,14 +130,18 @@ function onMouseRightClick(event: MouseEvent) {
       <span
         v-else
         class="ag-header-cell-text"
-        :class="{ virtualColumn: params.virtualColumn === true }"
+        :class="{ virtualColumn: params.type !== 'astColumn' }"
         >{{ params.displayName }}</span
       >
     </div>
   </div>
 </template>
 
-<style>
+<style scoped>
+.addColumnButton {
+  margin-left: 10px;
+}
+
 .virtualColumn {
   color: rgba(0, 0, 0, 0.5);
 }
diff --git a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/__tests__/tableNewArgument.test.ts b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/__tests__/tableNewArgument.test.ts
index 2f933ec9c0..9a77f79477 100644
--- a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/__tests__/tableNewArgument.test.ts
+++ b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/__tests__/tableNewArgument.test.ts
@@ -1,4 +1,7 @@
 import {
+  DEFAULT_COLUMN_PREFIX,
+  NEW_COLUMN_ID,
+  ROW_INDEX_HEADER,
   RowData,
   tableNewCallMayBeHandled,
   useTableNewArgument,
@@ -18,58 +21,61 @@ function suggestionDbWithNothing() {
   return db
 }
 
+const expectedRowIndexColumnDef = { headerName: ROW_INDEX_HEADER }
+const expectedNewColumnDef = { cellStyle: { display: 'none' } }
+
 test.each([
   {
     code: 'Table.new [["a", [1, 2, 3]], ["b", [4, 5, "six"]], ["empty", [Nothing, Standard.Base.Nothing, Nothing]]]',
     expectedColumnDefs: [
-      { headerName: '#' },
+      expectedRowIndexColumnDef,
       { headerName: 'a' },
       { headerName: 'b' },
       { headerName: 'empty' },
-      { headerName: 'New Column' },
+      expectedNewColumnDef,
     ],
     expectedRows: [
-      { '#': 0, a: 1, b: 4, empty: null, 'New Column': null },
-      { '#': 1, a: 2, b: 5, empty: null, 'New Column': null },
-      { '#': 2, a: 3, b: 'six', empty: null, 'New Column': null },
-      { '#': 3, a: null, b: null, empty: null, 'New Column': null },
+      { [ROW_INDEX_HEADER]: 0, a: 1, b: 4, empty: null, '': null },
+      { [ROW_INDEX_HEADER]: 1, a: 2, b: 5, empty: null, '': null },
+      { [ROW_INDEX_HEADER]: 2, a: 3, b: 'six', empty: null, '': null },
+      { [ROW_INDEX_HEADER]: 3, a: null, b: null, empty: null, '': null },
     ],
   },
   {
     code: 'Table.new []',
-    expectedColumnDefs: [{ headerName: '#' }, { headerName: 'New Column' }],
-    expectedRows: [{ '#': 0, 'New Column': null }],
+    expectedColumnDefs: [expectedRowIndexColumnDef, expectedNewColumnDef],
+    expectedRows: [{ [ROW_INDEX_HEADER]: 0, '': null }],
   },
   {
     code: 'Table.new',
-    expectedColumnDefs: [{ headerName: '#' }, { headerName: 'New Column' }],
-    expectedRows: [{ '#': 0, 'New Column': null }],
+    expectedColumnDefs: [expectedRowIndexColumnDef, expectedNewColumnDef],
+    expectedRows: [{ [ROW_INDEX_HEADER]: 0, '': null }],
   },
   {
     code: 'Table.new _',
-    expectedColumnDefs: [{ headerName: '#' }, { headerName: 'New Column' }],
-    expectedRows: [{ '#': 0, 'New Column': null }],
+    expectedColumnDefs: [expectedRowIndexColumnDef, expectedNewColumnDef],
+    expectedRows: [{ [ROW_INDEX_HEADER]: 0, '': null }],
   },
   {
     code: 'Table.new [["a", []]]',
-    expectedColumnDefs: [{ headerName: '#' }, { headerName: 'a' }, { headerName: 'New Column' }],
-    expectedRows: [{ '#': 0, a: null, 'New Column': null }],
+    expectedColumnDefs: [expectedRowIndexColumnDef, { headerName: 'a' }, expectedNewColumnDef],
+    expectedRows: [{ [ROW_INDEX_HEADER]: 0, a: null, '': null }],
   },
   {
     code: 'Table.new [["a", [1,,2]], ["b", [3, 4,]], ["c", [, 5, 6]], ["d", [,,]]]',
     expectedColumnDefs: [
-      { headerName: '#' },
+      expectedRowIndexColumnDef,
       { headerName: 'a' },
       { headerName: 'b' },
       { headerName: 'c' },
       { headerName: 'd' },
-      { headerName: 'New Column' },
+      expectedNewColumnDef,
     ],
     expectedRows: [
-      { '#': 0, a: 1, b: 3, c: null, d: null, 'New Column': null },
-      { '#': 1, a: null, b: 4, c: 5, d: null, 'New Column': null },
-      { '#': 2, a: 2, b: null, c: 6, d: null, 'New Column': null },
-      { '#': 3, a: null, b: null, c: null, d: null, 'New Column': null },
+      { [ROW_INDEX_HEADER]: 0, a: 1, b: 3, c: null, d: null, '': null },
+      { [ROW_INDEX_HEADER]: 1, a: null, b: 4, c: 5, d: null, '': null },
+      { [ROW_INDEX_HEADER]: 2, a: 2, b: null, c: 6, d: null, '': null },
+      { [ROW_INDEX_HEADER]: 3, a: null, b: null, c: null, d: null, '': null },
     ],
   },
 ])('Read table from $code', ({ code, expectedColumnDefs, expectedRows }) => {
@@ -172,32 +178,7 @@ test.each([
     expected: "Table.new [['a', [1, 2, 3]], ['b', [4, Nothing, 6]]]",
     importExpected: true,
   },
-  {
-    code: "Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]]]",
-    description: 'Add new column',
-    edit: { column: 3, row: 1, value: 8 },
-    expected:
-      "Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]], ['Column #2', [Nothing, 8, Nothing]]]",
-    importExpected: true,
-  },
-  {
-    code: 'Table.new []',
-    description: 'Add first column',
-    edit: { column: 1, row: 0, value: 8 },
-    expected: "Table.new [['Column #0', [8]]]",
-  },
-  {
-    code: 'Table.new',
-    description: 'Add parameter',
-    edit: { column: 1, row: 0, value: 8 },
-    expected: "Table.new [['Column #0', [8]]]",
-  },
-  {
-    code: 'Table.new _',
-    description: 'Update parameter',
-    edit: { column: 1, row: 0, value: 8 },
-    expected: "Table.new [['Column #0', [8]]]",
-  },
+
   {
     code: "Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]]]",
     description: 'Add new row',
@@ -212,14 +193,6 @@ test.each([
     expected: "Table.new [['a', [Nothing]], ['b', ['val']]]",
     importExpected: true,
   },
-  {
-    code: "Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]]]",
-    description: 'Add new row and column (the cell in the corner)',
-    edit: { column: 3, row: 3, value: 7 },
-    expected:
-      "Table.new [['a', [1, 2, 3, Nothing]], ['b', [4, 5, 6, Nothing]], ['Column #2', [Nothing, Nothing, Nothing, 7]]]",
-    importExpected: true,
-  },
   {
     code: "Table.new [['a', [1, ,3]]]",
     description: 'Set missing value',
@@ -257,6 +230,38 @@ test.each([
   else expect(addMissingImports).not.toHaveBeenCalled()
 })
 
+test.each([
+  {
+    code: "Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]]]",
+    expected: `Table.new [['a', [1, 2, 3]], ['b', [4, 5, 6]], ['${DEFAULT_COLUMN_PREFIX}3', [Nothing, Nothing, Nothing]]]`,
+    importExpected: true,
+  },
+  {
+    code: 'Table.new []',
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', []]]`,
+  },
+  {
+    code: 'Table.new',
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', []]]`,
+  },
+  {
+    code: 'Table.new _',
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', []]]`,
+  },
+])('Add column to table $code', ({ code, expected, importExpected }) => {
+  const { tableNewArgs, onUpdate, addMissingImports } = tableEditFixture(code, expected)
+  const newColumnDef = tableNewArgs.columnDefs.value.find(
+    (colDef) => colDef.colId === NEW_COLUMN_ID,
+  )
+  assert(newColumnDef != null)
+  assert(newColumnDef.headerComponentParams?.type === 'newColumn')
+  assert(newColumnDef.headerComponentParams.newColumnRequested != null)
+  newColumnDef.headerComponentParams.newColumnRequested()
+  expect(onUpdate).toHaveBeenCalledOnce()
+  if (importExpected) expect(addMissingImports).toHaveBeenCalled()
+  else expect(addMissingImports).not.toHaveBeenCalled()
+})
+
 function getCustomMenuItemByName(
   name: string,
   items:
@@ -405,7 +410,7 @@ test.each([
       ['1', '3'],
       ['2', '4'],
     ],
-    expected: "Table.new [['Column #0', [1, 2]], ['Column #1', [3, 4]]]",
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', [1, 2]], ['${DEFAULT_COLUMN_PREFIX}2', [3, 4]]]`,
   },
   {
     code: 'Table.new []',
@@ -414,13 +419,13 @@ test.each([
       ['1', '3'],
       ['2', '4'],
     ],
-    expected: "Table.new [['Column #0', [1, 2]], ['Column #1', [3, 4]]]",
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', [1, 2]], ['${DEFAULT_COLUMN_PREFIX}2', [3, 4]]]`,
   },
   {
     code: 'Table.new []',
     focused: { rowIndex: 0, colIndex: 1 },
     data: [['a single cell']],
-    expected: "Table.new [['Column #0', ['a single cell']]]",
+    expected: `Table.new [['${DEFAULT_COLUMN_PREFIX}1', ['a single cell']]]`,
   },
   {
     code: "Table.new [['a', [1, 2]], ['b', [3, 4]]]",
@@ -445,7 +450,7 @@ test.each([
     code: "Table.new [['a', [1, 2]], ['b', [3, 4]]]",
     focused: { rowIndex: 1, colIndex: 3 },
     data: [['a single cell']],
-    expected: "Table.new [['a', [1, 2]], ['b', [3, 4]], ['Column #2', [Nothing, 'a single cell']]]",
+    expected: `Table.new [['a', [1, 2]], ['b', [3, 4]], ['${DEFAULT_COLUMN_PREFIX}3', [Nothing, 'a single cell']]]`,
     importExpected: true,
   },
   {
@@ -473,7 +478,7 @@ test.each([
       ['5', '7'],
       ['6', '8'],
     ],
-    expected: "Table.new [['a', [1, 2]], ['b', [5, 6]], ['Column #2', [7, 8]]]",
+    expected: `Table.new [['a', [1, 2]], ['b', [5, 6]], ['${DEFAULT_COLUMN_PREFIX}3', [7, 8]]]`,
   },
   {
     code: "Table.new [['a', [1, 2]], ['b', [3, 4]]]",
@@ -482,19 +487,17 @@ test.each([
       ['5', '7'],
       ['6', '8'],
     ],
-    expected:
-      "Table.new [['a', [1, 2, Nothing]], ['b', [3, 5, 6]], ['Column #2', [Nothing, 7, 8]]]",
+    expected: `Table.new [['a', [1, 2, Nothing]], ['b', [3, 5, 6]], ['${DEFAULT_COLUMN_PREFIX}3', [Nothing, 7, 8]]]`,
     importExpected: true,
   },
   {
     code: "Table.new [['a', [1, 2]], ['b', [3, 4]]]",
-    focused: { rowIndex: 2, colIndex: 3 },
+    focused: { rowIndex: 2, colIndex: 2 },
     data: [
       ['5', '7'],
       ['6', '8'],
     ],
-    expected:
-      "Table.new [['a', [1, 2, Nothing, Nothing]], ['b', [3, 4, Nothing, Nothing]], ['Column #2', [Nothing, Nothing, 5, 6]], ['Column #3', [Nothing, Nothing, 7, 8]]]",
+    expected: `Table.new [['a', [1, 2, Nothing, Nothing]], ['b', [3, 4, 5, 6]], ['${DEFAULT_COLUMN_PREFIX}3', [Nothing, Nothing, 7, 8]]]`,
     importExpected: true,
   },
 ])(
@@ -502,7 +505,6 @@ test.each([
   ({ code, focused, data, expected, importExpected }) => {
     const { tableNewArgs, onUpdate, addMissingImports } = tableEditFixture(code, expected)
     const focusedCol = tableNewArgs.columnDefs.value[focused.colIndex]
-    console.log(focusedCol?.colId, focusedCol?.headerName)
     assert(focusedCol?.colId != null)
     tableNewArgs.pasteFromClipboard(data, {
       rowIndex: focused.rowIndex,
diff --git a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/tableNewArgument.ts b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/tableNewArgument.ts
index fbf8015c42..0d6b3f1ba8 100644
--- a/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/tableNewArgument.ts
+++ b/app/gui/src/project-view/components/GraphEditor/widgets/WidgetTableEditor/tableNewArgument.ts
@@ -12,12 +12,15 @@ import { qnLastSegment, type QualifiedName } from '@/util/qualifiedName'
 import type { ToValue } from '@/util/reactivity'
 import type { ColDef } from 'ag-grid-enterprise'
 import { computed, toValue } from 'vue'
+import { ColumnSpecificHeaderParams } from './TableHeader.vue'
 
-const NEW_COLUMN_ID = 'NewColumn'
+/** Id of a fake column with "Add new column" option. */
+export const NEW_COLUMN_ID = 'NewColumn'
 const ROW_INDEX_COLUMN_ID = 'RowIndex'
-const NEW_COLUMN_HEADER = 'New Column'
-const ROW_INDEX_HEADER = '#'
-const DEFAULT_COLUMN_PREFIX = 'Column #'
+/** A header of Row Index Column. */
+export const ROW_INDEX_HEADER = '#'
+/** A default prefix added to the column's index in newly created columns. */
+export const DEFAULT_COLUMN_PREFIX = 'Column #'
 const NOTHING_PATH = 'Standard.Base.Nothing.Nothing' as QualifiedName
 const NOTHING_NAME = qnLastSegment(NOTHING_PATH)
 
@@ -37,6 +40,7 @@ export interface ColumnDef extends ColDef<RowData> {
   mainMenuItems: (string | MenuItem<RowData>)[]
   contextMenuItems: (string | MenuItem<RowData>)[]
   rowDrag?: ({ data }: { data: RowData | undefined }) => boolean
+  headerComponentParams?: ColumnSpecificHeaderParams
 }
 
 namespace cellValueConversion {
@@ -260,34 +264,26 @@ export function useTableNewArgument(
 
   const newColumnDef = computed<ColumnDef>(() => ({
     colId: NEW_COLUMN_ID,
-    headerName: NEW_COLUMN_HEADER,
+    headerName: '',
     valueGetter: () => null,
-    valueSetter: ({ data, newValue }: { data: RowData; newValue: any }) => {
-      const edit = graph.startEdit()
-      if (data.index === rowCount.value) {
-        addRow(edit)
-      }
-      addColumn(
-        edit,
-        `${DEFAULT_COLUMN_PREFIX}${columns.value.length}`,
-        (index) => (index === data.index ? newValue : null),
-        Math.max(rowCount.value, data.index + 1),
-      )
-      onUpdate({ edit })
-      return true
-    },
+    editable: false,
+    resizable: false,
+    suppressNavigable: true,
+    width: 40,
+    maxWidth: 40,
     headerComponentParams: {
-      nameSetter: (newName: string) => {
+      type: 'newColumn',
+      newColumnRequested: () => {
         const edit = graph.startEdit()
         fixColumns(edit)
-        addColumn(edit, newName)
+        addColumn(edit, `${DEFAULT_COLUMN_PREFIX}${columns.value.length + 1}`)
         onUpdate({ edit })
       },
-      virtualColumn: true,
     },
     mainMenuItems: ['autoSizeThis', 'autoSizeAll'],
-    contextMenuItems: [commonContextMenuActions.paste, 'separator', removeRowMenuItem],
+    contextMenuItems: [removeRowMenuItem],
     lockPosition: 'right',
+    cellStyle: { display: 'none' },
   }))
 
   const rowIndexColumnDef = computed<ColumnDef>(() => ({
@@ -295,6 +291,11 @@ export function useTableNewArgument(
     headerName: ROW_INDEX_HEADER,
     valueGetter: ({ data }: { data: RowData | undefined }) => data?.index,
     editable: false,
+    resizable: false,
+    suppressNavigable: true,
+    headerComponentParams: {
+      type: 'rowIndexColumn',
+    },
     mainMenuItems: ['autoSizeThis', 'autoSizeAll'],
     contextMenuItems: [removeRowMenuItem],
     cellStyle: { color: 'rgba(0, 0, 0, 0.4)' },
@@ -338,11 +339,14 @@ export function useTableNewArgument(
             return true
           },
           headerComponentParams: {
-            nameSetter: (newName: string) => {
-              const edit = graph.startEdit()
-              fixColumns(edit)
-              edit.getVersion(col.name).setRawTextContent(newName)
-              onUpdate({ edit })
+            type: 'astColumn',
+            editHandlers: {
+              nameSetter: (newName: string) => {
+                const edit = graph.startEdit()
+                fixColumns(edit)
+                edit.getVersion(col.name).setRawTextContent(newName)
+                onUpdate({ edit })
+              },
             },
           },
           mainMenuItems: ['autoSizeThis', 'autoSizeAll', removeColumnMenuItem(col.id)],
@@ -357,7 +361,7 @@ export function useTableNewArgument(
             'separator',
             'export',
           ],
-        }) satisfies ColDef<RowData>,
+        }) satisfies ColumnDef,
     )
     cols.unshift(rowIndexColumnDef.value)
     cols.push(newColumnDef.value)
@@ -470,7 +474,7 @@ export function useTableNewArgument(
     for (let i = columns.value.length; i < newColCount; ++i) {
       modifiedColumnsAst = addColumn(
         edit,
-        `${DEFAULT_COLUMN_PREFIX}${i}`,
+        `${DEFAULT_COLUMN_PREFIX}${i + 1}`,
         (index) => newValueGetter(index, i),
         newRowCount,
         modifiedColumnsAst,
diff --git a/app/gui/src/project-view/components/TooltipTrigger.vue b/app/gui/src/project-view/components/TooltipTrigger.vue
index b4e00c30ef..7ba6561a15 100644
--- a/app/gui/src/project-view/components/TooltipTrigger.vue
+++ b/app/gui/src/project-view/components/TooltipTrigger.vue
@@ -6,7 +6,6 @@ import { toRef, useSlots } from 'vue'
 usePropagateScopesToAllRoots()
 
 const registry = useTooltipRegistry()
-
 const slots = useSlots()
 const tooltipSlot = toRef(slots, 'tooltip')
 const registered = registry.registerTooltip(tooltipSlot)
diff --git a/app/gui/src/project-view/providers/index.ts b/app/gui/src/project-view/providers/index.ts
index db8e23d42b..728b578ab5 100644
--- a/app/gui/src/project-view/providers/index.ts
+++ b/app/gui/src/project-view/providers/index.ts
@@ -46,9 +46,10 @@ export function createContextStore<F extends (...args: any[]) => any>(name: stri
 
   /**
    * An method allowing to directly provide a store value to any app or component context, possibly
-   * skipping invoking the factory function. **Do not use in application code.**
+   * skipping invoking the factory function. **Use in tests only, or in exceptional situations
+   * (like in TableHeader.vue)**.
    */
-  provideFn._mock = function (
+  provideFn.provideConstructed = function (
     valueOrArgs: Parameters<F> | ReturnType<F>,
     app?: App,
   ): ReturnType<F> {