diff --git a/src/lang/modifyAst.ts b/src/lang/modifyAst.ts
index e1627f699..c29845d1f 100644
--- a/src/lang/modifyAst.ts
+++ b/src/lang/modifyAst.ts
@@ -603,35 +603,30 @@ export function addOffsetPlane({
  */
 export function addShell({
   node,
-  selection,
+  extrudeNode,
+  selectedArtifact,
   thickness,
 }: {
   node: Node<Program>
-  selection: Selection
+  extrudeNode: {
+    node: VariableDeclarator
+    shallowPath: PathToNode
+    deepPath: PathToNode
+  }
+  selectedArtifact: Artifact
   thickness: Expr
-}): { modifiedAst: Node<Program>; pathToNode: PathToNode } | Error {
+}): { modifiedAst: Node<Program>; pathToNode: PathToNode } {
   const modifiedAst = structuredClone(node)
   const name = findUniqueName(node, KCL_DEFAULT_CONSTANT_PREFIXES.SHELL)
-  // TODO: change to what's needed for shell
-  console.log('selection, thickness', selection, thickness)
-  const baseNode = getNodeFromPath<VariableDeclarator>(
-    node,
-    selection.codeRef.pathToNode,
-    'VariableDeclarator'
-  )
-  console.log('baseNode', baseNode)
-  if (err(baseNode)) return baseNode
-  console.log("selection.artifact['subType']", selection.artifact['subType'])
   const shell = createCallExpressionStdLib('shell', [
     createObjectExpression({
       faces: createArrayExpression([
-        createLiteral(selection.artifact['subType']),
+        // TODO: make typescript happy
+        createLiteral(selectedArtifact['subType']),
       ]),
       thickness,
     }),
-    // createIdentifier(baseNode.node.id.name),
-    // TODO: replace with the query from sketch to extrude
-    createIdentifier('extrude001'),
+    createIdentifier(extrudeNode.node.id.name),
   ])
   const declaration = createVariableDeclaration(name, shell)
   modifiedAst.body.push(declaration)
diff --git a/src/machines/modelingMachine.ts b/src/machines/modelingMachine.ts
index e6bfab27e..f018ce5c6 100644
--- a/src/machines/modelingMachine.ts
+++ b/src/machines/modelingMachine.ts
@@ -52,6 +52,7 @@ import {
   applyEdgeTreatmentToSelection,
   EdgeTreatmentType,
   FilletParameters,
+  getPathToExtrudeForSegmentSelection,
 } from 'lang/modifyAst/addEdgeTreatment'
 import { getNodeFromPath } from '../lang/queryAst'
 import {
@@ -1594,12 +1595,10 @@ export const modelingMachine = setup({
         // Extract inputs
         const ast = kclManager.ast
         const { selection, thickness } = input
-
-        // TODO: extract selection
-        console.log('selection', selection)
         const graphSelection = selection.graphSelections[0]
-        if (!(graphSelection && graphSelection instanceof Object))
+        if (!(graphSelection && graphSelection instanceof Object)) {
           return trap('No plane selected')
+        }
 
         // Insert the thickness variable if it exists
         if (
@@ -1617,36 +1616,63 @@ export const modelingMachine = setup({
         }
 
         // Get the shell from the selection
-        if (err(selection.graphSelections[0])) return
+        const sketchNode = getNodeFromPath<VariableDeclarator>(
+          ast,
+          graphSelection.codeRef.pathToNode,
+          'VariableDeclarator'
+        )
+        if (err(sketchNode)) {
+          return sketchNode
+        }
+
+        const clonedAstForGetExtrude = structuredClone(ast)
+        const extrudeLookupResult = getPathToExtrudeForSegmentSelection(
+          clonedAstForGetExtrude,
+          selection,
+          engineCommandManager.artifactGraph
+        )
+        if (err(extrudeLookupResult)) {
+          return new Error("Couldn't find extrude")
+        }
+
+        const extrudeNode = getNodeFromPath<VariableDeclarator>(
+          ast,
+          extrudeLookupResult.pathToExtrudeNode,
+          'VariableDeclarator'
+        )
+        if (err(extrudeNode)) {
+          return extrudeNode
+        }
+
+        const selectedArtifact = graphSelection.artifact
+        if (!selectedArtifact || !selectedArtifact) {
+          return new Error('Bad artifact')
+        }
 
         const shellResult = addShell({
           node: ast,
-          selection: graphSelection,
+          extrudeNode,
+          selectedArtifact,
           thickness:
             'variableName' in thickness
               ? thickness.variableIdentifierAst
               : thickness.valueAst,
         })
 
-        if (trap(shellResult)) return
-        try {
-          const updateAstResult = await kclManager.updateAst(
-            shellResult.modifiedAst,
-            true,
-            {
-              focusPath: [shellResult.pathToNode],
-            }
-          )
-
-          await codeManager.updateEditorWithAstAndWriteToFile(
-            updateAstResult.newAst
-          )
-
-          if (updateAstResult?.selections) {
-            editorManager.selectRange(updateAstResult?.selections)
+        const updateAstResult = await kclManager.updateAst(
+          shellResult.modifiedAst,
+          true,
+          {
+            focusPath: [shellResult.pathToNode],
           }
-        } catch (e) {
-          console.error(e)
+        )
+
+        await codeManager.updateEditorWithAstAndWriteToFile(
+          updateAstResult.newAst
+        )
+
+        if (updateAstResult?.selections) {
+          editorManager.selectRange(updateAstResult?.selections)
         }
       }
     ),