Fix race condition with setExperimentalFeatures

e2e/playwright/point-click.spec.ts

@@ -3638,7 +3638,7 @@ solid001 = extrude(sketch001, length = 5)`
     toolbar,
     cmdBar,
   }) => {
-    const initialCode = `@settings(defaultLengthUnit = in, experimentalFeatures = allow)
+    const initialCode = `@settings(defaultLengthUnit = in)
 sketch001 = startSketchOn(XZ)
   |> circle(center = [0, 0], radius = 30)
 extrude001 = extrude(sketch001, length = 30)
@@ -3892,6 +3892,7 @@ extrude001 = extrude(sketch001, length = 30)
       await test.step('Submit and verify all parameters', async () => {
         await cmdBar.progressCmdBar()
         await scene.settled(cmdBar)
+        await editor.expectEditor.toContain('experimentalFeatures = allow')
         await editor.expectEditor.toContain('gdt::flatness(')
         await editor.expectEditor.toContain('faces = [capEnd001]')
         await editor.expectEditor.toContain('tolerance = 0.1mm')

src/components/ExperimentalFeaturesMenu.tsx

@@ -1,13 +1,22 @@
 import { Popover } from '@headlessui/react'
 import toast from 'react-hot-toast'
 
-import { DEFAULT_EXPERIMENTAL_FEATURES } from '@src/lib/constants'
-import { kclManager } from '@src/lib/singletons'
+import {
+  DEFAULT_EXPERIMENTAL_FEATURES,
+  EXECUTION_TYPE_REAL,
+} from '@src/lib/constants'
+import {
+  codeManager,
+  editorManager,
+  kclManager,
+  rustContext,
+} from '@src/lib/singletons'
 import { err, reportRejection } from '@src/lib/trap'
 import { CustomIcon } from '@src/components/CustomIcon'
 import { warningLevels } from '@src/lib/settings/settingsTypes'
 import type { WarningLevel } from '@rust/kcl-lib/bindings/WarningLevel'
-import { setExperimentalFeatures } from '@src/lib/kclHelpers'
+import { setExperimentalFeatures } from '@src/lang/modifyAst/settings'
+import { updateModelingState } from '@src/lang/modelingWorkflows'
 
 export function ExperimentalFeaturesMenu() {
   const currentLevel: WarningLevel =
@@ -41,20 +50,35 @@ export function ExperimentalFeaturesMenu() {
                     <button
                       className="flex items-center gap-2 m-0 py-1.5 px-2 cursor-pointer hover:bg-chalkboard-20 dark:hover:bg-chalkboard-80 border-none text-left"
                       onClick={() => {
-                        setExperimentalFeatures(level)
-                          .then((result) => {
-                            if (err(result)) {
-                              toast.error(
-                                `Failed to set file experimental features level: ${result.message}`
-                              )
-                              return
-                            }
-
-                            toast.success(
-                              `Updated file experimental features level to ${level.type}`
-                            )
+                        const newAst = setExperimentalFeatures(
+                          codeManager.code,
+                          level
+                        )
+                        if (err(newAst)) {
+                          toast.error(
+                            `Failed to set file experimental features level: ${newAst.message}`
+                          )
+                        } else {
+                          updateModelingState(newAst, EXECUTION_TYPE_REAL, {
+                            kclManager,
+                            editorManager,
+                            codeManager,
+                            rustContext,
                           })
-                          .catch(reportRejection)
+                            .then((result) => {
+                              if (err(result)) {
+                                toast.error(
+                                  `Failed to set file experimental features level: ${result.message}`
+                                )
+                                return
+                              }
+
+                              toast.success(
+                                `Updated file experimental features level to ${level.type}`
+                              )
+                            })
+                            .catch(reportRejection)
+                        }
                         close()
                       }}
                     >

src/lang/modifyAst/settings.spec.ts

@@ -0,0 +1,33 @@
+import { recast } from '@src/lang/wasm'
+import { err } from '@src/lib/trap'
+import { join } from 'path'
+import { loadAndInitialiseWasmInstance } from '@src/lang/wasmUtilsNode'
+import { setExperimentalFeatures } from '@src/lang/modifyAst/settings'
+const WASM_PATH = join(process.cwd(), 'public/kcl_wasm_lib_bg.wasm')
+
+describe('settings.spec.ts', () => {
+  describe('Testing setExperimentalFeatures', () => {
+    it('should add the annotation and set the flag if not present', async () => {
+      const instance = await loadAndInitialiseWasmInstance(WASM_PATH)
+      const newAst = setExperimentalFeatures('', { type: 'Allow' }, instance)
+      if (err(newAst)) {
+        throw newAst
+      }
+
+      const newCode = recast(newAst, instance)
+      expect(newCode).toBe(`@settings(experimentalFeatures = allow)\n`)
+    })
+
+    it('should set the flag if the annotation exists', async () => {
+      const instance = await loadAndInitialiseWasmInstance(WASM_PATH)
+      const code = `@settings(experimentalFeatures = warn)\n`
+      const newAst = setExperimentalFeatures(code, { type: 'Deny' }, instance)
+      if (err(newAst)) {
+        throw newAst
+      }
+
+      const newCode = recast(newAst, instance)
+      expect(newCode).toBe(`@settings(experimentalFeatures = deny)\n`)
+    })
+  })
+})

src/lang/modifyAst/settings.ts

@@ -0,0 +1,27 @@
+import type { Node } from '@rust/kcl-lib/bindings/Node'
+import type { WarningLevel } from '@rust/kcl-lib/bindings/WarningLevel'
+import { err } from '@src/lib/trap'
+import { changeExperimentalFeatures, parse, type Program } from '@src/lang/wasm'
+import type { ModuleType } from '@src/lib/wasm_lib_wrapper'
+
+export function setExperimentalFeatures(
+  code: string,
+  level: WarningLevel,
+  instance?: ModuleType
+): Node<Program> | Error {
+  const newCode = changeExperimentalFeatures(code, level, instance)
+  if (err(newCode)) {
+    return newCode
+  }
+
+  const result = parse(newCode, instance)
+  if (err(result)) {
+    return result
+  }
+
+  if (result.program === null) {
+    return new Error('Empty program returned')
+  }
+
+  return result.program
+}

src/lang/wasm.ts

@@ -816,10 +816,14 @@ export function changeDefaultUnits(
  */
 export function changeExperimentalFeatures(
   kcl: string,
-  warningLevel: WarningLevel | null = null
+  warningLevel: WarningLevel | null = null,
+  instance?: ModuleType
 ): string | Error {
   try {
-    return change_experimental_features(kcl, JSON.stringify(warningLevel))
+    const level = JSON.stringify(warningLevel)
+    return instance
+      ? instance.change_experimental_features(kcl, level)
+      : change_experimental_features(kcl, level)
   } catch (e) {
     console.error('Caught error changing kcl settings', e)
     return new Error('Caught error changing kcl settings', { cause: e })

src/lib/kclCommands.ts

@@ -32,7 +32,7 @@ import { getNodeFromPath } from '@src/lang/queryAst'
 import type { Node } from '@rust/kcl-lib/bindings/Node'
 import { getVariableDeclaration } from '@src/lang/queryAst/getVariableDeclaration'
 import { getNodePathFromSourceRange } from '@src/lang/queryAstNodePathUtils'
-import { setExperimentalFeatures } from '@src/lib/kclHelpers'
+import { setExperimentalFeatures } from '@src/lang/modifyAst/settings'
 
 interface KclCommandConfig {
   // TODO: find a different approach that doesn't require
@@ -128,14 +128,31 @@ export function kclCommands(commandProps: KclCommandConfig): Command[] {
       },
       onSubmit: (data) => {
         if (typeof data === 'object' && 'level' in data) {
-          setExperimentalFeatures({ type: data.level })
+          const newAst = setExperimentalFeatures(codeManager.code, {
+            type: data.level,
+          })
+          if (err(newAst)) {
+            toast.error(
+              `Failed to set file experimental features level: ${newAst.message}`
+            )
+            return
+          }
+          updateModelingState(newAst, EXECUTION_TYPE_REAL, {
+            kclManager,
+            editorManager,
+            codeManager,
+            rustContext,
+          })
             .then((result) => {
               if (err(result)) {
-                reportRejection(result)
+                toast.error(
+                  `Failed to set file experimental features level: ${result.message}`
+                )
                 return
               }
+
               toast.success(
-                `Updated experimental features level to ${data.level}`
+                `Updated file experimental features level to ${data.level}`
               )
             })
             .catch(reportRejection)

src/lib/kclHelpers.ts

@@ -1,15 +1,13 @@
-import type { WarningLevel } from '@rust/kcl-lib/bindings/WarningLevel'
 import { executeAstMock } from '@src/lang/langHelpers'
 import {
   type SourceRange,
   type KclValue,
   formatNumberValue,
   parse,
   resultIsOk,
-  changeExperimentalFeatures,
 } from '@src/lang/wasm'
 import type { KclExpression } from '@src/lib/commandTypes'
-import { codeManager, kclManager, rustContext } from '@src/lib/singletons'
+import { rustContext } from '@src/lib/singletons'
 import { err } from '@src/lib/trap'
 import type { ModuleType } from '@src/lib/wasm_lib_wrapper'
 import type RustContext from '@src/lib/rustContext'
@@ -153,15 +151,3 @@ export async function stringToKclExpression(
 export function getStringValue(code: string, range: SourceRange): string {
   return code.slice(range[0], range[1]).replaceAll(`'`, ``).replaceAll(`"`, ``)
 }
