Add edit flow for Revolve point-and-click (#5951)

* 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

* WIP: Add edit flow for Revolve
Fixes #5504

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

* Clean up, add edit steps to three e2e tests

* Fix up tests after ccw change

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

* Fix tsc

* Remove uneeded return

* Update 2020 snapshots after helix change

* Update 2020 snapshots after helix change

* 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! 📷🐛

* Another one :djkhaled:

* 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! 📷🐛

* Move to fromPromise actor, fix test

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

* Clean up

* 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! 📷🐛

* Lint

* Fix third test

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

* Lint

* Fix edit insert order

* Fix axis, edge cases with new var, edit new var

* Add test case for new variable creation

* Clean up addRevolve with getSafeInsertIndex

---------

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

e2e/playwright/point-click.spec.ts

@@ -3476,6 +3476,39 @@ segAng(rectangleSegmentA002),
 
       const newCodeToFind = `revolve001 = revolve(sketch002, angle = 360, axis = 'X')`
       expect(editor.expectEditor.toContain(newCodeToFind)).toBeTruthy()
+
+      // Edit flow
+      const newAngle = '90'
+      await toolbar.openPane('feature-tree')
+      const operationButton = await toolbar.getFeatureTreeOperation(
+        'Revolve',
+        0
+      )
+      await operationButton.dblclick({ button: 'left' })
+      await cmdBar.expectState({
+        commandName: 'Revolve',
+        currentArgKey: 'angle',
+        currentArgValue: '360',
+        headerArguments: {
+          Angle: '360',
+        },
+        highlightedHeaderArg: 'angle',
+        stage: 'arguments',
+      })
+      await page.keyboard.insertText(newAngle)
+      await cmdBar.progressCmdBar()
+      await cmdBar.expectState({
+        stage: 'review',
+        headerArguments: {
+          Angle: newAngle,
+        },
+        commandName: 'Revolve',
+      })
+      await cmdBar.progressCmdBar()
+      await toolbar.closePane('feature-tree')
+      await editor.expectEditor.toContain(
+        newCodeToFind.replace('angle = 360', 'angle = ' + newAngle)
+      )
     })
     test('revolve surface around edge from an extruded solid2d', async ({
       context,
@@ -3486,26 +3519,22 @@ segAng(rectangleSegmentA002),
       toolbar,
       cmdBar,
     }) => {
-      const initialCode = `
+      const initialCode = `sketch001 = startSketchOn(XZ)
-sketch001 = startSketchOn(XZ)
+  |> startProfileAt([-102.57, 101.72], %)
-|> startProfileAt([-102.57, 101.72], %)
+  |> angledLine([0, 202.6], %, $rectangleSegmentA001)
-|> angledLine([0, 202.6], %, $rectangleSegmentA001)
+  |> angledLine([
-|> angledLine([
+       segAng(rectangleSegmentA001) - 90,
-segAng(rectangleSegmentA001) - 90,
+       202.6
-202.6
+     ], %, $rectangleSegmentB001)
-], %, $rectangleSegmentB001)
+  |> angledLine([
-|> angledLine([
+       segAng(rectangleSegmentA001),
-segAng(rectangleSegmentA001),
+       -segLen(rectangleSegmentA001)
--segLen(rectangleSegmentA001)
+     ], %, $rectangleSegmentC001)
-], %, $rectangleSegmentC001)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
-|> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
-|> close()
 extrude001 = extrude(sketch001, length = 50)
 sketch002 = startSketchOn(extrude001, rectangleSegmentA001)
-|> circle(
+  |> circle(center = [-11.34, 10.0], radius = 8.69)
-center = [-11.34, 10.0],
-radius = 8.69
-)
 `
       await context.addInitScript((initialCode) => {
         localStorage.setItem('persistCode', initialCode)
@@ -3523,9 +3552,49 @@ radius = 8.69
       const lineCodeToSelection = `|> angledLine([0, 202.6], %, $rectangleSegmentA001)`
       await page.getByText(lineCodeToSelection).click()
       await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
 
-      const newCodeToFind = `revolve001 = revolve(sketch002, angle = 360, axis = getOppositeEdge(rectangleSegmentA001)) `
+      const newCodeToFind = `revolve001 = revolve(sketch002, angle = 360, axis = rectangleSegmentA001)`
-      expect(editor.expectEditor.toContain(newCodeToFind)).toBeTruthy()
+      await editor.expectEditor.toContain(newCodeToFind)
+
+      // Edit flow
+      const newAngle = '180'
+      await toolbar.openPane('feature-tree')
+      const operationButton = await toolbar.getFeatureTreeOperation(
+        'Revolve',
+        0
+      )
+      await operationButton.dblclick({ button: 'left' })
+      await cmdBar.expectState({
+        commandName: 'Revolve',
+        currentArgKey: 'angle',
+        currentArgValue: '360',
+        headerArguments: {
+          Angle: '360',
+        },
+        highlightedHeaderArg: 'angle',
+        stage: 'arguments',
+      })
+      await page.keyboard.insertText(newAngle)
+      await page.getByRole('button', { name: 'Create new variable' }).click()
+      await expect(page.getByPlaceholder('Variable name')).toHaveValue(
+        'angle001'
+      )
+      await cmdBar.progressCmdBar()
+      await cmdBar.expectState({
+        stage: 'review',
+        headerArguments: {
+          Angle: newAngle,
+        },
+        commandName: 'Revolve',
+      })
+      await cmdBar.progressCmdBar()
+      await toolbar.closePane('feature-tree')
+      await editor.expectEditor.toContain('angle001 = ' + newAngle)
+      await editor.expectEditor.toContain(
+        newCodeToFind.replace('angle = 360', 'angle = angle001')
+      )
     })
     test('revolve sketch circle around line segment from startProfileAt sketch', async ({
       context,
@@ -3536,11 +3605,10 @@ radius = 8.69
       toolbar,
       cmdBar,
     }) => {
-      const initialCode = `
+      const initialCode = `sketch002 = startSketchOn(XY)
-    sketch002 = startSketchOn(XY)
   |> startProfileAt([-2.02, 1.79], %)
   |> xLine(length = 2.6)
-    sketch001 = startSketchOn('-XY')
+sketch001 = startSketchOn(-XY)
   |> startProfileAt([-0.48, 1.25], %)
   |> angledLine([0, 2.38], %, $rectangleSegmentA001)
   |> angledLine([segAng(rectangleSegmentA001) - 90, 2.4], %, $rectangleSegmentB001)
@@ -3550,12 +3618,9 @@ radius = 8.69
      ], %, $rectangleSegmentC001)
   |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
   |> close()
-    extrude001 = extrude(sketch001, length = 5)
+extrude001 = extrude(sketch001, length = 5)
-    sketch003 = startSketchOn(extrude001, 'START')
+sketch003 = startSketchOn(extrude001, 'START')
-      |> circle(
+  |> circle(center = [-0.69, 0.56], radius = 0.28)
-        center = [-0.69, 0.56],
-        radius = 0.28
-      )
 `
 
       await context.addInitScript((initialCode) => {
@@ -3574,9 +3639,44 @@ radius = 8.69
       const lineCodeToSelection = `|> xLine(length = 2.6)`
       await page.getByText(lineCodeToSelection).click()
       await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
 
       const newCodeToFind = `revolve001 = revolve(sketch003, angle = 360, axis = seg01)`
       expect(editor.expectEditor.toContain(newCodeToFind)).toBeTruthy()
+
+      // Edit flow
+      const newAngle = '270'
+      await toolbar.openPane('feature-tree')
+      const operationButton = await toolbar.getFeatureTreeOperation(
+        'Revolve',
+        0
+      )
+      await operationButton.dblclick({ button: 'left' })
+      await cmdBar.expectState({
+        commandName: 'Revolve',
+        currentArgKey: 'angle',
+        currentArgValue: '360',
+        headerArguments: {
+          Angle: '360',
+        },
+        highlightedHeaderArg: 'angle',
+        stage: 'arguments',
+      })
+      await page.keyboard.insertText(newAngle)
+      await cmdBar.progressCmdBar()
+      await cmdBar.expectState({
+        stage: 'review',
+        headerArguments: {
+          Angle: newAngle,
+        },
+        commandName: 'Revolve',
+      })
+      await cmdBar.progressCmdBar()
+      await toolbar.closePane('feature-tree')
+      await editor.expectEditor.toContain(
+        newCodeToFind.replace('angle = 360', 'angle = ' + newAngle)
+      )
     })
   })
 

