Add edge and segment selection in point-and-click Helix flow (#5866)

* WIP: Add edge and segment selection in point-and-click Helix flow
Fixes #5393

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* Working edge based helix edit

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* Add utility function for shared code between revolve and helix

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* Use updateModelingState in codemod

* A snapshot a day keeps the bugs away! 📷🐛

* Add skip: true for edge helix to be consistent with axis as options

* A snapshot a day keeps the bugs away! 📷🐛

* Add support for sweepEdge and tests

* Lint

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

* Clean up snapshots

* Fix up tests after ccw change

* Use displayName: 'CounterClockWise' cause ccw not cutting it

* Fix tsc

* Update 2020 snapshots after helix change

* Update 2020 snapshots after helix change

* A snapshot a day keeps the bugs away! 📷🐛

* Another one :djkhaled:

* Clean up

* A snapshot a day keeps the bugs away! 📷🐛

* A snapshot a day keeps the bugs away! 📷🐛

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

e2e/playwright/point-click.spec.ts

@@ -1071,7 +1071,7 @@ openSketch = startSketchOn(XY)
     })
   })
 
-  test('Helix point-and-click', async ({
+  test('Helix point-and-click on default axis', async ({
     context,
     page,
     homePage,
@@ -1087,24 +1087,21 @@ openSketch = startSketchOn(XY)
 
     await homePage.goToModelingScene()
 
-    // await test.step(`Look for the red of the default plane`, async () => {
-    //   await scene.expectPixelColor([96, 52, 52], testPoint, 15)
-    // })
     await test.step(`Go through the command bar flow`, async () => {
       await toolbar.helixButton.click()
       await cmdBar.expectState({
         stage: 'arguments',
-        currentArgKey: 'revolutions',
+        currentArgKey: 'axisOrEdge',
-        currentArgValue: '1',
+        currentArgValue: '',
         headerArguments: {
           AngleStart: '',
-          Axis: '',
+          AxisOrEdge: '',
-          Ccw: '',
+          CounterClockWise: '',
           Length: '',
           Radius: '',
           Revolutions: '',
         },
-        highlightedHeaderArg: 'revolutions',
+        highlightedHeaderArg: 'axisOrEdge',
         commandName: 'Helix',
       })
       await cmdBar.progressCmdBar()
@@ -1114,6 +1111,7 @@ openSketch = startSketchOn(XY)
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
     })
 
     await test.step(`Confirm code is added to the editor, scene has changed`, async () => {
@@ -1141,7 +1139,7 @@ openSketch = startSketchOn(XY)
         headerArguments: {
           AngleStart: '360',
           Axis: 'X',
-          Ccw: '',
+          CounterClockWise: '',
           Length: initialInput,
           Radius: '5',
           Revolutions: '1',
@@ -1156,7 +1154,7 @@ openSketch = startSketchOn(XY)
         headerArguments: {
           AngleStart: '360',
           Axis: 'X',
-          Ccw: '',
+          CounterClockWise: '',
           Length: newInput,
           Radius: '5',
           Revolutions: '1',
@@ -1179,6 +1177,165 @@ openSketch = startSketchOn(XY)
     })
   })
 
+  const helixCases = [
+    {
+      selectionType: 'segment',
+      testPoint: { x: 513, y: 221 },
+      expectedOutput: `helix001 = helix(  revolutions = 20,  angleStart = 0,  ccw = true,  radius = 1,  axis = seg01,  length = 100,)`,
+      expectedEditedOutput: `helix001 = helix(  revolutions = 20,  angleStart = 0,  ccw = true,  radius = 1,  axis = seg01,  length = 50,)`,
+    },
+    {
+      selectionType: 'sweepEdge',
+      testPoint: { x: 564, y: 364 },
+      expectedOutput: `helix001 = helix(  revolutions = 20,  angleStart = 0,  ccw = true,  radius = 1,  axis =   getOppositeEdge(seg01),  length = 100,)`,
+      expectedEditedOutput: `helix001 = helix(  revolutions = 20,  angleStart = 0,  ccw = true,  radius = 1,  axis =   getOppositeEdge(seg01),  length = 50,)`,
+    },
+  ]
+  helixCases.map(
+    ({ selectionType, testPoint, expectedOutput, expectedEditedOutput }) => {
+      test(`Helix point-and-click around ${selectionType}`, async ({
+        context,
+        page,
+        homePage,
+        scene,
+        editor,
+        toolbar,
+        cmdBar,
+      }) => {
+        page.on('console', console.log)
+        const initialCode = `sketch001 = startSketchOn('XZ')
+  profile001 = startProfileAt([0, 0], sketch001)
+    |> yLine(length = 100)
+    |> line(endAbsolute = [100, 0])
+    |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+    |> close()
+  extrude001 = extrude(profile001, length = 100)`
+
+        // One dumb hardcoded screen pixel value
+        const [clickOnEdge] = scene.makeMouseHelpers(testPoint.x, testPoint.y)
+
+        await context.addInitScript((initialCode) => {
+          localStorage.setItem('persistCode', initialCode)
+        }, initialCode)
+        await page.setBodyDimensions({ width: 1000, height: 500 })
+        await homePage.goToModelingScene()
+
+        await test.step(`Go through the command bar flow`, async () => {
+          await toolbar.closePane('code')
+          await toolbar.helixButton.click()
+          await cmdBar.expectState({
+            stage: 'arguments',
+            currentArgKey: 'axisOrEdge',
+            currentArgValue: '',
+            headerArguments: {
+              AngleStart: '',
+              AxisOrEdge: '',
+              CounterClockWise: '',
+              Length: '',
+              Radius: '',
+              Revolutions: '',
+            },
+            highlightedHeaderArg: 'axisOrEdge',
+            commandName: 'Helix',
+          })
+          await cmdBar.selectOption({ name: 'Edge' }).click()
+          await clickOnEdge()
+          await cmdBar.progressCmdBar()
+          await cmdBar.argumentInput.focus()
+          await page.keyboard.insertText('20')
+          await cmdBar.progressCmdBar()
+          await page.keyboard.insertText('0')
+          await cmdBar.progressCmdBar()
+          await cmdBar.selectOption({ name: 'True' }).click()
+          await page.keyboard.insertText('1')
+          await cmdBar.progressCmdBar()
+          await page.keyboard.insertText('100')
+          await cmdBar.progressCmdBar()
+          await cmdBar.expectState({
+            stage: 'review',
+            headerArguments: {
+              AngleStart: '0',
+              AxisOrEdge: 'Edge',
+              Edge: `1 ${selectionType}`,
+              CounterClockWise: '',
+              Length: '100',
+              Radius: '1',
+              Revolutions: '20',
+            },
+            commandName: 'Helix',
+          })
+          await cmdBar.progressCmdBar()
+        })
+
+        await test.step(`Confirm code is added to the editor, scene has changed`, async () => {
+          await toolbar.openPane('code')
+          await editor.expectEditor.toContain(expectedOutput)
+          await toolbar.closePane('code')
+        })
+
+        await test.step(`Edit helix through the feature tree`, async () => {
+          await toolbar.openPane('feature-tree')
+          const operationButton = await toolbar.getFeatureTreeOperation(
+            'Helix',
+            0
+          )
+          await operationButton.dblclick()
+          const initialInput = '100'
+          const newInput = '50'
+          await cmdBar.expectState({
+            commandName: 'Helix',
+            stage: 'arguments',
+            currentArgKey: 'length',
+            currentArgValue: initialInput,
+            headerArguments: {
+              AngleStart: '0',
+              CounterClockWise: '',
+              Length: initialInput,
+              Radius: '1',
+              Revolutions: '20',
+            },
+            highlightedHeaderArg: 'length',
+          })
+          await expect(cmdBar.currentArgumentInput).toBeVisible()
+          await cmdBar.currentArgumentInput
+            .locator('.cm-content')
+            .fill(newInput)
+          await cmdBar.progressCmdBar()
+          await cmdBar.expectState({
+            stage: 'review',
+            headerArguments: {
+              AngleStart: '0',
+              CounterClockWise: '',
+              Length: newInput,
+              Radius: '1',
+              Revolutions: '20',
+            },
+            commandName: 'Helix',
+          })
+          await cmdBar.progressCmdBar()
+          await toolbar.closePane('feature-tree')
+          await toolbar.openPane('code')
+          await editor.expectEditor.toContain(expectedEditedOutput)
+          await toolbar.closePane('code')
+        })
+
+        await test.step('Delete helix via feature tree selection', async () => {
+          await toolbar.openPane('feature-tree')
+          const operationButton = await toolbar.getFeatureTreeOperation(
+            'Helix',
+            0
+          )
+          await operationButton.click({ button: 'left' })
+          await page.keyboard.press('Delete')
+          await editor.expectEditor.not.toContain(expectedEditedOutput)
+          await expect(
+            await toolbar.getFeatureTreeOperation('Helix', 0)
+          ).not.toBeVisible()
+        })
+      })
+    }
+  )
+
   const loftPointAndClickCases = [
     { shouldPreselect: true },
     { shouldPreselect: false },

src/lang/modifyAst.ts

@@ -822,7 +822,7 @@ export function addHelix({
   angleStart: Expr
   ccw: boolean
   radius: Expr
-  axis: string
+  axis: Node<Literal> | Node<Name | CallExpression | CallExpressionKw>
   length: Expr
   insertIndex?: number
   variableName?: string
@@ -840,7 +840,7 @@ export function addHelix({
         createLabeledArg('angleStart', angleStart),
         createLabeledArg('ccw', createLiteral(ccw)),
         createLabeledArg('radius', radius),
-        createLabeledArg('axis', createLiteral(axis)),
+        createLabeledArg('axis', axis),
         createLabeledArg('length', length),
       ]
     )

src/lang/modifyAst/addRevolve.ts

@@ -32,11 +32,66 @@ import {
 import { Artifact, getPathsFromArtifact } from 'lang/std/artifactGraph'
 import { kclManager } from 'lib/singletons'
 
+export function getAxisExpressionAndIndex(
+  axisOrEdge: 'Axis' | 'Edge',
+  axis: string | undefined,
+  edge: Selections | undefined,
+  ast: Node<Program>
+) {
+  let generatedAxis
+  let axisDeclaration: PathToNode | null = null
+  let axisIndexIfAxis: number | undefined = undefined
+
+  if (axisOrEdge === 'Edge' && edge) {
+    const pathToAxisSelection = getNodePathFromSourceRange(
+      ast,
+      edge.graphSelections[0]?.codeRef.range
+    )
+    const lineNode = getNodeFromPath<CallExpression | CallExpressionKw>(
+      ast,
+      pathToAxisSelection,
+      ['CallExpression', 'CallExpressionKw']
+    )
+    if (err(lineNode)) return lineNode
+
+    const tagResult = mutateAstWithTagForSketchSegment(ast, pathToAxisSelection)
+
+    // Have the tag whether it is already created or a new one is generated
+    if (err(tagResult)) return tagResult
+    const { tag } = tagResult
+    const axisSelection = edge?.graphSelections[0]?.artifact
+    if (!axisSelection) return new Error('Generated axis selection is missing.')
+    generatedAxis = getEdgeTagCall(tag, axisSelection)
+    if (
+      axisSelection.type === 'segment' ||
+      axisSelection.type === 'path' ||
+      axisSelection.type === 'edgeCut'
+    ) {
+      axisDeclaration = axisSelection.codeRef.pathToNode
+      if (!axisDeclaration)
+        return new Error('Expected to fine axis declaration')
+      const axisIndexInPathToNode =
+        axisDeclaration.findIndex((a) => a[0] === 'body') + 1
+      const value = axisDeclaration[axisIndexInPathToNode][0]
+      if (typeof value !== 'number')
+        return new Error('expected axis index value to be a number')
+      axisIndexIfAxis = value
+    }
+  } else if (axisOrEdge === 'Axis' && axis) {
+    generatedAxis = createLiteral(axis)
+  }
+
+  return {
+    generatedAxis,
+    axisIndexIfAxis,
+  }
+}
+
 export function revolveSketch(
   ast: Node<Program>,
   pathToSketchNode: PathToNode,
   angle: Expr = createLiteral(4),
-  axisOrEdge: string,
+  axisOrEdge: 'Axis' | 'Edge',
   axis: string,
   edge: Selections,
   artifactGraph: ArtifactGraph,
@@ -58,44 +113,6 @@ export function revolveSketch(
   const clonedAst = structuredClone(ast)
   const sketchNode = getNodeFromPath(clonedAst, pathToSketchNode)
   if (err(sketchNode)) return sketchNode
-
-  let generatedAxis
-  let axisDeclaration: PathToNode | null = null
-
-  if (axisOrEdge === 'Edge') {
-    const pathToAxisSelection = getNodePathFromSourceRange(
-      clonedAst,
-      edge.graphSelections[0]?.codeRef.range
-    )
-    const lineNode = getNodeFromPath<CallExpression | CallExpressionKw>(
-      clonedAst,
-      pathToAxisSelection,
-      ['CallExpression', 'CallExpressionKw']
-    )
-    if (err(lineNode)) return lineNode
-
-    const tagResult = mutateAstWithTagForSketchSegment(
-      clonedAst,
-      pathToAxisSelection
-    )
-
-    // Have the tag whether it is already created or a new one is generated
-    if (err(tagResult)) return tagResult
-    const { tag } = tagResult
-    const axisSelection = edge?.graphSelections[0]?.artifact
-    if (!axisSelection) return new Error('Generated axis selection is missing.')
-    generatedAxis = getEdgeTagCall(tag, axisSelection)
-    if (
-      axisSelection.type === 'segment' ||
-      axisSelection.type === 'path' ||
-      axisSelection.type === 'edgeCut'
-    ) {
-      axisDeclaration = axisSelection.codeRef.pathToNode
-    }
-  } else {
-    generatedAxis = createLiteral(axis)
-  }
-
   const sketchVariableDeclaratorNode = getNodeFromPath<VariableDeclarator>(
     clonedAst,
     pathToSketchNode,
@@ -104,6 +121,9 @@ export function revolveSketch(
   if (err(sketchVariableDeclaratorNode)) return sketchVariableDeclaratorNode
   const { node: sketchVariableDeclarator } = sketchVariableDeclaratorNode
 
+  const getAxisResult = getAxisExpressionAndIndex(axisOrEdge, axis, edge, ast)
+  if (err(getAxisResult)) return getAxisResult
+  const { generatedAxis, axisIndexIfAxis } = getAxisResult
   if (!generatedAxis) return new Error('Generated axis selection is missing.')
 
   const revolveCall = createCallExpressionStdLibKw(
@@ -124,15 +144,8 @@ export function revolveSketch(
   }
 
   // If an axis was selected in KCL, find the max index to insert the revolve command
-  if (axisDeclaration) {
+  if (axisIndexIfAxis) {
-    const axisIndexInPathToNode =
+    sketchIndexInBody = Math.max(sketchIndexInBody, axisIndexIfAxis)
-      axisDeclaration.findIndex((a) => a[0] === 'body') + 1
-    const axisIndex = axisDeclaration[axisIndexInPathToNode][0]
-
-    if (typeof axisIndex !== 'number')
-      return new Error('expected axisIndex to be a number')
-
-    sketchIndexInBody = Math.max(sketchIndexInBody, axisIndex)
   }
 
   clonedAst.body.splice(sketchIndexInBody + 1, 0, VariableDeclaration)

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -96,12 +96,15 @@ export type ModelingCommandSchema = {
   Helix: {
     // Enables editing workflow
     nodeToEdit?: PathToNode
+    // Flow arg
+    axisOrEdge: 'Axis' | 'Edge'
     // KCL stdlib arguments
+    axis: string | undefined
+    edge: Selections | undefined
     revolutions: KclCommandValue
     angleStart: KclCommandValue
     ccw: boolean
     radius: KclCommandValue
-    axis: string
     length: KclCommandValue
   }
   'event.parameter.create': {
@@ -532,6 +535,38 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
         required: false,
         hidden: true,
       },
+      axisOrEdge: {
+        inputType: 'options',
+        required: true,
+        defaultValue: 'Axis',
+        options: [
+          { name: 'Axis', isCurrent: true, value: 'Axis' },
+          { name: 'Edge', isCurrent: false, value: 'Edge' },
+        ],
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
+      },
+      axis: {
+        inputType: 'options',
+        required: (commandContext) =>
+          ['Axis'].includes(
+            commandContext.argumentsToSubmit.axisOrEdge as string
+          ),
+        options: [
+          { name: 'X Axis', value: 'X' },
+          { name: 'Y Axis', value: 'Y' },
+          { name: 'Z Axis', value: 'Z' },
+        ],
+      },
+      edge: {
+        required: (commandContext) =>
+          ['Edge'].includes(
+            commandContext.argumentsToSubmit.axisOrEdge as string
+          ),
+        inputType: 'selection',
+        selectionTypes: ['segment', 'sweepEdge'],
+        multiple: false,
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
+      },
       revolutions: {
         inputType: 'kcl',
         defaultValue: '1',
@@ -547,6 +582,7 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
       ccw: {
         inputType: 'options',
         required: true,
+        displayName: 'CounterClockWise',
         defaultValue: false,
         options: [
           { name: 'False', value: false },
@@ -558,16 +594,6 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
         defaultValue: KCL_DEFAULT_LENGTH,
         required: true,
       },
-      axis: {
-        inputType: 'options',
-        required: true,
-        defaultValue: 'X',
-        options: [
-          { name: 'X Axis', value: 'X' },
-          { name: 'Y Axis', value: 'Y' },
-          { name: 'Z Axis', value: 'Z' },
-        ],
-      },
       length: {
         inputType: 'kcl',
         defaultValue: KCL_DEFAULT_LENGTH,

src/lib/operations.ts

@@ -3,6 +3,7 @@ import {
   Artifact,
   getArtifactOfTypes,
   getCapCodeRef,
+  getSweepEdgeCodeRef,
 } from 'lang/std/artifactGraph'
 import { Operation } from '@rust/kcl-lib/bindings/Operation'
 import { codeManager, engineCommandManager, kclManager } from './singletons'
@@ -458,12 +459,85 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   }
 
   // TODO: find a way to loop over the arguments while keeping it safe
+
+  // axis options string arg
+  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis) {
+    return baseCommand
+  }
+
+  const axisValue = operation.labeledArgs.axis.value
+  let axisOrEdge: 'Axis' | 'Edge' | undefined
+  let axis: string | undefined
+  let edge: Selections | undefined
+  if (axisValue.type === 'String') {
+    // default axis casee
+    axisOrEdge = 'Axis'
+    axis = axisValue.value
+  } else if (axisValue.type === 'TagIdentifier' && axisValue.artifact_id) {
+    // segment case
+    axisOrEdge = 'Edge'
+    const artifact = getArtifactOfTypes(
+      {
+        key: axisValue.artifact_id,
+        types: ['segment'],
+      },
+      engineCommandManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return baseCommand
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef: artifact.codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else if (axisValue.type === 'Uuid') {
+    // sweepEdge case
+    axisOrEdge = 'Edge'
+    const artifact = getArtifactOfTypes(
+      {
+        key: axisValue.value,
+        types: ['sweepEdge'],
+      },
+      engineCommandManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return baseCommand
+    }
+
+    const codeRef = getSweepEdgeCodeRef(
+      artifact,
+      engineCommandManager.artifactGraph
+    )
+    if (err(codeRef)) {
+      return baseCommand
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else {
+    return baseCommand
+  }
+
   // revolutions kcl arg
   if (
     !('revolutions' in operation.labeledArgs) ||
     !operation.labeledArgs.revolutions
-  )
+  ) {
     return baseCommand
+  }
   const revolutions = await stringToKclExpression(
     codeManager.code.slice(
       operation.labeledArgs.revolutions.sourceRange[0],
@@ -476,19 +550,23 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   if (
     !('angleStart' in operation.labeledArgs) ||
     !operation.labeledArgs.angleStart
-  )
+  ) {
     return baseCommand
+  }
   const angleStart = await stringToKclExpression(
     codeManager.code.slice(
       operation.labeledArgs.angleStart.sourceRange[0],
       operation.labeledArgs.angleStart.sourceRange[1]
     )
   )
-  if (err(angleStart) || 'errors' in angleStart) return baseCommand
+  if (err(angleStart) || 'errors' in angleStart) {
+    return baseCommand
+  }
 
   // counterClockWise options boolean arg
-  if (!('ccw' in operation.labeledArgs) || !operation.labeledArgs.ccw)
+  if (!('ccw' in operation.labeledArgs) || !operation.labeledArgs.ccw) {
     return baseCommand
+  }
   const ccw =
     codeManager.code.slice(
       operation.labeledArgs.ccw.sourceRange[0],
@@ -496,46 +574,47 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
     ) === 'true'
 
   // radius kcl arg
-  if (!('radius' in operation.labeledArgs) || !operation.labeledArgs.radius)
+  if (!('radius' in operation.labeledArgs) || !operation.labeledArgs.radius) {
+    console.log(
+      "!('radius' in operation.labeledArgs) || !operation.labeledArgs.radius"
+    )
     return baseCommand
+  }
   const radius = await stringToKclExpression(
     codeManager.code.slice(
       operation.labeledArgs.radius.sourceRange[0],
       operation.labeledArgs.radius.sourceRange[1]
     )
   )
-  if (err(radius) || 'errors' in radius) return baseCommand
+  if (err(radius) || 'errors' in radius) {
-
-  // axis options string arg
-  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis)
     return baseCommand
-  const axis = codeManager.code
+  }
-    .slice(
-      operation.labeledArgs.axis.sourceRange[0],
-      operation.labeledArgs.axis.sourceRange[1]
-    )
-    .replaceAll("'", '') // TODO: fix this crap
 
   // length kcl arg
-  if (!('length' in operation.labeledArgs) || !operation.labeledArgs.length)
+  if (!('length' in operation.labeledArgs) || !operation.labeledArgs.length) {
     return baseCommand
+  }
   const length = await stringToKclExpression(
     codeManager.code.slice(
       operation.labeledArgs.length.sourceRange[0],
       operation.labeledArgs.length.sourceRange[1]
     )
   )
-  if (err(length) || 'errors' in length) return baseCommand
+  if (err(length) || 'errors' in length) {
+    return baseCommand
+  }
 
   // Assemble the default argument values for the Offset Plane command,
   // with `nodeToEdit` set, which will let the Offset Plane actor know
   // to edit the node that corresponds to the StdLibCall.
   const argDefaultValues: ModelingCommandSchema['Helix'] = {
+    axisOrEdge,
+    axis,
+    edge,
     revolutions,
     angleStart,
     ccw,
     radius,
-    axis,
     length,
     nodeToEdit: getNodePathFromSourceRange(
       kclManager.ast,

src/machines/modelingMachine.ts

@@ -41,7 +41,10 @@ import {
   applyConstraintEqualLength,
   setEqualLengthInfo,
 } from 'components/Toolbar/EqualLength'
-import { revolveSketch } from 'lang/modifyAst/addRevolve'
+import {
+  getAxisExpressionAndIndex,
+  revolveSketch,
+} from 'lang/modifyAst/addRevolve'
 import {
   addHelix,
   addOffsetPlane,
@@ -1904,11 +1907,13 @@ export const modelingMachine = setup({
         // Extract inputs
         const ast = kclManager.ast
         const {
+          axisOrEdge,
+          axis,
+          edge,
           revolutions,
           angleStart,
           ccw,
           radius,
-          axis,
           length,
           nodeToEdit,
         } = input
@@ -1936,6 +1941,25 @@ export const modelingMachine = setup({
           opInsertIndex = nodeToEdit[1][0]
         }
 
+        const getAxisResult = getAxisExpressionAndIndex(
+          axisOrEdge,
+          axis,
+          edge,
+          ast
+        )
+        if (err(getAxisResult)) return getAxisResult
+        const { generatedAxis } = getAxisResult
+        if (!generatedAxis) {
+          return new Error('Generated axis selection is missing.')
+        }
+
+        // TODO: figure out if we want to smart insert after the sketch as below
+        // *or* after the sweep that consumes the sketch, in which case the below code doesn't work
+        // If an axis was selected in KCL, find the max index to insert the revolve command
+        // if (axisIndexIfAxis) {
+        // opInsertIndex = axisIndexIfAxis + 1
+        // }
+
         for (const variable of [revolutions, angleStart, radius, length]) {
           // Insert the variable if it exists
           if (
@@ -1958,33 +1982,28 @@ export const modelingMachine = setup({
             ? variable.variableIdentifierAst
             : variable.valueAst
 
-        const result = addHelix({
+        const { modifiedAst, pathToNode } = addHelix({
           node: ast,
           revolutions: valueOrVariable(revolutions),
           angleStart: valueOrVariable(angleStart),
           ccw,
           radius: valueOrVariable(radius),
-          axis,
+          axis: generatedAxis,
           length: valueOrVariable(length),
           insertIndex: opInsertIndex,
           variableName: opVariableName,
         })
-
+        await updateModelingState(
-        const updateAstResult = await kclManager.updateAst(
+          modifiedAst,
-          result.modifiedAst,
-          true,
           {
-            focusPath: [result.pathToNode],
+            kclManager,
+            editorManager,
+            codeManager,
+          },
+          {
+            focusPath: [pathToNode],
           }
         )
-
-        await codeManager.updateEditorWithAstAndWriteToFile(
-          updateAstResult.newAst
-        )
-
-        if (updateAstResult?.selections) {
-          editorManager.selectRange(updateAstResult?.selections)
-        }
       }
     ),
     sweepAstMod: fromPromise(