Prompt to edit (#4830)

* initial plumbing for getting the new option into the cmd-bar

* start of prompt edit

* update AI poll

* add spinner

* more prompt engineering

* add success toast, allowing user's to reject code

* select code that changed in prompt to edit

* selection in scene should not disappear when opening prompt cmd

* tweak

* fmt

* add tests

* some clean up

* clean up

* fix tests

e2e/playwright/fixtures/cmdBarFixture.ts

@@ -1,4 +1,4 @@
-import type { Page } from '@playwright/test'
+import type { Page, Locator } from '@playwright/test'
 import { expect } from '@playwright/test'
 
 type CmdBarSerialised =
@@ -26,9 +26,11 @@ type CmdBarSerialised =
 
 export class CmdBarFixture {
   public page: Page
+  cmdBarOpenBtn!: Locator
 
   constructor(page: Page) {
     this.page = page
+    this.cmdBarOpenBtn = page.getByTestId('command-bar-open-button')
   }
   reConstruct = (page: Page) => {
     this.page = page
@@ -116,4 +118,21 @@ export class CmdBarFixture {
       await this.page.keyboard.press('Enter')
     }
   }
+
+  openCmdBar = async (selectCmd?: 'promptToEdit') => {
+    // TODO why does this button not work in electron tests?
+    // await this.cmdBarOpenBtn.click()
+    await this.page.keyboard.down('ControlOrMeta')
+    await this.page.keyboard.press('KeyK')
+    await this.page.keyboard.up('ControlOrMeta')
+    await expect(this.page.getByPlaceholder('Search commands')).toBeVisible()
+    if (selectCmd === 'promptToEdit') {
+      const promptEditCommand = this.page.getByText(
+        'Use Zoo AI to edit your kcl'
+      )
+      await expect(promptEditCommand.first()).toBeVisible()
+      await promptEditCommand.first().scrollIntoViewIfNeeded()
+      await promptEditCommand.first().click()
+    }
+  }
 }

e2e/playwright/prompt-to-edit.spec.ts

@@ -0,0 +1,190 @@
+import { test, expect } from './zoo-test'
+
+/* eslint-disable jest/no-conditional-expect */
+
+const file = `sketch001 = startSketchOn('XZ')
+profile001 = startProfileAt([57.81, 250.51], sketch001)
+  |> line([121.13, 56.63], %, $seg02)
+  |> line([83.37, -34.61], %, $seg01)
+  |> line([19.66, -116.4], %)
+  |> line([-221.8, -41.69], %)
+  |> lineTo([profileStartX(%), profileStartY(%)], %)
+  |> close(%)
+extrude001 = extrude(200, profile001)
+sketch002 = startSketchOn('XZ')
+  |> startProfileAt([-73.64, -42.89], %)
+  |> xLine(173.71, %)
+  |> line([-22.12, -94.4], %)
+  |> xLine(-156.98, %)
+  |> lineTo([profileStartX(%), profileStartY(%)], %)
+  |> close(%)
+extrude002 = extrude(50, sketch002)
+sketch003 = startSketchOn('XY')
+  |> startProfileAt([52.92, 157.81], %)
+  |> angledLine([0, 176.4], %, $rectangleSegmentA001)
+  |> angledLine([
+       segAng(rectangleSegmentA001) - 90,
+       53.4
+     ], %, $rectangleSegmentB001)
+  |> angledLine([
+       segAng(rectangleSegmentA001),
+       -segLen(rectangleSegmentA001)
+     ], %, $rectangleSegmentC001)
+  |> lineTo([profileStartX(%), profileStartY(%)], %)
+  |> close(%)
+extrude003 = extrude(20, sketch003)
+`
+
+test.describe('Check the happy path, for basic changing color', () => {
+  const cases = [
+    {
+      desc: 'User accepts change',
+      shouldReject: false,
+    },
+    {
+      desc: 'User rejects change',
+      shouldReject: true,
+    },
+  ] as const
+  for (const { desc, shouldReject } of cases) {
+    test(`${desc}`, async ({
+      context,
+      homePage,
+      cmdBar,
+      editor,
+      page,
+      scene,
+    }) => {
+      await context.addInitScript((file) => {
+        localStorage.setItem('persistCode', file)
+      }, file)
+      await homePage.goToModelingScene()
+
+      const body1CapCoords = { x: 571, y: 351 }
+      const greenCheckCoords = { x: 565, y: 345 }
+      const body2WallCoords = { x: 609, y: 153 }
+      const [clickBody1Cap] = scene.makeMouseHelpers(
+        body1CapCoords.x,
+        body1CapCoords.y
+      )
+      const yellow: [number, number, number] = [179, 179, 131]
+      const green: [number, number, number] = [108, 152, 75]
+      const notGreen: [number, number, number] = [132, 132, 132]
+      const body2NotGreen: [number, number, number] = [88, 88, 88]
+      const submittingToast = page.getByText('Submitting to Edit AI...')
+      const successToast = page.getByText('Prompt to edit successful')
+      const acceptBtn = page.getByRole('button', { name: 'checkmark Accept' })
+      const rejectBtn = page.getByRole('button', { name: 'close Reject' })
+
+      await test.step('wait for scene to load select body and check selection came through', async () => {
+        await scene.expectPixelColor([134, 134, 134], body1CapCoords, 15)
+        await clickBody1Cap()
+        await scene.expectPixelColor(yellow, body1CapCoords, 20)
+        await editor.expectState({
+          highlightedCode: '',
+          activeLines: ['|>startProfileAt([-73.64,-42.89],%)'],
+          diagnostics: [],
+        })
+      })
+
+      await test.step('fire off edit prompt', async () => {
+        await cmdBar.openCmdBar('promptToEdit')
+        // being specific about the color with a hex means asserting pixel color is more stable
+        await page
+          .getByTestId('cmd-bar-arg-value')
+          .fill('make this neon green please, use #39FF14')
+        await page.waitForTimeout(100)
+        await cmdBar.progressCmdBar()
+        await expect(submittingToast).toBeVisible()
+        await expect(submittingToast).not.toBeVisible({ timeout: 2 * 60_000 }) // can take a while
+        await expect(successToast).toBeVisible()
+      })
+
+      await test.step('verify initial change', async () => {
+        await scene.expectPixelColor(green, greenCheckCoords, 15)
+        await scene.expectPixelColor(body2NotGreen, body2WallCoords, 15)
+        await editor.expectEditor.toContain('appearance({')
+      })
+
+      if (!shouldReject) {
+        await test.step('check accept works and can be "undo"ed', async () => {
+          await acceptBtn.click()
+          await expect(successToast).not.toBeVisible()
+
+          await scene.expectPixelColor(green, greenCheckCoords, 15)
+          await editor.expectEditor.toContain('appearance({')
+
+          // ctrl-z works after accepting
+          await page.keyboard.down('ControlOrMeta')
+          await page.keyboard.press('KeyZ')
+          await page.keyboard.up('ControlOrMeta')
+          await editor.expectEditor.not.toContain('appearance({')
+          await scene.expectPixelColor(notGreen, greenCheckCoords, 15)
+        })
+      } else {
+        await test.step('check reject works', async () => {
+          await rejectBtn.click()
+          await expect(successToast).not.toBeVisible()
+
+          await scene.expectPixelColor(notGreen, greenCheckCoords, 15)
+          await editor.expectEditor.not.toContain('appearance({')
+        })
+      }
+    })
+  }
+})
+
+test.describe('bad path', () => {
+  test(`bad edit prompt`, async ({
+    context,
+    homePage,
+    cmdBar,
+    editor,
+    toolbar,
+    page,
+    scene,
+  }) => {
+    await context.addInitScript((file) => {
+      localStorage.setItem('persistCode', file)
+    }, file)
+    await homePage.goToModelingScene()
+
+    const body1CapCoords = { x: 571, y: 351 }
+    const [clickBody1Cap] = scene.makeMouseHelpers(
+      body1CapCoords.x,
+      body1CapCoords.y
+    )
+    const yellow: [number, number, number] = [179, 179, 131]
+    const submittingToast = page.getByText('Submitting to Edit AI...')
+    const failToast = page.getByText(
+      'Failed to edit your KCL code, please try again with a different prompt or selection'
+    )
+
+    await test.step('wait for scene to load and select body', async () => {
+      await scene.expectPixelColor([134, 134, 134], body1CapCoords, 15)
+
+      await clickBody1Cap()
+      await scene.expectPixelColor(yellow, body1CapCoords, 20)
+
+      await editor.expectState({
+        highlightedCode: '',
+        activeLines: ['|>startProfileAt([-73.64,-42.89],%)'],
+        diagnostics: [],
+      })
+    })
+
+    await test.step('fire of bad prompt', async () => {
+      await cmdBar.openCmdBar('promptToEdit')
+      await page
+        .getByTestId('cmd-bar-arg-value')
+        .fill('ansheusha asnthuatshoeuhtaoetuhthaeu laughs in dvorak')
+      await page.waitForTimeout(100)
+      await cmdBar.progressCmdBar()
+      await expect(submittingToast).toBeVisible()
+    })
+    await test.step('check fail toast appeared', async () => {
+      await expect(submittingToast).not.toBeVisible({ timeout: 2 * 60_000 }) // can take a while
+      await expect(failToast).toBeVisible()
+    })
+  })
+})

package.json

@@ -39,6 +39,7 @@
     "chokidar": "^4.0.1",
     "codemirror": "^6.0.1",
     "decamelize": "^6.0.0",
+    "diff": "^7.0.0",
     "electron-updater": "6.3.0",
     "fuse.js": "^7.0.0",
     "html2canvas-pro": "^1.5.8",
@@ -154,6 +155,7 @@
     "@testing-library/jest-dom": "^5.14.1",
     "@testing-library/react": "^15.0.2",
     "@types/d3-force": "^3.0.10",
+    "@types/diff": "^6.0.0",
     "@types/electron": "^1.6.10",
     "@types/isomorphic-fetch": "^0.0.39",
     "@types/minimist": "^1.2.5",

src/components/CommandBar/CommandBarSelectionInput.tsx

@@ -78,7 +78,7 @@ function CommandBarSelectionInput({
     return () => {
       toSync(() => {
         const promises = [
-          new Promise(() => kclManager.defaultSelectionFilter()),
+          new Promise(() => kclManager.defaultSelectionFilter(selection)),
         ]
         if (!kclManager._isAstEmpty(kclManager.ast)) {
           promises.push(kclManager.hidePlanes())

src/components/ModelingMachineProvider.tsx

@@ -87,6 +87,7 @@ import { useFileContext } from 'hooks/useFileContext'
 import { uuidv4 } from 'lib/utils'
 import { IndexLoaderData } from 'lib/types'
 import { Node } from 'wasm-lib/kcl/bindings/Node'
+import { promptToEditFlow } from 'lib/promptToEdit'
 
 type MachineContext<T extends AnyStateMachine> = {
   state: StateFrom<T>
@@ -1104,6 +1105,15 @@ export const ModelingMachineProvider = ({
             }
           }
         ),
+        'submit-prompt-edit': fromPromise(async ({ input }) => {
+          return await promptToEditFlow({
+            code: codeManager.code,
+            prompt: input.prompt,
+            selections: input.selection,
+            token,
+            artifactGraph: engineCommandManager.artifactGraph,
+          })
+        }),
       },
     }),
     {

src/components/ToastTextToCad.tsx

@@ -4,7 +4,10 @@ import { useFileContext } from 'hooks/useFileContext'
 import { isDesktop } from 'lib/isDesktop'
 import { PATHS } from 'lib/paths'
 import toast from 'react-hot-toast'
-import { TextToCad_type } from '@kittycad/lib/dist/types/src/models'
+import {
+  TextToCad_type,
+  TextToCadIteration_type,
+} from '@kittycad/lib/dist/types/src/models'
 import { useCallback, useEffect, useRef, useState } from 'react'
 import {
   Box3,
@@ -29,6 +32,7 @@ import { commandBarMachine } from 'machines/commandBarMachine'
 import { EventFrom } from 'xstate'
 import { fileMachine } from 'machines/fileMachine'
 import { reportRejection } from 'lib/trap'
+import { codeManager, kclManager } from 'lib/singletons'
 
 const CANVAS_SIZE = 128
 const PROMPT_TRUNCATE_LENGTH = 128
@@ -411,3 +415,69 @@ function traverseSceneToStyleObjects({
     }
   })
 }
+
+export function ToastPromptToEditCadSuccess({
+  toastId,
+  token,
+  data,
+  oldCode,
+}: {
+  toastId: string
+  oldCode: string
+  data: TextToCadIteration_type
+  token?: string
+}) {
+  const modelId = data.id
+
+  return (
+    <div className="flex gap-4 min-w-80">
+      <div className="flex flex-col justify-between gap-6">
+        <section>
+          <h2>Prompt to edit successful</h2>
+          <p className="text-sm text-chalkboard-70 dark:text-chalkboard-30">
+            Prompt: "
+            {data?.prompt && data?.prompt?.length > PROMPT_TRUNCATE_LENGTH
+              ? data.prompt.slice(0, PROMPT_TRUNCATE_LENGTH) + '...'
+              : data.prompt}
+            "
+          </p>
+          <p>Do you want to keep the change?</p>
+        </section>
+        <div className="flex justify-between gap-8">
+          <ActionButton
+            Element="button"
+            iconStart={{
+              icon: 'close',
+            }}
+            data-negative-button={'reject'}
+            name={'Reject'}
+            onClick={() => {
+              // TODO add telemetry when we know how sendTelemetry is setup for /user/text-to-cad/
+              // sendTelemetry(modelId, 'rejected', token).catch(reportRejection)
+              codeManager.updateCodeEditor(oldCode)
+              kclManager.executeCode().catch(reportRejection)
+              toast.dismiss(toastId)
+            }}
+          >
+            {'Reject'}
+          </ActionButton>
+
+          <ActionButton
+            Element="button"
+            iconStart={{
+              icon: 'checkmark',
+            }}
+            name="Accept"
+            onClick={() => {
+              // TODO add telemetry when we know how sendTelemetry is setup for /user/text-to-cad/
+              sendTelemetry(modelId, 'accepted', token).catch(reportRejection)
+              toast.dismiss(toastId)
+            }}
+          >
+            Accept
+          </ActionButton>
+        </div>
+      </div>
+    </div>
+  )
+}

src/lang/KclSingleton.ts

@@ -1,5 +1,5 @@
 import { executeAst, lintAst } from 'lang/langHelpers'
-import { Selections } from 'lib/selections'
+import { handleSelectionBatch, Selections } from 'lib/selections'
 import {
   KCLError,
   complilationErrorsToDiagnostics,
@@ -28,7 +28,10 @@ import { codeManager, editorManager, sceneInfra } from 'lib/singletons'
 import { Diagnostic } from '@codemirror/lint'
 import { markOnce } from 'lib/performance'
 import { Node } from 'wasm-lib/kcl/bindings/Node'
-import { EntityType_type } from '@kittycad/lib/dist/types/src/models'
+import {
+  EntityType_type,
+  ModelingCmdReq_type,
+} from '@kittycad/lib/dist/types/src/models'
 
 interface ExecuteArgs {
   ast?: Node<Program>
@@ -311,7 +314,7 @@ export class KclManager {
     // Do not send send scene commands if the program was interrupted, go to clean up
     if (!isInterrupted) {
       this.addDiagnostics(await lintAst({ ast: ast }))
-      setSelectionFilterToDefault(execState.memory, this.engineCommandManager)
+      setSelectionFilterToDefault(this.engineCommandManager)
 
       if (args.zoomToFit) {
         let zoomObjectId: string | undefined = ''
@@ -603,8 +606,8 @@ export class KclManager {
     return Promise.all(thePromises)
   }
   /** TODO: this function is hiding unawaited asynchronous work */
-  defaultSelectionFilter() {
+  defaultSelectionFilter(selectionsToRestore?: Selections) {
-    setSelectionFilterToDefault(this.programMemory, this.engineCommandManager)
+    setSelectionFilterToDefault(this.engineCommandManager, selectionsToRestore)
   }
   /** TODO: this function is hiding unawaited asynchronous work */
   setSelectionFilter(filter: EntityType_type[]) {
@@ -640,25 +643,65 @@ const defaultSelectionFilter: EntityType_type[] = [
 
 /** TODO: This function is not synchronous but is currently treated as such */
 function setSelectionFilterToDefault(
-  programMemory: ProgramMemory,
+  engineCommandManager: EngineCommandManager,
-  engineCommandManager: EngineCommandManager
+  selectionsToRestore?: Selections
 ) {
   // eslint-disable-next-line @typescript-eslint/no-floating-promises
-  setSelectionFilter(defaultSelectionFilter, engineCommandManager)
+  setSelectionFilter(
+    defaultSelectionFilter,
+    engineCommandManager,
+    selectionsToRestore
+  )
 }
 
 /** TODO: This function is not synchronous but is currently treated as such */
 function setSelectionFilter(
   filter: EntityType_type[],
-  engineCommandManager: EngineCommandManager
+  engineCommandManager: EngineCommandManager,
+  selectionsToRestore?: Selections
 ) {
-  // eslint-disable-next-line @typescript-eslint/no-floating-promises
+  const { engineEvents } = selectionsToRestore
-  engineCommandManager.sendSceneCommand({
+    ? handleSelectionBatch({
-    type: 'modeling_cmd_req',
+        selections: selectionsToRestore,
-    cmd_id: uuidv4(),
+      })
-    cmd: {
+    : { engineEvents: undefined }
-      type: 'set_selection_filter',
+  if (!selectionsToRestore || !engineEvents) {
-      filter,
+    // eslint-disable-next-line @typescript-eslint/no-floating-promises
-    },
+    engineCommandManager.sendSceneCommand({
+      type: 'modeling_cmd_req',
+      cmd_id: uuidv4(),
+      cmd: {
+        type: 'set_selection_filter',
+        filter,
+      },
+    })
+    return
+  }
+  const modelingCmd: ModelingCmdReq_type[] = []
+  engineEvents.forEach((event) => {
+    if (event.type === 'modeling_cmd_req') {
+      modelingCmd.push({
+        cmd_id: uuidv4(),
+        cmd: event.cmd,
+      })
+    }
   })
+  // batch is needed other wise the selection flickers.
+  engineCommandManager
+    .sendSceneCommand({
+      type: 'modeling_cmd_batch_req',
+      batch_id: uuidv4(),
+      requests: [
+        {
+          cmd_id: uuidv4(),
+          cmd: {
+            type: 'set_selection_filter',
+            filter,
+          },
+        },
+        ...modelingCmd,
+      ],
+      responses: false,
+    })
+    .catch(reportError)
 }

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -76,6 +76,10 @@ export type ModelingCommandSchema = {
   'Text-to-CAD': {
     prompt: string
   }
+  'Prompt-to-edit': {
+    prompt: string
+    selection: Selections
+  }
 }
 
 export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
@@ -479,4 +483,29 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
       },
     },
   },
+  'Prompt-to-edit': {
+    description: 'Use Zoo AI to edit your kcl',
+    icon: 'chat',
+    args: {
+      selection: {
+        inputType: 'selection',
+        selectionTypes: [
+          'solid2D',
+          'segment',
+          'sweepEdge',
+          'cap',
+          'wall',
+          'edgeCut',
+          'edgeCutEdge',
+        ],
+        multiple: true,
+        required: true,
+        skip: true,
+      },
+      prompt: {
+        inputType: 'text',
+        required: true,
+      },
+    },
+  },
 }

src/lib/promptToEdit.ts

@@ -0,0 +1,345 @@
+import { Models } from '@kittycad/lib'
+import { VITE_KC_API_BASE_URL } from 'env'
+import crossPlatformFetch from './crossPlatformFetch'
+import { err, reportRejection } from './trap'
+import { Selections } from './selections'
+import { ArtifactGraph, getArtifactOfTypes } from 'lang/std/artifactGraph'
+import { SourceRange } from 'lang/wasm'
+import toast from 'react-hot-toast'
+import { codeManager, editorManager, kclManager } from './singletons'
+import { ToastPromptToEditCadSuccess } from 'components/ToastTextToCad'
+import { uuidv4 } from './utils'
+import { diffLines } from 'diff'
+import { Transaction, EditorSelection, SelectionRange } from '@codemirror/state'
+import { modelingMachineEvent } from 'editor/manager'
+
+function sourceIndexToLineColumn(
+  code: string,
+  index: number
+): { line: number; column: number } {
+  const codeStart = code.slice(0, index)
+  const lines = codeStart.split('\n')
+  const line = lines.length
+  const column = lines[lines.length - 1].length
+  return { line, column }
+}
+
+function convertAppRangeToApiRange(
+  range: SourceRange,
+  code: string
+): Models['SourceRange_type'] {
+  return {
+    start: sourceIndexToLineColumn(code, range[0]),
+    end: sourceIndexToLineColumn(code, range[1]),
+  }
+}
+
+export async function submitPromptToEditToQueue({
+  prompt,
+  selections,
+  code,
+  token,
+  artifactGraph,
+}: {
+  prompt: string
+  selections: Selections
+  code: string
+  token?: string
+  artifactGraph: ArtifactGraph
+}): Promise<Models['TextToCadIteration_type'] | Error> {
+  const ranges: Models['TextToCadIterationBody_type']['source_ranges'] =
+    selections.graphSelections.flatMap((selection) => {
+      const artifact = selection.artifact
+      const prompts: Models['TextToCadIterationBody_type']['source_ranges'] = []
+
+      if (artifact?.type === 'cap') {
+        prompts.push({
+          prompt: `The users main selection is the end cap of a general-sweep (that is an extrusion, revolve, sweep or loft).
+The source range most likely refers to "startProfileAt" simply because this is the start of the profile that was swept.
+If you need to operate on this cap, for example for sketching on the face, you can use the special string ${
+            artifact.subType === 'end' ? 'END' : 'START'
+          } i.e. \`startSketchOn(someSweepVariable, ${
+            artifact.subType === 'end' ? 'END' : 'START'
+          })\`
+When they made this selection they main have intended this surface directly or meant something more general like the sweep body.
+See later source ranges for more context.`,
+          range: convertAppRangeToApiRange(selection.codeRef.range, code),
+        })
+        let sweep = getArtifactOfTypes(
+          { key: artifact.sweepId, types: ['sweep'] },
+          artifactGraph
+        )
+        if (!err(sweep)) {
+          prompts.push({
+            prompt: `This is the sweep's source range from the user's main selection of the end cap.`,
+            range: convertAppRangeToApiRange(sweep.codeRef.range, code),
+          })
+        }
+      }
+      if (artifact?.type === 'wall') {
+        prompts.push({
+          prompt: `The users main selection is the wall of a general-sweep (that is an extrusion, revolve, sweep or loft).
+The source range though is for the original segment before it was extruded, you can add a tag to that segment in order to refer to this wall, for example "startSketchOn(someSweepVariable, segmentTag)"
+But it's also worth bearing in mind that the user may have intended to select the sweep itself, not this individual wall, see later source ranges for more context. about the sweep`,
+          range: convertAppRangeToApiRange(selection.codeRef.range, code),
+        })
+        let sweep = getArtifactOfTypes(
+          { key: artifact.sweepId, types: ['sweep'] },
+          artifactGraph
+        )
+        if (!err(sweep)) {
+          prompts.push({
+            prompt: `This is the sweep's source range from the user's main selection of the end cap.`,
+            range: convertAppRangeToApiRange(sweep.codeRef.range, code),
+          })
+        }
+      }
+      if (artifact?.type === 'sweepEdge') {
+        prompts.push({
+          prompt: `The users main selection is the edge of a general-sweep (that is an extrusion, revolve, sweep or loft).
+it is an ${
+            artifact.subType
+          } edge, in order to refer to this edge you should add a tag to the segment function in this source range,
+and then use the function ${
+            artifact.subType === 'adjacent'
+              ? 'getAdjacentEdge'
+              : 'getOppositeEdge'
+          }
+See later source ranges for more context. about the sweep`,
+          range: convertAppRangeToApiRange(selection.codeRef.range, code),
+        })
+        let sweep = getArtifactOfTypes(
+          { key: artifact.sweepId, types: ['sweep'] },
+          artifactGraph
+        )
+        if (!err(sweep)) {
+          prompts.push({
+            prompt: `This is the sweep's source range from the user's main selection of the end cap.`,
+            range: convertAppRangeToApiRange(sweep.codeRef.range, code),
+          })
+        }
+      }
+      if (artifact?.type === 'segment') {
+        if (!artifact.surfaceId) {
+          prompts.push({
+            prompt: `This selection is of a segment, likely an individual part of a profile. Segments are often "constrained" by the use of variables and relationships with other segments. Adding tags to segments helps refer to their length, angle or other properties`,
+            range: convertAppRangeToApiRange(selection.codeRef.range, code),
+          })
+        } else {
+          prompts.push({
+            prompt: `This selection is for a segment (line, xLine, angledLine etc) that has been swept (a general-sweep, either an extrusion, revolve, sweep or loft).
+Because it now refers to an edge the way to refer to this edge is to add a tag to the segment, and then use that tag directly.
+i.e. \`fillet({ radius = someInteger, tags = [newTag] }, %)\` will work in the case of filleting this edge
+See later source ranges for more context. about the sweep`,
+            range: convertAppRangeToApiRange(selection.codeRef.range, code),
+          })
+          let path = getArtifactOfTypes(
+            { key: artifact.pathId, types: ['path'] },
+            artifactGraph
+          )
+          if (!err(path)) {
+            const sweep = getArtifactOfTypes(
+              { key: path.sweepId, types: ['sweep'] },
+              artifactGraph
+            )
+            if (!err(sweep)) {
+              prompts.push({
+                prompt: `This is the sweep's source range from the user's main selection of the edge.`,
+                range: convertAppRangeToApiRange(sweep.codeRef.range, code),
+              })
+            }
+          }
+        }
+      }
+      return prompts
+    })
+  const body: Models['TextToCadIterationBody_type'] = {
+    original_source_code: code,
+    prompt,
+    source_ranges: ranges,
+  }
+  const url = VITE_KC_API_BASE_URL + '/ml/text-to-cad/iteration'
+  const data: Models['TextToCadIteration_type'] | Error =
+    await crossPlatformFetch(
+      url,
+      {
+        method: 'POST',
+        body: JSON.stringify(body),
+      },
+      token
+    )
+
+  // Make sure we have an id.
+  if (data instanceof Error) {
+    return data
+  }
+
+  if (!data.id) {
+    return new Error('No id returned from Text-to-CAD API')
+  }
+  return data
+}
+
+export async function getPromptToEditResult(
+  id: string,
+  token?: string
+): Promise<Models['TextToCadIteration_type'] | Error> {
+  const url = VITE_KC_API_BASE_URL + '/async/operations/' + id
+  const data: Models['TextToCadIteration_type'] | Error =
+    await crossPlatformFetch(
+      url,
+      {
+        method: 'GET',
+      },
+      token
+    )
+
+  return data
+}
+
+export async function doPromptEdit({
+  prompt,
+  selections,
+  code,
+  token,
+  artifactGraph,
+}: {
+  prompt: string
+  selections: Selections
+  code: string
+  token?: string
+  artifactGraph: ArtifactGraph
+}): Promise<Models['TextToCadIteration_type'] | Error> {
+  const toastId = toast.loading('Submitting to Edit AI...')
+  const submitResult = await submitPromptToEditToQueue({
+    prompt,
+    selections,
+    code,
+    token,
+    artifactGraph,
+  })
+  if (err(submitResult)) return submitResult
+
+  const textToCadComplete = new Promise<Models['TextToCadIteration_type']>(
+    (resolve, reject) => {
+      ;(async () => {
+        const MAX_CHECK_TIMEOUT = 3 * 60_000
+        const CHECK_DELAY = 200
+
+        let timeElapsed = 0
+
+        while (timeElapsed < MAX_CHECK_TIMEOUT) {
+          const check = await getPromptToEditResult(submitResult.id, token)
+          if (check instanceof Error || check.status === 'failed') {
+            reject(check)
+            return
+          } else if (check.status === 'completed') {
+            resolve(check)
+            return
+          }
+
+          await new Promise((r) => setTimeout(r, CHECK_DELAY))
+          timeElapsed += CHECK_DELAY
+        }
+
+        reject(new Error('Text-to-CAD API timed out'))
+      })().catch(reportRejection)
+    }
+  )
+
+  try {
+    const result = await textToCadComplete
+    toast.dismiss(toastId)
+    return result
+  } catch (e) {
+    toast.dismiss(toastId)
+    toast.error(
+      'Failed to edit your KCL code, please try again with a different prompt or selection'
+    )
+    console.error('textToCadComplete', e)
+  }
+
+  return textToCadComplete
+}
+
+/** takes care of the whole submit prompt to endpoint flow including the accept-reject toast once the result is back */
+export async function promptToEditFlow({
+  prompt,
+  selections,
+  code,
+  token,
+  artifactGraph,
+}: {
+  prompt: string
+  selections: Selections
+  code: string
+  token?: string
+  artifactGraph: ArtifactGraph
+}) {
+  const result = await doPromptEdit({
+    prompt,
+    selections,
+    code,
+    token,
+    artifactGraph,
+  })
+  if (err(result)) return Promise.reject(result)
+  const oldCode = codeManager.code
+  const { code: newCode } = result
+  codeManager.updateCodeEditor(newCode)
+  const diff = reBuildNewCodeWithRanges(oldCode, newCode)
+  const ranges: SelectionRange[] = diff.insertRanges.map((range) =>
+    EditorSelection.range(range[0], range[1])
+  )
+  editorManager?.editorView?.dispatch({
+    selection: EditorSelection.create(
+      ranges,
+      selections.graphSelections.length - 1
+    ),
+    annotations: [modelingMachineEvent, Transaction.addToHistory.of(false)],
+  })
+  await kclManager.executeCode()
+  const toastId = uuidv4()
+
+  toast.success(
+    () =>
+      ToastPromptToEditCadSuccess({
+        toastId,
+        data: result,
+        token,
+        oldCode,
+      }),
+    {
+      id: toastId,
+      duration: Infinity,
+      icon: null,
+    }
+  )
+}
+
+const reBuildNewCodeWithRanges = (
+  oldCode: string,
+  newCode: string
+): {
+  newCode: string
+  insertRanges: SourceRange[]
+} => {
+  let insertRanges: SourceRange[] = []
+  const changes = diffLines(oldCode, newCode)
+  let newCodeWithRanges = ''
+  for (const change of changes) {
+    if (!change.added && !change.removed) {
+      // no change add it to newCodeWithRanges
+      newCodeWithRanges += change.value
+    } else if (change.added && !change.removed) {
+      const start = newCodeWithRanges.length
+      const end = start + change.value.length
+      insertRanges.push([start, end, true])
+      newCodeWithRanges += change.value
+    }
+  }
+  return {
+    newCode: newCodeWithRanges,
+    insertRanges,
+  }
+}

src/lib/textToCad.ts

@@ -17,7 +17,7 @@ import { getNextFileName } from './desktopFS'
 import { reportRejection } from './trap'
 import { toSync } from './utils'
 
-export async function submitTextToCadPrompt(
+async function submitTextToCadPrompt(
   prompt: string,
   token?: string
 ): Promise<Models['TextToCad_type'] | Error> {
@@ -45,7 +45,7 @@ export async function submitTextToCadPrompt(
   return data
 }
 
-export async function getTextToCadResult(
+async function getTextToCadResult(
   id: string,
   token?: string
 ): Promise<Models['TextToCad_type'] | Error> {

yarn.lock

@@ -2422,6 +2422,11 @@
   dependencies:
     "@types/ms" "*"
 
+"@types/diff@^6.0.0":
+  version "6.0.0"
+  resolved "https://registry.yarnpkg.com/@types/diff/-/diff-6.0.0.tgz#031f27cf57564f3cce825f38fb19fdd4349ad07a"
+  integrity sha512-dhVCYGv3ZSbzmQaBSagrv1WJ6rXCdkyTcDyoNu1MD8JohI7pR7k8wdZEm+mvdxRKXyHVwckFzWU1vJc+Z29MlA==
+
 "@types/electron@^1.6.10":
   version "1.6.10"
   resolved "https://registry.yarnpkg.com/@types/electron/-/electron-1.6.10.tgz#7e87888ed3888767cca68e92772c2c8ea46bc873"
@@ -4247,6 +4252,11 @@ diff@^4.0.1:
   resolved "https://registry.yarnpkg.com/diff/-/diff-4.0.2.tgz#60f3aecb89d5fae520c11aa19efc2bb982aade7d"
   integrity sha512-58lmxKSA4BNyLz+HHMUzlOEpg09FV+ev6ZMe3vJihgdxzgcwZ8VoEEPmALCZG9LmqfVoNMMKpttIYTVG6uDY7A==
 
+diff@^7.0.0:
+  version "7.0.0"
+  resolved "https://registry.yarnpkg.com/diff/-/diff-7.0.0.tgz#3fb34d387cd76d803f6eebea67b921dab0182a9a"
+  integrity sha512-PJWHUb1RFevKCwaFA9RlG5tCd+FO5iRh9A8HEtkmBH2Li03iJriB6m6JIN4rGz3K3JLawI7/veA1xzRKP6ISBw==
+
 dir-compare@^3.0.0:
   version "3.3.0"
   resolved "https://registry.yarnpkg.com/dir-compare/-/dir-compare-3.3.0.tgz#2c749f973b5c4b5d087f11edaae730db31788416"