Add test snapshots on two tests

Fixes #5504

e2e/playwright/point-click.spec.ts

@@ -3409,17 +3409,16 @@ sweep001 = sweep(sketch001, path = sketch002)
     })
   })
 
-  test.describe('Revolve point and click workflows', () => {
-    test('Base case workflow, auto spam continue in command bar', async ({
-      context,
-      page,
-      homePage,
-      scene,
-      editor,
-      toolbar,
-      cmdBar,
-    }) => {
-      const initialCode = `
+  test('Revolve point-and-click base case workflow with auto spam continue in command bar', async ({
+    context,
+    page,
+    homePage,
+    scene,
+    editor,
+    toolbar,
+    cmdBar,
+  }) => {
+    const initialCode = `
 sketch001 = startSketchOn(XZ)
 |> startProfileAt([-100.0, 100.0], %)
 |> angledLine([0, 200.0], %, $rectangleSegmentA001)
@@ -3446,13 +3445,14 @@ segAng(rectangleSegmentA002),
 |> close()
 `
 
-      await context.addInitScript((initialCode) => {
-        localStorage.setItem('persistCode', initialCode)
-      }, initialCode)
-      await page.setBodyDimensions({ width: 1000, height: 500 })
-      await homePage.goToModelingScene()
-      await scene.waitForExecutionDone()
+    await context.addInitScript((initialCode) => {
+      localStorage.setItem('persistCode', initialCode)
+    }, initialCode)
+    await page.setBodyDimensions({ width: 1000, height: 500 })
+    await homePage.goToModelingScene()
+    await scene.waitForExecutionDone()
 
+    await test.step('Add revolve through the command bar and code selection', async () => {
       // select line of code
       const codeToSelecton = `segAng(rectangleSegmentA002) - 90,`
       // revolve
@@ -3462,47 +3462,76 @@ segAng(rectangleSegmentA002),
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
       await cmdBar.progressCmdBar()
-
-      const newCodeToFind = `revolve001 = revolve(sketch002, angle = 360, axis = 'X')`
-      expect(editor.expectEditor.toContain(newCodeToFind)).toBeTruthy()
+      await editor.snapshot({ name: 'revolve-base-add' })
     })
-    test('revolve surface around edge from an extruded solid2d', async ({
-      context,
-      page,
-      homePage,
-      scene,
-      editor,
-      toolbar,
-      cmdBar,
-    }) => {
-      const initialCode = `
-sketch001 = startSketchOn(XZ)
-|> startProfileAt([-102.57, 101.72], %)
-|> angledLine([0, 202.6], %, $rectangleSegmentA001)
-|> angledLine([
-segAng(rectangleSegmentA001) - 90,
-202.6
-], %, $rectangleSegmentB001)
-|> angledLine([
-segAng(rectangleSegmentA001),
--segLen(rectangleSegmentA001)
-], %, $rectangleSegmentC001)
-|> line(endAbsolute = [profileStartX(%), profileStartY(%)])
-|> close()
+
+    await test.step('Edit flow', async () => {
+      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.snapshot({ name: 'revolve-base-edit' })
+    })
+  })
+
+  test('Revolve point-and-click surface around edge from an extruded solid2d', async ({
+    context,
+    page,
+    homePage,
+    scene,
+    editor,
+    toolbar,
+    cmdBar,
+  }) => {
+    const initialCode = `sketch001 = startSketchOn(XZ)
+  |> startProfileAt([-102.57, 101.72], %)
+  |> angledLine([0, 202.6], %, $rectangleSegmentA001)
+  |> angledLine([
+       segAng(rectangleSegmentA001) - 90,
+       202.6
+     ], %, $rectangleSegmentB001)
+  |> angledLine([
+       segAng(rectangleSegmentA001),
+       -segLen(rectangleSegmentA001)
+     ], %, $rectangleSegmentC001)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
 extrude001 = extrude(sketch001, length = 50)
 sketch002 = startSketchOn(extrude001, rectangleSegmentA001)
-|> circle(
-center = [-11.34, 10.0],
-radius = 8.69
-)
+  |> circle(center = [-11.34, 10.0], radius = 8.69)
 `
-      await context.addInitScript((initialCode) => {
-        localStorage.setItem('persistCode', initialCode)
-      }, initialCode)
-      await page.setBodyDimensions({ width: 1000, height: 500 })
-      await homePage.goToModelingScene()
-      await scene.waitForExecutionDone()
+    await context.addInitScript((initialCode) => {
+      localStorage.setItem('persistCode', initialCode)
+    }, initialCode)
+    await page.setBodyDimensions({ width: 1000, height: 500 })
+    await homePage.goToModelingScene()
+    await scene.waitForExecutionDone()
 
+    await test.step('Add through the command bar and code selection', async () => {
       // select line of code
       const codeToSelecton = `center = [-11.34, 10.0]`
       // revolve
@@ -3512,61 +3541,118 @@ radius = 8.69
       const lineCodeToSelection = `|> angledLine([0, 202.6], %, $rectangleSegmentA001)`
       await page.getByText(lineCodeToSelection).click()
       await cmdBar.progressCmdBar()
-
-      const newCodeToFind = `revolve001 = revolve(sketch002, angle = 360, axis = getOppositeEdge(rectangleSegmentA001)) `
-      expect(editor.expectEditor.toContain(newCodeToFind)).toBeTruthy()
+      await cmdBar.progressCmdBar()
+      await cmdBar.progressCmdBar()
+      await editor.snapshot({ name: 'revolve-solid-edge-add' })
     })
-    test('revolve sketch circle around line segment from startProfileAt sketch', async ({
-      context,
-      page,
-      homePage,
-      scene,
-      editor,
-      toolbar,
-      cmdBar,
-    }) => {
-      const initialCode = `
-    sketch002 = startSketchOn(XY)
-      |> startProfileAt([-2.02, 1.79], %)
-      |> xLine(length = 2.6)
-    sketch001 = startSketchOn('-XY')
-      |> startProfileAt([-0.48, 1.25], %)
-      |> angledLine([0, 2.38], %, $rectangleSegmentA001)
-      |> angledLine([segAng(rectangleSegmentA001) - 90, 2.4], %, $rectangleSegmentB001)
-      |> angledLine([
-        segAng(rectangleSegmentA001),
-          -segLen(rectangleSegmentA001)
-      ], %, $rectangleSegmentC001)
-      |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
-      |> close()
-    extrude001 = extrude(sketch001, length = 5)
-    sketch003 = startSketchOn(extrude001, 'START')
-      |> circle(
-        center = [-0.69, 0.56],
-        radius = 0.28
-      )
+
+    await test.step('Edit flow', async () => {
+      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 cmdBar.progressCmdBar()
+      await cmdBar.expectState({
+        stage: 'review',
+        headerArguments: {
+          Angle: newAngle,
+        },
+        commandName: 'Revolve',
+      })
+      await cmdBar.progressCmdBar()
+      await toolbar.closePane('feature-tree')
+      await editor.snapshot({ name: 'revolve-solid-edge-edit' })
+    })
+  })
+  test('Revolve point-and-click sketch circle around line segment from startProfileAt sketch', async ({
+    context,
+    page,
+    homePage,
+    scene,
+    editor,
+    toolbar,
+    cmdBar,
+  }) => {
+    const initialCode = `sketch002 = startSketchOn(XY)
+  |> startProfileAt([-2.02, 1.79], %)
+  |> xLine(length = 2.6)
+sketch001 = startSketchOn(-XY)
+  |> startProfileAt([-0.48, 1.25], %)
+  |> angledLine([0, 2.38], %, $rectangleSegmentA001)
+  |> angledLine([segAng(rectangleSegmentA001) - 90, 2.4], %, $rectangleSegmentB001)
+  |> angledLine([
+       segAng(rectangleSegmentA001),
+       -segLen(rectangleSegmentA001)
+     ], %, $rectangleSegmentC001)
+  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])
+  |> close()
+extrude001 = extrude(sketch001, length = 5)
+sketch003 = startSketchOn(extrude001, 'START')
+  |> circle(center = [-0.69, 0.56], radius = 0.28)
 `
 
-      await context.addInitScript((initialCode) => {
-        localStorage.setItem('persistCode', initialCode)
-      }, initialCode)
-      await page.setBodyDimensions({ width: 1000, height: 500 })
-      await homePage.goToModelingScene()
-      await scene.waitForExecutionDone()
+    await context.addInitScript((initialCode) => {
+      localStorage.setItem('persistCode', initialCode)
+    }, initialCode)
+    await page.setBodyDimensions({ width: 1000, height: 500 })
+    await homePage.goToModelingScene()
+    await scene.waitForExecutionDone()
 
-      // select line of code
-      const codeToSelecton = `center = [-0.69, 0.56]`
-      // revolve
-      await page.getByText(codeToSelecton).click()
-      await toolbar.revolveButton.click()
-      await page.getByText('Edge', { exact: true }).click()
-      const lineCodeToSelection = `|> xLine(length = 2.6)`
-      await page.getByText(lineCodeToSelection).click()
-      await cmdBar.progressCmdBar()
+    // select line of code
+    const codeToSelecton = `center = [-0.69, 0.56]`
+    // revolve
+    await page.getByText(codeToSelecton).click()
+    await toolbar.revolveButton.click()
+    await page.getByText('Edge', { exact: true }).click()
+    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()
+    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)
+    )
   })
 
   test(`Set appearance`, async ({

e2e/playwright/point-click.spec.ts-snapshots/revolve-base-add-chromium-darwin

@@ -0,0 +1 @@
+sketch001 = startSketchOn(XZ)  |> startProfileAt([-100.0, 100.0], %)  |> angledLine([0, 200.0], %, $rectangleSegmentA001)  |> angledLine([segAng(rectangleSegmentA001) - 90, 200], %, $rectangleSegmentB001)  |> angledLine([       segAng(rectangleSegmentA001),       -segLen(rectangleSegmentA001)     ], %, $rectangleSegmentC001)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()extrude001 = extrude(sketch001, length = 200)sketch002 = startSketchOn(extrude001, rectangleSegmentA001)  |> startProfileAt([-66.77, 84.81], %)  |> angledLine([180, 27.08], %, $rectangleSegmentA002)  |> angledLine([       segAng(rectangleSegmentA002) - 90,       27.8     ], %, $rectangleSegmentB002)  |> angledLine([       segAng(rectangleSegmentA002),       -segLen(rectangleSegmentA002)     ], %, $rectangleSegmentC002)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()revolve001 = revolve(sketch002, angle = 360, axis = 'X')

e2e/playwright/point-click.spec.ts-snapshots/revolve-base-edit-chromium-darwin

@@ -0,0 +1 @@
+sketch001 = startSketchOn(XZ)  |> startProfileAt([-100.0, 100.0], %)  |> angledLine([0, 200.0], %, $rectangleSegmentA001)  |> angledLine([segAng(rectangleSegmentA001) - 90, 200], %, $rectangleSegmentB001)  |> angledLine([       segAng(rectangleSegmentA001),       -segLen(rectangleSegmentA001)     ], %, $rectangleSegmentC001)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()extrude001 = extrude(sketch001, length = 200)sketch002 = startSketchOn(extrude001, rectangleSegmentA001)  |> startProfileAt([-66.77, 84.81], %)  |> angledLine([180, 27.08], %, $rectangleSegmentA002)  |> angledLine([       segAng(rectangleSegmentA002) - 90,       27.8     ], %, $rectangleSegmentB002)  |> angledLine([       segAng(rectangleSegmentA002),       -segLen(rectangleSegmentA002)     ], %, $rectangleSegmentC002)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()revolve001 = revolve(sketch002, angle = 360, axis = 'X')

e2e/playwright/point-click.spec.ts-snapshots/revolve-solid-edge-add-chromium-darwin

@@ -0,0 +1 @@
+sketch001 = startSketchOn(XZ)  |> startProfileAt([-102.57, 101.72], %)  |> angledLine([0, 202.6], %, $rectangleSegmentA001)  |> angledLine([       segAng(rectangleSegmentA001) - 90,       202.6     ], %, $rectangleSegmentB001)  |> angledLine([       segAng(rectangleSegmentA001),       -segLen(rectangleSegmentA001)     ], %, $rectangleSegmentC001)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()extrude001 = extrude(sketch001, length = 50)sketch002 = startSketchOn(extrude001, rectangleSegmentA001)  |> circle(center = [-11.34, 10.0], radius = 8.69)revolve001 = revolve(sketch002, angle = 360, axis = rectangleSegmentA001)

e2e/playwright/point-click.spec.ts-snapshots/revolve-solid-edge-edit-chromium-darwin

@@ -0,0 +1 @@
+sketch001 = startSketchOn(XZ)  |> startProfileAt([-102.57, 101.72], %)  |> angledLine([0, 202.6], %, $rectangleSegmentA001)  |> angledLine([       segAng(rectangleSegmentA001) - 90,       202.6     ], %, $rectangleSegmentB001)  |> angledLine([       segAng(rectangleSegmentA001),       -segLen(rectangleSegmentA001)     ], %, $rectangleSegmentC001)  |> line(endAbsolute = [profileStartX(%), profileStartY(%)])  |> close()extrude001 = extrude(sketch001, length = 50)sketch002 = startSketchOn(extrude001, rectangleSegmentA001)  |> circle(center = [-11.34, 10.0], radius = 8.69)revolve001 = revolve(sketch002, angle = 360, axis = rectangleSegmentA001)

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

@@ -4,9 +4,7 @@ import {
   Program,
   PathToNode,
   Expr,
-  CallExpression,
   VariableDeclarator,
-  CallExpressionKw,
   ArtifactGraph,
 } from 'lang/wasm'
 import { Selections } from 'lib/selections'
@@ -47,13 +45,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 +81,14 @@ export function getAxisExpressionAndIndex(
 export function revolveSketch(
   ast: Node<Program>,
   pathToSketchNode: PathToNode,
-  angle: Expr = createLiteral(4),
+  angle: Expr,
   axisOrEdge: 'Axis' | 'Edge',
-  axis: string,
-  edge: Selections,
+  axis: string | undefined,
+  edge: Selections | undefined,
   artifactGraph: ArtifactGraph,
-  artifact?: Artifact
+  artifact?: Artifact,
+  variableName?: string,
+  insertIndex?: number
 ):
   | {
       modifiedAst: Node<Program>
@@ -121,7 +114,12 @@ 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
   if (!generatedAxis) return new Error('Generated axis selection is missing.')
@@ -134,21 +132,30 @@ 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 VariableDeclaration = createVariableDeclaration(name, revolveCall)
-  const lastSketchNodePath =
-    orderedSketchNodePaths[orderedSketchNodePaths.length - 1]
-  let sketchIndexInBody = Number(lastSketchNodePath[1][0])
-  if (typeof sketchIndexInBody !== 'number') {
-    return new Error('expected sketchIndexInBody to be a number')
+  const name =
+    variableName ??
+    findUniqueName(clonedAst, KCL_DEFAULT_CONSTANT_PREFIXES.REVOLVE)
+  const declaration = createVariableDeclaration(name, revolveCall)
+
+  let sketchIndexInBody: number | undefined
+  // If it's an edit flow (no change on selection yet)
+  if (insertIndex) {
+    sketchIndexInBody = insertIndex
+  } else {
+    const lastSketchNodePath =
+      orderedSketchNodePaths[orderedSketchNodePaths.length - 1]
+    sketchIndexInBody = Number(lastSketchNodePath[1][0])
+    if (typeof sketchIndexInBody !== 'number') {
+      return new Error('expected sketchIndexInBody to be a number')
+    }
+
+    // If an axis was selected in KCL, find the max index to insert the revolve command
+    if (axisIndexIfAxis) {
+      sketchIndexInBody = Math.max(sketchIndexInBody, axisIndexIfAxis)
+    }
   }
 
-  // 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)
+  clonedAst.body.splice(sketchIndexInBody + 1, 0, declaration)
 
   const pathToRevolveArg: PathToNode = [
     ['body', ''],

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
-    edge: Selections
+    axis: string | undefined
+    edge: Selections | undefined
   }
   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

@@ -810,6 +810,166 @@ 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'],
+    },
+    engineCommandManager.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'],
+    },
+    engineCommandManager.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'],
+      },
+      engineCommandManager.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'],
+      },
+      engineCommandManager.artifactGraph
+    )
+    if (err(artifact)) {
+      return { reason: "Couldn't find related edge artifact" }
+    }
+
+    const codeRef = getSweepEdgeCodeRef(
+      artifact,
+      engineCommandManager.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.
@@ -882,6 +1042,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,
-          engineCommandManager.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({
@@ -1839,6 +1783,85 @@ 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
+      }
+
+      // 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,
+        engineCommandManager.artifactGraph,
+        selection.graphSelections[0]?.artifact,
+        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,
@@ -2769,9 +2792,8 @@ export const modelingMachine = setup({
         },
 
         Revolve: {
-          target: 'idle',
-          actions: ['AST revolve'],
-          reenter: false,
+          target: 'Applying revolve',
+          reenter: true,
         },
 
         Sweep: {
@@ -4025,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',