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>
2025-03-24 16:08:19 -04:00
parent fdeb2b3f49
commit 89dd4fb039
7 changed files with 401 additions and 107 deletions
--- a/e2e/playwright/point-click.spec.ts
+++ b/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',
-        currentArgValue: '1',
+        currentArgKey: 'axisOrEdge',
+        currentArgValue: '',
        headerArguments: {
          AngleStart: '',
-          Axis: '',
-          Ccw: '',
+          AxisOrEdge: '',
+          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 },
--- a/e2e/playwright/snapshot-tests.spec.ts-snapshots/extrude-on-default-planes-should-be-stable-XZ-1-Google-Chrome-linux.png
+++ b/e2e/playwright/snapshot-tests.spec.ts-snapshots/extrude-on-default-planes-should-be-stable-XZ-1-Google-Chrome-linux.png
--- a/src/lang/modifyAst.ts
+++ b/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),
      ]
    )
--- a/src/lang/modifyAst/addRevolve.ts
+++ b/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) {
-    const axisIndexInPathToNode =
-      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)
+  if (axisIndexIfAxis) {
+    sketchIndexInBody = Math.max(sketchIndexInBody, axisIndexIfAxis)
  }

  clonedAst.body.splice(sketchIndexInBody + 1, 0, VariableDeclaration)
--- a/src/lib/commandBarConfigs/modelingCommandConfig.ts
+++ b/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,
--- a/src/lib/operations.ts
+++ b/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
-
-  // axis options string arg
-  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis)
+  if (err(radius) || 'errors' in radius) {
    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,
--- a/src/machines/modelingMachine.ts
+++ b/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,
        })
-
-        const updateAstResult = await kclManager.updateAst(
-          result.modifiedAst,
-          true,
+        await updateModelingState(
+          modifiedAst,
          {
-            focusPath: [result.pathToNode],
+            kclManager,
+            editorManager,
+            codeManager,
+          },
+          {
+            focusPath: [pathToNode],
          }
        )
-
-        await codeManager.updateEditorWithAstAndWriteToFile(
-          updateAstResult.newAst
-        )
-
-        if (updateAstResult?.selections) {
-          editorManager.selectRange(updateAstResult?.selections)
-        }
      }
    ),
    sweepAstMod: fromPromise(