-
-export async function setExperimentalFeatures(
-  level: WarningLevel
-): Promise<void | Error> {
-  const newCode = changeExperimentalFeatures(codeManager.code, level)
-  if (err(newCode)) {
-    return new Error(`Failed to set experimental features: ${newCode.message}`)
-  }
-  codeManager.updateCodeStateEditor(newCode)
-  await codeManager.writeToFile()
-  await kclManager.executeCode()
-}

src/machines/modelingMachine.ts

@@ -163,7 +163,7 @@ import {
   type EquipTool,
   sketchSolveMachine,
 } from '@src/machines/sketchSolve/sketchSolveMode'
-import { setExperimentalFeatures } from '@src/lib/kclHelpers'
+import { setExperimentalFeatures } from '@src/lang/modifyAst/settings'
 import type CodeManager from '@src/lang/codeManager'
 import type EditorManager from '@src/editor/manager'
 import type { KclManager } from '@src/lang/KclSingleton'
@@ -3780,19 +3780,23 @@ export const modelingMachine = setup({
           return Promise.reject(new Error(NO_INPUT_PROVIDED_MESSAGE))
         }
 
-        // Remove once it isn't experimental anymore
+        // Remove once this command isn't experimental anymore
+        let astWithNewSetting: Node<Program> | undefined
         if (kclManager.fileSettings.experimentalFeatures?.type !== 'Allow') {
-          const result = await setExperimentalFeatures({ type: 'Allow' })
-          if (err(result)) {
-            return Promise.reject(result)
+          const ast = setExperimentalFeatures(codeManager.code, {
+            type: 'Allow',
+          })
+          if (err(ast)) {
+            return Promise.reject(ast)
           }
+
+          astWithNewSetting = ast
         }
 
-        const { ast, artifactGraph } = kclManager
         const result = addFlatnessGdt({
           ...input,
-          ast,
-          artifactGraph,
+          ast: astWithNewSetting ?? kclManager.ast,
+          artifactGraph: kclManager.artifactGraph,
         })
         if (err(result)) {
           return Promise.reject(result)