Fix edit issue in e2e

e2e/playwright/point-click.spec.ts

@@ -1406,7 +1406,7 @@ sketch002 = startSketchOn('XZ')
       })
       await clickOnSketch2()
       await page.waitForTimeout(500)
-      await cmdBar.selectOption({ name: 'False' }).click()
+      await cmdBar.progressCmdBar()
       await cmdBar.expectState({
         commandName: 'Sweep',
         headerArguments: {
@@ -1432,23 +1432,20 @@ sketch002 = startSketchOn('XZ')
       await toolbar.closePane('code')
     })
 
-    // TODO: this test hits a case that happens sometimes in manual testing where
-    // we can't edit the sweep right after its creation. As if the codemod didn't fully
-    // register during the initall add flow.
     await test.step('Edit sweep via feature tree selection works', async () => {
       await toolbar.openPane('feature-tree')
       const operationButton = await toolbar.getFeatureTreeOperation('Sweep', 0)
       await operationButton.dblclick({ button: 'left' })
-      await cmdBar.selectOption({ name: 'True' }).click()
       await cmdBar.expectState({
         commandName: 'Sweep',
+        currentArgKey: 'sectional',
+        currentArgValue: '',
         headerArguments: {
           Sectional: '',
         },
-        stage: 'review',
+        highlightedHeaderArg: 'sectional',
+        stage: 'arguments',
       })
-      const submitButton = page.getByRole('button', { name: 'Submit command' })
-      await submitButton.press('Shift+Backspace')
       await cmdBar.selectOption({ name: 'True' }).click()
       await cmdBar.expectState({
         commandName: 'Sweep',

src/lang/modifyAst.ts

@@ -475,9 +475,24 @@ export function addSweep({
   modifiedAst: Node<Program>
   pathToNode: PathToNode
 } {
+  console.log(
+    'addSweep',
+    JSON.stringify({
+      node,
+      targetDeclarator,
+      trajectoryDeclarator,
+      sectional,
+      variableName,
+      insertIndex,
+    })
+  )
   const modifiedAst = structuredClone(node)
+  console.log('modifiedAst')
   const name =
     variableName ?? findUniqueName(node, KCL_DEFAULT_CONSTANT_PREFIXES.SWEEP)
+  console.log('name', name)
+  console.log('targetDeclarator', JSON.stringify(targetDeclarator))
+  console.log('targetDeclarator.id.name', targetDeclarator.id.name)
   const call = createCallExpressionStdLibKw(
     'sweep',
     createIdentifier(targetDeclarator.id.name),
@@ -486,7 +501,9 @@ export function addSweep({
       createLabeledArg('sectional', createLiteral(sectional)),
     ]
   )
+  console.log('call')
   const variable = createVariableDeclaration(name, call)
+  console.log('variable')
   const insertAt =
     insertIndex !== undefined
       ? insertIndex
@@ -494,6 +511,9 @@ export function addSweep({
       ? modifiedAst.body.length
       : 0
 
+  console.log('insertAt', insertAt)
+  console.log('modifiedAst.body.length', modifiedAst.body.length)
+
   modifiedAst.body.length
     ? modifiedAst.body.splice(insertAt, 0, variable)
     : modifiedAst.body.push(variable)

src/machines/modelingMachine.ts

@@ -70,7 +70,7 @@ import {
 import { ModelingCommandSchema } from 'lib/commandBarConfigs/modelingCommandConfig'
 import { err, reportRejection, trap } from 'lib/trap'
 import { DefaultPlaneStr } from 'lib/planes'
-import { uuidv4 } from 'lib/utils'
+import { isArray, uuidv4 } from 'lib/utils'
 import { Coords2d } from 'lang/std/sketch'
 import { deleteSegment } from 'clientSideScene/ClientSideSceneComp'
 import toast from 'react-hot-toast'
@@ -1891,7 +1891,7 @@ export const modelingMachine = setup({
         let variableName: string | undefined = undefined
         let insertIndex: number | undefined = undefined
 
-        // If this is an edit flow, first we're going to remove the old extrusion
+        // If this is an edit flow, first we're going to remove the old one
         if (nodeToEdit !== undefined && typeof nodeToEdit[1][0] === 'number') {
           // Extract the plane name from the node to edit
           const variableNode = getNodeFromPath<VariableDeclaration>(
@@ -1918,30 +1918,36 @@ export const modelingMachine = setup({
           ast,
           target.graphSelections[0].codeRef.range
         )
-        const targetNode = getNodeFromPath<VariableDeclarator>(
-          ast,
-          targetNodePath,
-          'VariableDeclarator'
-        )
+        // Gotchas, not sure why
+        // - it seems like in some cases we get a list on edit, especially the state that e2e hits
+        // - looking for a VariableDeclaration seems more robust than VariableDeclarator
+        const targetNode = getNodeFromPath<
+          VariableDeclaration | VariableDeclaration[]
+        >(ast, targetNodePath, 'VariableDeclaration')
         if (err(targetNode)) {
           return new Error("Couldn't parse profile selection")
         }
-        const targetDeclarator = targetNode.node
+
+        const targetDeclarator = isArray(targetNode.node)
+          ? targetNode.node[0].declaration
+          : targetNode.node.declaration
 
         // Find the trajectory (or path) declaration
         const trajectoryNodePath = getNodePathFromSourceRange(
           ast,
           trajectory.graphSelections[0].codeRef.range
         )
-        const trajectoryNode = getNodeFromPath<VariableDeclarator>(
+        // Also looking for VariableDeclaration for consistency here
+        const trajectoryNode = getNodeFromPath<VariableDeclaration>(
           ast,
           trajectoryNodePath,
-          'VariableDeclarator'
+          'VariableDeclaration'
         )
         if (err(trajectoryNode)) {
           return new Error("Couldn't parse path selection")
         }
-        const trajectoryDeclarator = trajectoryNode.node
+
+        const trajectoryDeclarator = trajectoryNode.node.declaration
 
         // Perform the sweep
         const addResult = addSweep({
@@ -1966,6 +1972,8 @@ export const modelingMachine = setup({
         if (updatedAst?.selections) {
           editorManager.selectRange(updatedAst?.selections)
         }
+
+        console.log('end of sweepastmod')
       }
     ),
     loftAstMod: fromPromise(