Migrate Text-to-CAD (edit) to multi-file endpoint (#6066)

* start of migrate to multi file endpoint

* get some relative path stuff sorted

* blobifying files, and making selections work with imports working

* add write to disk

* warn about big projects

* update known circular

* update snapshot

* remove log

* tweak selection filters

* Update src/components/ModelingMachineProvider.tsx

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

* fmt

* fix one thing

* typo

* raw dog form data like a fucking peasant

* remove fake data

* fmt

* steal Kevin's stuff

* good progress

* clean up

* fix writing to files when response returns

* comment the terriable code

* push fix of sorts

* better fix

* spot of clean up

* fix: Needed to support the bad request flow, the toast will hang forever, the return control flows don't dismiss a forever toast

* fix: handling more error flows by dismissing the toast

* chore: leaving a comment for a confusing workflow

* fix: trying to clean up some async logic

* fix: trying to fix a few things at once...

* fix: fixing toast success

* fix: how did this desync?

* fix: removing useless logic, we write to disk ahead of time, the continue is to say ya no problem

* fix: typo

* Change back to `spawnChild`, forego `actors` by reference

* fix: updating PR comments

* fix: found a bug with paths from rust! it is actually OS paths!

* fix: updated type still is failing tsc

* fix: the type of the machine was wrong, we always set it to at least ''

* fix: idk man

* Fix happy path test (button labels)

---------

Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>
Co-authored-by: Kevin Nadro <kevin@zoo.dev>
Co-authored-by: Frank Noirot <frankjohnson1993@gmail.com>
Co-authored-by: Kevin Nadro <nadr0@users.noreply.github.com>
Co-authored-by: Pierre Jacquier <pierre@zoo.dev>

src/components/ModelingMachineProvider.tsx

@@ -21,30 +21,19 @@ import type { Plane } from '@rust/kcl-lib/bindings/Plane'
 
 import { useAppState } from '@src/AppState'
 import { letEngineAnimateAndSyncCamAfter } from '@src/clientSideScene/CameraControls'
-import {
-  SEGMENT_BODIES,
-  getParentGroup,
-} from '@src/clientSideScene/sceneConstants'
-import type { MachineManager } from '@src/components/MachineManagerProvider'
-import { MachineManagerContext } from '@src/components/MachineManagerProvider'
-import type { SidebarType } from '@src/components/ModelingSidebar/ModelingPanes'
-import { applyConstraintIntersect } from '@src/components/Toolbar/Intersect'
-import { applyConstraintAbsDistance } from '@src/components/Toolbar/SetAbsDistance'
-import {
-  angleBetweenInfo,
-  applyConstraintAngleBetween,
-} from '@src/components/Toolbar/SetAngleBetween'
-import { applyConstraintHorzVertDistance } from '@src/components/Toolbar/SetHorzVertDistance'
 import {
   applyConstraintAngleLength,
   applyConstraintLength,
 } from '@src/components/Toolbar/setAngleLength'
+import {
+  SEGMENT_BODIES,
+  getParentGroup,
+} from '@src/clientSideScene/sceneConstants'
 import { useFileContext } from '@src/hooks/useFileContext'
 import {
   useMenuListener,
   useSketchModeMenuEnableDisable,
 } from '@src/hooks/useMenu'
-import { useNetworkContext } from '@src/hooks/useNetworkContext'
 import useStateMachineCommands from '@src/hooks/useStateMachineCommands'
 import { useKclContext } from '@src/lang/KclProvider'
 import { updateModelingState } from '@src/lang/modelingWorkflows'
@@ -70,34 +59,33 @@ import {
   getPlaneFromArtifact,
 } from '@src/lang/std/artifactGraph'
 import {
-  EngineConnectionEvents,
   EngineConnectionStateType,
+  EngineConnectionEvents,
 } from '@src/lang/std/engineConnection'
+import { err, reportRejection, trap, reject } from '@src/lib/trap'
+import { isNonNullable, platform, uuidv4 } from '@src/lib/utils'
+import { promptToEditFlow } from '@src/lib/promptToEdit'
+import type { FileMeta } from '@src/lib/types'
+import { kclEditorActor } from '@src/machines/kclEditorMachine'
+import { commandBarActor } from '@src/lib/singletons'
+import { useToken, useSettings } from '@src/lib/singletons'
+import type { IndexLoaderData } from '@src/lib/types'
 import {
-  crossProduct,
-  isCursorInSketchCommandRange,
-  updateSketchDetailsNodePaths,
-} from '@src/lang/util'
-import type {
-  KclValue,
-  PathToNode,
-  PipeExpression,
-  Program,
-  VariableDeclaration,
-} from '@src/lang/wasm'
-import { parse, recast, resultIsOk } from '@src/lang/wasm'
-import type { ModelingCommandSchema } from '@src/lib/commandBarConfigs/modelingCommandConfig'
-import { modelingMachineCommandConfig } from '@src/lib/commandBarConfigs/modelingCommandConfig'
-import {
-  EXECUTION_TYPE_MOCK,
   EXPORT_TOAST_MESSAGES,
   MAKE_TOAST_MESSAGES,
+  EXECUTION_TYPE_MOCK,
+  FILE_EXT,
 } from '@src/lib/constants'
 import { exportMake } from '@src/lib/exportMake'
 import { exportSave } from '@src/lib/exportSave'
-import { promptToEditFlow } from '@src/lib/promptToEdit'
-import type { Selections } from '@src/lib/selections'
-import { handleSelectionBatch, updateSelections } from '@src/lib/selections'
+import { isDesktop } from '@src/lib/isDesktop'
+import type { FileEntry } from '@src/lib/project'
+import type { WebContentSendPayload } from '@src/menu/channels'
+import {
+  getPersistedContext,
+  modelingMachine,
+  modelingMachineDefaultContext,
+} from '@src/machines/modelingMachine'
 import {
   codeManager,
   editorManager,
@@ -107,18 +95,39 @@ import {
   sceneEntitiesManager,
   sceneInfra,
 } from '@src/lib/singletons'
-import { err, reject, reportRejection, trap } from '@src/lib/trap'
-import type { IndexLoaderData } from '@src/lib/types'
-import { platform, uuidv4 } from '@src/lib/utils'
-import { useSettings, useToken } from '@src/lib/singletons'
-import { commandBarActor } from '@src/lib/singletons'
-import { kclEditorActor } from '@src/machines/kclEditorMachine'
+import type { MachineManager } from '@src/components/MachineManagerProvider'
+import { MachineManagerContext } from '@src/components/MachineManagerProvider'
 import {
-  getPersistedContext,
-  modelingMachine,
-  modelingMachineDefaultContext,
-} from '@src/machines/modelingMachine'
-import type { WebContentSendPayload } from '@src/menu/channels'
+  handleSelectionBatch,
+  updateSelections,
+  type Selections,
+} from '@src/lib/selections'
+import {
+  crossProduct,
+  isCursorInSketchCommandRange,
+  updateSketchDetailsNodePaths,
+} from '@src/lang/util'
+import {
+  modelingMachineCommandConfig,
+  type ModelingCommandSchema,
+} from '@src/lib/commandBarConfigs/modelingCommandConfig'
+import type {
+  KclValue,
+  PathToNode,
+  PipeExpression,
+  Program,
+  VariableDeclaration,
+} from '@src/lang/wasm'
+import { parse, recast, resultIsOk } from '@src/lang/wasm'
+import { applyConstraintHorzVertDistance } from '@src/components/Toolbar/SetHorzVertDistance'
+import {
+  angleBetweenInfo,
+  applyConstraintAngleBetween,
+} from '@src/components/Toolbar/SetAngleBetween'
+import { applyConstraintIntersect } from '@src/components/Toolbar/Intersect'
+import { applyConstraintAbsDistance } from '@src/components/Toolbar/SetAbsDistance'
+import type { SidebarType } from '@src/components/ModelingSidebar/ModelingPanes'
+import { useNetworkContext } from '@src/hooks/useNetworkContext'
 
 export const ModelingMachineContext = createContext(
   {} as {
@@ -1726,8 +1735,94 @@ export const ModelingMachineProvider = ({
           }
         ),
         'submit-prompt-edit': fromPromise(async ({ input }) => {
+          let projectFiles: FileMeta[] = [
+            {
+              type: 'kcl',
+              relPath: 'main.kcl',
+              absPath: 'main.kcl',
+              fileContents: codeManager.code,
+              execStateFileNamesIndex: 0,
+            },
+          ]
+          const execStateNameToIndexMap: { [fileName: string]: number } = {}
+          Object.entries(kclManager.execState.filenames).forEach(
+            ([index, val]) => {
+              if (val?.type === 'Local') {
+                execStateNameToIndexMap[val.value] = Number(index)
+              }
+            }
+          )
+          let basePath = ''
+          if (isDesktop() && context?.project?.children) {
+            basePath = context?.selectedDirectory?.path
+            const filePromises: Promise<FileMeta | null>[] = []
+            let uploadSize = 0
+            const recursivelyPushFilePromises = (files: FileEntry[]) => {
+              // mutates filePromises declared above, so this function definition should stay here
+              // if pulled out, it would need to be refactored.
+              for (const file of files) {
+                if (file.children !== null) {
+                  // is directory
+                  recursivelyPushFilePromises(file.children)
+                  continue
+                }
+
+                const absolutePathToFileNameWithExtension = file.path
+                const fileNameWithExtension = window.electron.path.relative(
+                  basePath,
+                  absolutePathToFileNameWithExtension
+                )
+
+                const filePromise = window.electron
+                  .readFile(absolutePathToFileNameWithExtension)
+                  .then((file): FileMeta => {
+                    uploadSize += file.byteLength
+                    const decoder = new TextDecoder('utf-8')
+                    const fileType = window.electron.path.extname(
+                      absolutePathToFileNameWithExtension
+                    )
+                    if (fileType === FILE_EXT) {
+                      return {
+                        type: 'kcl',
+                        absPath: absolutePathToFileNameWithExtension,
+                        relPath: fileNameWithExtension,
+                        fileContents: decoder.decode(file),
+                        execStateFileNamesIndex:
+                          execStateNameToIndexMap[
+                            absolutePathToFileNameWithExtension
+                          ],
+                      }
+                    }
+                    const blob = new Blob([file], {
+                      type: 'application/octet-stream',
+                    })
+                    return {
+                      type: 'other',
+                      relPath: fileNameWithExtension,
+                      data: blob,
+                    }
+                  })
+                  .catch((e) => {
+                    console.error('error reading file', e)
+                    return null
+                  })
+
+                filePromises.push(filePromise)
+              }
+            }
+            recursivelyPushFilePromises(context?.project?.children)
+            projectFiles = (await Promise.all(filePromises)).filter(
+              isNonNullable
+            )
+            const MB20 = 2 ** 20 * 20
+            if (uploadSize > MB20) {
+              toast.error(
+                'Your project exceeds 20Mb, this will slow down Text-to-CAD\nPlease remove any unnecessary files'
+              )
+            }
+          }
           return await promptToEditFlow({
-            code: codeManager.code,
+            projectFiles,
             prompt: input.prompt,
             selections: input.selection,
             token,

src/components/ToastTextToCad.tsx

@@ -1,6 +1,6 @@
 import type {
-  TextToCadIteration_type,
   TextToCad_type,
+  TextToCadMultiFileIteration_type,
 } from '@kittycad/lib/dist/types/src/models'
 import { useCallback, useEffect, useRef } from 'react'
 import toast from 'react-hot-toast'
@@ -29,9 +29,17 @@ import { PATHS } from '@src/lib/paths'
 import { codeManager, kclManager, systemIOActor } from '@src/lib/singletons'
 import { sendTelemetry } from '@src/lib/textToCadTelemetry'
 import { reportRejection } from '@src/lib/trap'
+import {
+  SystemIOMachineEvents,
+  waitForIdleState,
+} from '@src/machines/systemIO/utils'
+import {
+  useProjectDirectoryPath,
+  useRequestedProjectName,
+} from '@src/machines/systemIO/hooks'
 import { commandBarActor } from '@src/lib/singletons'
-import { SystemIOMachineEvents } from '@src/machines/systemIO/utils'
-import { useProjectDirectoryPath } from '@src/machines/systemIO/hooks'
+import type { FileMeta } from '@src/lib/types'
+import type { RequestedKCLFile } from '@src/machines/systemIO/utils'
 
 const CANVAS_SIZE = 128
 const PROMPT_TRUNCATE_LENGTH = 128
@@ -478,14 +486,17 @@ export function ToastPromptToEditCadSuccess({
   toastId,
   token,
   data,
-  oldCode,
+  oldCodeWebAppOnly: oldCode,
+  oldFiles,
 }: {
   toastId: string
-  oldCode: string
-  data: TextToCadIteration_type
+  oldCodeWebAppOnly: string
+  oldFiles: FileMeta[]
+  data: TextToCadMultiFileIteration_type
   token?: string
 }) {
   const modelId = data.id
+  const requestedProjectName = useRequestedProjectName()
 
   return (
     <div className="flex gap-4 min-w-80">
@@ -514,13 +525,37 @@ export function ToastPromptToEditCadSuccess({
             data-negative-button={'reject'}
             name={'Reject'}
             onClick={() => {
-              sendTelemetry(modelId, 'rejected', token).catch(reportRejection)
-              codeManager.updateCodeEditor(oldCode)
-              kclManager.executeCode().catch(reportRejection)
-              toast.dismiss(toastId)
+              void (async () => {
+                sendTelemetry(modelId, 'rejected', token).catch(reportRejection)
+                // revert to before the prompt-to-edit
+                if (isDesktop()) {
+                  const requestedFiles: RequestedKCLFile[] = []
+                  for (const file of oldFiles) {
+                    if (file.type !== 'kcl') {
+                      // only need to write the kcl files
+                      // as the endpoint would have never overwritten other file types
+                      continue
+                    }
+                    requestedFiles.push({
+                      requestedCode: file.fileContents,
+                      requestedFileName: file.relPath,
+                      requestedProjectName: requestedProjectName.name,
+                    })
+                  }
+                  toast.dismiss(toastId)
+                  await writeOverFilesAndExecute({
+                    requestedFiles: requestedFiles,
+                    projectName: requestedProjectName.name,
+                  })
+                } else {
+                  codeManager.updateCodeEditor(oldCode)
+                  kclManager.executeCode().catch(reportRejection)
+                  toast.dismiss(toastId)
+                }
+              })()
             }}
           >
-            {'Reject'}
+            {'Revert'}
           </ActionButton>
 
           <ActionButton
@@ -532,23 +567,38 @@ export function ToastPromptToEditCadSuccess({
             onClick={() => {
               sendTelemetry(modelId, 'accepted', token).catch(reportRejection)
               toast.dismiss(toastId)
-
-              // Write new content to disk since they have accepted.
-              codeManager
-                .writeToFile()
-                .then(() => {
-                  // no-op
-                })
-                .catch((e) => {
-                  console.error('Failed to save prompt-to-edit to disk')
-                  console.error(e)
-                })
+              /**
+               * NO OP. Do not rewrite code to disk, we already do this ahead of time. This will dismiss the toast.
+               * All of the files were already written! Don't rewrite the current code editor.
+               * If this prompt to edit makes 5 new files, the code manager is only watching 1 of these files, why
+               * would it rewrite the current file selected when this is completed?
+               */
             }}
           >
-            Accept
+            Continue
           </ActionButton>
         </div>
       </div>
     </div>
   )
 }
+
+export const writeOverFilesAndExecute = async ({
+  requestedFiles,
+  projectName,
+}: {
+  requestedFiles: RequestedKCLFile[]
+  projectName: string
+}) => {
+  systemIOActor.send({
+    type: SystemIOMachineEvents.bulkCreateKCLFilesAndNavigateToProject,
+    data: {
+      files: requestedFiles,
+      requestedProjectName: projectName,
+      override: true,
+    },
+  })
+
+  // to await the result of the send event above
+  await waitForIdleState({ systemIOActor })
+}