src/lang/modifyAst/addEdgeTreatment.ts

@@ -320,6 +320,7 @@ export function mutateAstWithTagForSketchSegment(
 
   // Check whether selection is a valid segment
   if (
+    !segmentNode.node.callee ||
     !(
       segmentNode.node.callee.name.name in sketchLineHelperMap ||
       segmentNode.node.callee.name.name in sketchLineHelperMapKw

src/lang/modifyAst/addRevolve.ts

@@ -1,14 +1,6 @@
 import { err } from 'lib/trap'
 import { KCL_DEFAULT_CONSTANT_PREFIXES } from 'lib/constants'
-import {
+import { Program, PathToNode, Expr, VariableDeclarator } from 'lang/wasm'
-  Program,
-  PathToNode,
-  Expr,
-  CallExpression,
-  VariableDeclarator,
-  CallExpressionKw,
-  ArtifactGraph,
-} from 'lang/wasm'
 import { Selections } from 'lib/selections'
 import { Node } from '@rust/kcl-lib/bindings/Node'
 import {
@@ -29,8 +21,7 @@ import {
   mutateAstWithTagForSketchSegment,
   getEdgeTagCall,
 } from 'lang/modifyAst/addEdgeTreatment'
-import { Artifact, getPathsFromArtifact } from 'lang/std/artifactGraph'
+import { getSafeInsertIndex } from 'lang/queryAst/getSafeInsertIndex'
-import { kclManager } from 'lib/singletons'
 
 export function getAxisExpressionAndIndex(
   axisOrEdge: 'Axis' | 'Edge',
@@ -47,13 +38,6 @@ export function getAxisExpressionAndIndex(
       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
@@ -90,12 +74,12 @@ export function getAxisExpressionAndIndex(
 export function revolveSketch(
   ast: Node<Program>,
   pathToSketchNode: PathToNode,
-  angle: Expr = createLiteral(4),
+  angle: Expr,
   axisOrEdge: 'Axis' | 'Edge',
-  axis: string,
+  axis: string | undefined,
-  edge: Selections,
+  edge: Selections | undefined,
-  artifactGraph: ArtifactGraph,
+  variableName?: string,
-  artifact?: Artifact
+  insertIndex?: number
 ):
   | {
       modifiedAst: Node<Program>
@@ -103,16 +87,7 @@ export function revolveSketch(
       pathToRevolveArg: PathToNode
     }
   | Error {
-  const orderedSketchNodePaths = getPathsFromArtifact({
-    artifact: artifact,
-    sketchPathToNode: pathToSketchNode,
-    artifactGraph,
-    ast: kclManager.ast,
-  })
-  if (err(orderedSketchNodePaths)) return orderedSketchNodePaths
   const clonedAst = structuredClone(ast)
-  const sketchNode = getNodeFromPath(clonedAst, pathToSketchNode)
-  if (err(sketchNode)) return sketchNode
   const sketchVariableDeclaratorNode = getNodeFromPath<VariableDeclarator>(
     clonedAst,
     pathToSketchNode,
@@ -121,9 +96,14 @@ export function revolveSketch(
   if (err(sketchVariableDeclaratorNode)) return sketchVariableDeclaratorNode
   const { node: sketchVariableDeclarator } = sketchVariableDeclaratorNode
 
-  const getAxisResult = getAxisExpressionAndIndex(axisOrEdge, axis, edge, ast)
+  const getAxisResult = getAxisExpressionAndIndex(
+    axisOrEdge,
+    axis,
+    edge,
+    clonedAst
+  )
   if (err(getAxisResult)) return getAxisResult
-  const { generatedAxis, axisIndexIfAxis } = getAxisResult
+  const { generatedAxis } = getAxisResult
   if (!generatedAxis) return new Error('Generated axis selection is missing.')
 
   const revolveCall = createCallExpressionStdLibKw(
@@ -134,29 +114,21 @@ export function revolveSketch(
 
   // We're not creating a pipe expression,
   // but rather a separate constant for the extrusion
-  const name = findUniqueName(clonedAst, KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE)
+  const name =
-  const VariableDeclaration = createVariableDeclaration(name, revolveCall)
+    variableName ??
-  const lastSketchNodePath =
+    findUniqueName(clonedAst, KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE)
-    orderedSketchNodePaths[orderedSketchNodePaths.length - 1]
+  const variableDeclaration = createVariableDeclaration(name, revolveCall)
-  let sketchIndexInBody = Number(lastSketchNodePath[1][0])
+  const bodyInsertIndex =
-  if (typeof sketchIndexInBody !== 'number') {
+    insertIndex ?? getSafeInsertIndex(revolveCall, clonedAst)
-    return new Error('expected sketchIndexInBody to be a number')
+  clonedAst.body.splice(bodyInsertIndex, 0, variableDeclaration)
-  }
+  const argIndex = 0
-
-  // If an axis was selected in KCL, find the max index to insert the revolve command
-  if (axisIndexIfAxis) {
-    sketchIndexInBody = Math.max(sketchIndexInBody, axisIndexIfAxis)
-  }
-
-  clonedAst.body.splice(sketchIndexInBody + 1, 0, VariableDeclaration)
-
   const pathToRevolveArg: PathToNode = [
     ['body', ''],
-    [sketchIndexInBody + 1, 'index'],
+    [bodyInsertIndex, 'index'],
     ['declaration', 'VariableDeclaration'],
     ['init', 'VariableDeclarator'],
     ['arguments', 'CallExpressionKw'],
-    [0, ARG_INDEX_FIELD],
+    [argIndex, ARG_INDEX_FIELD],
     ['arg', LABELED_ARG_FIELD],
   ]
 

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -77,11 +77,15 @@ export type ModelingCommandSchema = {
     thickness: KclCommandValue
   }
   Revolve: {
+    // Enables editing workflow
+    nodeToEdit?: PathToNode
+    // Flow arg
+    axisOrEdge: 'Axis' | 'Edge'
+    // KCL stdlib arguments
     selection: Selections
     angle: KclCommandValue
-    axisOrEdge: 'Axis' | 'Edge'
+    axis: string | undefined
-    axis: string
+    edge: Selections | undefined
-    edge: Selections
   }
   Fillet: {
     // Enables editing workflow
@@ -472,6 +476,13 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
     icon: 'revolve',
     needsReview: true,
     args: {
+      nodeToEdit: {
+        description:
+          'Path to the node in the AST to edit. Never shown to the user.',
+        skip: true,
+        inputType: 'text',
+        required: false,
+      },
       selection: {
         inputType: 'selection',
         selectionTypes: ['solid2d', 'segment'],
@@ -500,6 +511,7 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
           { name: 'X Axis', isCurrent: true, value: 'X' },
           { name: 'Y Axis', isCurrent: false, value: 'Y' },
         ],
+        hidden: (context) => Boolean(context.argumentsToSubmit.nodeToEdit),
       },
       edge: {
         required: (commandContext) =>

src/lib/operations.ts

@@ -800,6 +800,163 @@ const prepareToEditHelix: PrepareToEditCallback = async ({ operation }) => {
   }
 }
 
+const prepareToEditRevolve: PrepareToEditCallback = async ({
+  operation,
+  artifact,
+}) => {
+  const baseCommand = {
+    name: 'Revolve',
+    groupId: 'modeling',
+  }
+  if (
+    !artifact ||
+    !('pathId' in artifact) ||
+    operation.type !== 'StdLibCall' ||
+    !operation.labeledArgs
+  ) {
+    return { reason: 'Wrong operation type or artifact' }
+  }
+
+  // We have to go a little roundabout to get from the original artifact
+  // to the solid2DId that we need to pass to the command.
+  const pathArtifact = getArtifactOfTypes(
+    {
+      key: artifact.pathId,
+      types: ['path'],
+    },
+    kclManager.artifactGraph
+  )
+  if (
+    err(pathArtifact) ||
+    pathArtifact.type !== 'path' ||
+    !pathArtifact.solid2dId
+  ) {
+    return { reason: "Couldn't find related path artifact" }
+  }
+
+  const solid2DArtifact = getArtifactOfTypes(
+    {
+      key: pathArtifact.solid2dId,
+      types: ['solid2d'],
+    },
+    kclManager.artifactGraph
+  )
+  if (err(solid2DArtifact) || solid2DArtifact.type !== 'solid2d') {
+    return { reason: "Couldn't find related solid2d artifact" }
+  }
+
+  const selection = {
+    graphSelections: [
+      {
+        artifact: solid2DArtifact,
+        codeRef: pathArtifact.codeRef,
+      },
+    ],
+    otherSelections: [],
+  }
+
+  // axis options string arg
+  if (!('axis' in operation.labeledArgs) || !operation.labeledArgs.axis) {
+    return { reason: "Couldn't find axis argument" }
+  }
+
+  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'],
+      },
+      kclManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return { reason: "Couldn't find related edge artifact" }
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef: artifact.codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else if (axisValue.type === 'Uuid') {
+    // sweepEdge case
+    axisOrEdge = 'Edge'
+    const artifact = getArtifactOfTypes(
+      {
+        key: axisValue.value,
+        types: ['sweepEdge'],
+      },
+      kclManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return { reason: "Couldn't find related edge artifact" }
+    }
+
+    const codeRef = getSweepEdgeCodeRef(artifact, kclManager.artifactGraph)
+    if (err(codeRef)) {
+      return { reason: "Couldn't find related edge code ref" }
+    }
+
+    edge = {
+      graphSelections: [
+        {
+          artifact,
+          codeRef,
+        },
+      ],
+      otherSelections: [],
+    }
+  } else {
+    return { reason: 'The type of the axis argument is unsupported' }
+  }
+
+  // angle kcl arg
+  if (!('angle' in operation.labeledArgs) || !operation.labeledArgs.angle) {
+    return { reason: "Couldn't find angle argument" }
+  }
+  const angle = await stringToKclExpression(
+    codeManager.code.slice(
+      operation.labeledArgs.angle.sourceRange[0],
+      operation.labeledArgs.angle.sourceRange[1]
+    )
+  )
+  if (err(angle) || 'errors' in angle) {
+    return { reason: 'Error in angle argument retrieval' }
+  }
+
+  // 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['Revolve'] = {
+    axisOrEdge,
+    axis,
+    edge,
+    selection,
+    angle,
+    nodeToEdit: getNodePathFromSourceRange(
+      kclManager.ast,
+      sourceRangeFromRust(operation.sourceRange)
+    ),
+  }
+  return {
+    ...baseCommand,
+    argDefaultValues,
+  }
+}
+
 /**
  * A map of standard library calls to their corresponding information
  * for use in the feature tree UI.
@@ -872,6 +1029,7 @@ export const stdLibMap: Record<string, StdLibCallInfo> = {
   revolve: {
     label: 'Revolve',
     icon: 'revolve',
+    prepareToEdit: prepareToEditRevolve,
     supportsAppearance: true,
   },
   shell: {

src/machines/modelingMachine.ts

@@ -735,62 +735,6 @@ export const modelingMachine = setup({
         sketchDetails: event.output,
       }
     }),
-    'AST revolve': ({ context: { store }, event }) => {
-      if (event.type !== 'Revolve') return
-      ;(async () => {
-        if (!event.data) return
-        const { selection, angle, axis, edge, axisOrEdge } = event.data
-        let ast = kclManager.ast
-        if (
-          'variableName' in angle &&
-          angle.variableName &&
-          angle.insertIndex !== undefined
-        ) {
-          const newBody = [...ast.body]
-          newBody.splice(angle.insertIndex, 0, angle.variableDeclarationAst)
-          ast.body = newBody
-        }
-
-        // This is the selection of the sketch that will be revolved
-        const pathToNode = getNodePathFromSourceRange(
-          ast,
-          selection.graphSelections[0]?.codeRef.range
-        )
-
-        const revolveSketchRes = revolveSketch(
-          ast,
-          pathToNode,
-          'variableName' in angle
-            ? angle.variableIdentifierAst
-            : angle.valueAst,
-          axisOrEdge,
-          axis,
-          edge,
-          kclManager.artifactGraph,
-          selection.graphSelections[0]?.artifact
-        )
-        if (trap(revolveSketchRes)) return
-        const { modifiedAst, pathToRevolveArg } = revolveSketchRes
-
-        await updateModelingState(
-          modifiedAst,
-          EXECUTION_TYPE_REAL,
-          {
-            kclManager,
-            editorManager,
-            codeManager,
-          },
-          {
-            focusPath: [pathToRevolveArg],
-            zoomToFit: true,
-            zoomOnRangeAndType: {
-              range: selection.graphSelections[0]?.codeRef.range,
-              type: 'path',
-            },
-          }
-        )
-      })().catch(reportRejection)
-    },
     'set selection filter to curves only': () => {
       ;(async () => {
         await engineCommandManager.sendSceneCommand({
@@ -1837,6 +1781,87 @@ export const modelingMachine = setup({
         }
       )
     }),
+    revolveAstMod: fromPromise<
+      unknown,
+      ModelingCommandSchema['Revolve'] | undefined
+    >(async ({ input }) => {
+      if (!input) return new Error('No input provided')
+      const { nodeToEdit, selection, angle, axis, edge, axisOrEdge } = input
+      let ast = kclManager.ast
+      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 (nodeToEdit && typeof nodeToEdit[1][0] === 'number') {
+        // Extract the plane name from the node to edit
+        const nameNode = getNodeFromPath<VariableDeclaration>(
+          ast,
+          nodeToEdit,
+          'VariableDeclaration'
+        )
+        if (err(nameNode)) {
+          console.error('Error extracting plane name')
+        } else {
+          variableName = nameNode.node.declaration.id.name
+        }
+
+        // Removing the old extrusion statement
+        const newBody = [...ast.body]
+        newBody.splice(nodeToEdit[1][0], 1)
+        ast.body = newBody
+        insertIndex = nodeToEdit[1][0]
+      }
+
+      if (
+        'variableName' in angle &&
+        angle.variableName &&
+        angle.insertIndex !== undefined
+      ) {
+        const newBody = [...ast.body]
+        newBody.splice(angle.insertIndex, 0, angle.variableDeclarationAst)
+        ast.body = newBody
+        if (insertIndex) {
+          // if editing need to offset that new var
+          insertIndex += 1
+        }
+      }
+
+      // This is the selection of the sketch that will be revolved
+      const pathToNode = getNodePathFromSourceRange(
+        ast,
+        selection.graphSelections[0]?.codeRef.range
+      )
+
+      const revolveSketchRes = revolveSketch(
+        ast,
+        pathToNode,
+        'variableName' in angle ? angle.variableIdentifierAst : angle.valueAst,
+        axisOrEdge,
+        axis,
+        edge,
+        variableName,
+        insertIndex
+      )
+      if (trap(revolveSketchRes)) return
+      const { modifiedAst, pathToRevolveArg } = revolveSketchRes
+      await updateModelingState(
+        modifiedAst,
+        EXECUTION_TYPE_REAL,
+        {
+          kclManager,
+          editorManager,
+          codeManager,
+        },
+        {
+          focusPath: [pathToRevolveArg],
+          zoomToFit: true,
+          zoomOnRangeAndType: {
+            range: selection.graphSelections[0]?.codeRef.range,
+            type: 'path',
+          },
+        }
+      )
+    }),
     offsetPlaneAstMod: fromPromise(
       async ({
         input,
@@ -2767,9 +2792,8 @@ export const modelingMachine = setup({
         },
 
         Revolve: {
-          target: 'idle',
+          target: 'Applying revolve',
-          actions: ['AST revolve'],
+          reenter: true,
-          reenter: false,
         },
 
         Sweep: {
@@ -4023,6 +4047,22 @@ export const modelingMachine = setup({
       },
     },
 
+    'Applying revolve': {
+      invoke: {
+        src: 'revolveAstMod',
+        id: 'revolveAstMod',
+        input: ({ event }) => {
+          if (event.type !== 'Revolve') return undefined
+          return event.data
+        },
+        onDone: ['idle'],
+        onError: {
+          target: 'idle',
+          actions: 'toastError',
+        },
+      },
+    },
+
     'Applying offset plane': {
       invoke: {
         src: 'offsetPlaneAstMod',