Multiple prompt-to-edit selection, plus direct editor selections (#5478)

* Add multiple selections and editor selections for promptToEdit

* remove unused

* re-enable prompt to edit tests

* add test for manual code selection

* at test for multi-selection

* clean up

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* typo

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

e2e/playwright/fixtures/editorFixture.ts

@@ -171,4 +171,22 @@ export class EditorFixture {
       { text, placeCursor }
     )
   }
+  async selectText(text: string) {
+    // First make sure the code pane is open
+    const wasPaneOpen = await this.checkIfPaneIsOpen()
+    if (!wasPaneOpen) {
+      await this.openPane()
+    }
+
+    // Use Playwright's built-in text selection on the code content
+    // it seems to only select whole divs, which works out to align with syntax highlighting
+    // for code mirror, so you can probably select "sketch002 = startSketchOn('XZ')"
+    // but less so for exactly "sketch002 = startS"
+    await this.codeContent.getByText(text).first().selectText()
+
+    // Reset pane state if needed
+    if (!wasPaneOpen) {
+      await this.closePane()
+    }
+  }
 }

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

@@ -36,7 +36,7 @@ extrude003 = extrude(sketch003, length = 20)
 `
 
 test.describe('Prompt-to-edit tests', { tag: '@skipWin' }, () => {
-  test.fixme('Check the happy path, for basic changing color', () => {
+  test.describe('Check the happy path, for basic changing color', () => {
     const cases = [
       {
         desc: 'User accepts change',
@@ -70,7 +70,7 @@ test.describe('Prompt-to-edit tests', { tag: '@skipWin' }, () => {
           body1CapCoords.y
         )
         const yellow: [number, number, number] = [179, 179, 131]
-        const green: [number, number, number] = [108, 152, 75]
+        const green: [number, number, number] = [128, 194, 88]
         const notGreen: [number, number, number] = [132, 132, 132]
         const body2NotGreen: [number, number, number] = [88, 88, 88]
         const submittingToast = page.getByText(
@@ -109,7 +109,7 @@ test.describe('Prompt-to-edit tests', { tag: '@skipWin' }, () => {
         })
 
         await test.step('verify initial change', async () => {
-          await scene.expectPixelColor(green, greenCheckCoords, 15)
+          await scene.expectPixelColor(green, greenCheckCoords, 20)
           await scene.expectPixelColor(body2NotGreen, body2WallCoords, 15)
           await editor.expectEditor.toContain('appearance(')
         })
@@ -142,7 +142,7 @@ test.describe('Prompt-to-edit tests', { tag: '@skipWin' }, () => {
     }
   })
 
-  test(`bad edit prompt`, async ({
+  test('bad edit prompt', async ({
     context,
     homePage,
     cmdBar,
@@ -195,4 +195,150 @@ test.describe('Prompt-to-edit tests', { tag: '@skipWin' }, () => {
       await expect(failToast).toBeVisible()
     })
   })
+
+  test(`manual code selection rename`, async ({
+    context,
+    homePage,
+    cmdBar,
+    editor,
+    page,
+    scene,
+  }) => {
+    const body1CapCoords = { x: 571, y: 351 }
+
+    await context.addInitScript((file) => {
+      localStorage.setItem('persistCode', file)
+    }, file)
+    await homePage.goToModelingScene()
+    await scene.waitForExecutionDone()
+
+    const submittingToast = page.getByText('Submitting to Text-to-CAD API...')
+    const successToast = page.getByText('Prompt to edit successful')
+    const acceptBtn = page.getByRole('button', { name: 'checkmark Accept' })
+
+    await test.step('wait for scene to load and select code in editor', async () => {
+      // Find and select the text "sketch002" in the editor
+      await editor.selectText('sketch002')
+
+      // Verify the selection was made
+      await editor.expectState({
+        highlightedCode: '',
+        activeLines: ["sketch002 = startSketchOn('XZ')"],
+        diagnostics: [],
+      })
+    })
+
+    await test.step('fire off edit prompt', async () => {
+      await scene.expectPixelColor([134, 134, 134], body1CapCoords, 15)
+      await cmdBar.openCmdBar('promptToEdit')
+      await page
+        .getByTestId('cmd-bar-arg-value')
+        .fill('Please rename to mySketch')
+      await page.waitForTimeout(100)
+      await cmdBar.progressCmdBar()
+      await expect(submittingToast).toBeVisible()
+      await expect(submittingToast).not.toBeVisible({
+        timeout: 2 * 60_000,
+      })
+      await expect(successToast).toBeVisible()
+    })
+
+    await test.step('verify rename change and accept it', async () => {
+      await editor.expectEditor.toContain('mySketch = startSketchOn')
+      await editor.expectEditor.not.toContain('sketch002 = startSketchOn')
+      await editor.expectEditor.toContain(
+        'extrude002 = extrude(mySketch, length = 50)'
+      )
+
+      await acceptBtn.click()
+      await expect(successToast).not.toBeVisible()
+    })
+  })
+
+  test('multiple body selections', async ({
+    context,
+    homePage,
+    cmdBar,
+    editor,
+    page,
+    scene,
+  }) => {
+    const body1CapCoords = { x: 571, y: 351 }
+    const body2WallCoords = { x: 620, y: 152 }
+    const [clickBody1Cap] = scene.makeMouseHelpers(
+      body1CapCoords.x,
+      body1CapCoords.y
+    )
+    const [clickBody2Cap] = scene.makeMouseHelpers(
+      body2WallCoords.x,
+      body2WallCoords.y
+    )
+    const grey: [number, number, number] = [132, 132, 132]
+
+    await context.addInitScript((file) => {
+      localStorage.setItem('persistCode', file)
+    }, file)
+    await homePage.goToModelingScene()
+    await scene.waitForExecutionDone()
+
+    const submittingToast = page.getByText('Submitting to Text-to-CAD API...')
+    const successToast = page.getByText('Prompt to edit successful')
+    const acceptBtn = page.getByRole('button', { name: 'checkmark Accept' })
+
+    await test.step('select multiple bodies and fire prompt', async () => {
+      // Initial color check
+      await scene.expectPixelColor(grey, body1CapCoords, 15)
+
+      // Open command bar first (without selection)
+      await cmdBar.openCmdBar('promptToEdit')
+
+      // Select first body
+      await page.waitForTimeout(100)
+      await clickBody1Cap()
+
+      // Hold shift and select second body
+      await editor.expectState({
+        highlightedCode: '',
+        activeLines: ['|>startProfileAt([-73.64,-42.89],%)'],
+        diagnostics: [],
+      })
+      await page.keyboard.down('Shift')
+      await page.waitForTimeout(100)
+      await clickBody2Cap()
+      await editor.expectState({
+        highlightedCode:
+          'line(end=[121.13,56.63],tag=$seg02)extrude(profile001,length=200)',
+        activeLines: [
+          '|>line(end=[121.13,56.63],tag=$seg02)',
+          '|>startProfileAt([-73.64,-42.89],%)',
+        ],
+        diagnostics: [],
+      })
+      await page.keyboard.up('Shift')
+      await page.waitForTimeout(100)
+      await cmdBar.progressCmdBar()
+
+      // Enter prompt and submit
+      await page
+        .getByTestId('cmd-bar-arg-value')
+        .fill('make these neon green please, use #39FF14')
+      await page.waitForTimeout(100)
+      await cmdBar.progressCmdBar()
+
+      // Wait for API response
+      await expect(submittingToast).toBeVisible()
+      await expect(submittingToast).not.toBeVisible({
+        timeout: 2 * 60_000,
+      })
+      await expect(successToast).toBeVisible()
+    })
+
+    await test.step('verify code changed', async () => {
+      await editor.expectEditor.toContain('appearance(')
+
+      // Accept changes
+      await acceptBtn.click()
+      await expect(successToast).not.toBeVisible()
+    })
+  })
 })

src/components/CommandBar/CommandBar.tsx

@@ -17,7 +17,9 @@ export const CommandBar = () => {
   const {
     context: { selectedCommand, currentArgument, commands },
   } = commandBarState
-  const isSelectionArgument = currentArgument?.inputType === 'selection'
+  const isSelectionArgument =
+    currentArgument?.inputType === 'selection' ||
+    currentArgument?.inputType === 'selectionMixed'
   const WrapperComponent = isSelectionArgument ? Popover : Dialog
 
   // Close the command bar when navigating

src/components/CommandBar/CommandBarArgument.tsx

@@ -1,6 +1,7 @@
 import CommandArgOptionInput from './CommandArgOptionInput'
 import CommandBarBasicInput from './CommandBarBasicInput'
 import CommandBarSelectionInput from './CommandBarSelectionInput'
+import CommandBarSelectionMixedInput from './CommandBarSelectionMixedInput'
 import { CommandArgument } from 'lib/commandTypes'
 import CommandBarHeader from './CommandBarHeader'
 import CommandBarKclInput from './CommandBarKclInput'
@@ -84,6 +85,14 @@ function ArgumentInput({
           onSubmit={onSubmit}
         />
       )
+    case 'selectionMixed':
+      return (
+        <CommandBarSelectionMixedInput
+          arg={arg}
+          stepBack={stepBack}
+          onSubmit={onSubmit}
+        />
+      )
     case 'kcl':
       return (
         <CommandBarKclInput arg={arg} stepBack={stepBack} onSubmit={onSubmit} />

src/components/CommandBar/CommandBarHeader.tsx

@@ -124,7 +124,8 @@ function CommandBarHeader({ children }: React.PropsWithChildren<{}>) {
                     <span className="sr-only">:&nbsp;</span>
                     <span data-testid="header-arg-value">
                       {argValue ? (
-                        arg.inputType === 'selection' ? (
+                        arg.inputType === 'selection' ||
+                        arg.inputType === 'selectionMixed' ? (
                           getSelectionTypeDisplayText(argValue as Selections)
                         ) : arg.inputType === 'kcl' ? (
                           roundOff(

src/components/CommandBar/CommandBarSelectionMixedInput.tsx

@@ -0,0 +1,135 @@
+import { useEffect, useMemo, useRef, useState } from 'react'
+import { CommandArgument } from 'lib/commandTypes'
+import {
+  Selections,
+  canSubmitSelectionArg,
+  getSelectionCountByType,
+  getSelectionTypeDisplayText,
+} from 'lib/selections'
+import { useSelector } from '@xstate/react'
+import { commandBarActor, useCommandBarState } from 'machines/commandBarMachine'
+
+const selectionSelector = (snapshot: any) => snapshot?.context.selectionRanges
+
+export default function CommandBarSelectionMixedInput({
+  arg,
+  stepBack,
+  onSubmit,
+}: {
+  arg: CommandArgument<unknown> & { inputType: 'selectionMixed'; name: string }
+  stepBack: () => void
+  onSubmit: (data: unknown) => void
+}) {
+  const inputRef = useRef<HTMLInputElement>(null)
+  const commandBarState = useCommandBarState()
+  const [hasSubmitted, setHasSubmitted] = useState(false)
+  const [hasAutoSkipped, setHasAutoSkipped] = useState(false)
+  const selection: Selections = useSelector(arg.machineActor, selectionSelector)
+
+  const selectionsByType = useMemo(() => {
+    return getSelectionCountByType(selection)
+  }, [selection])
+
+  const canSubmitSelection = useMemo<boolean>(() => {
+    if (!selection) return false
+    const isNonZeroRange = selection.graphSelections.some((sel) => {
+      const range = sel.codeRef.range
+      return range[1] - range[0] !== 0 // Non-zero range is always valid
+    })
+    if (isNonZeroRange) return true
+    return canSubmitSelectionArg(selectionsByType, arg)
+  }, [selectionsByType, selection])
+
+  useEffect(() => {
+    inputRef.current?.focus()
+  }, [selection, inputRef])
+
+  // Only auto-skip on initial mount if we have a valid selection
+  // different from the component CommandBarSelectionInput in the the dependency array
+  // is empty
+  useEffect(() => {
+    if (!hasAutoSkipped && canSubmitSelection && arg.skip) {
+      const argValue = commandBarState.context.argumentsToSubmit[arg.name]
+      if (argValue === undefined) {
+        handleSubmit()
+        setHasAutoSkipped(true)
+      }
+    }
+  }, [])
+
+  function handleChange() {
+    inputRef.current?.focus()
+  }
+
+  function handleSubmit(e?: React.FormEvent<HTMLFormElement>) {
+    e?.preventDefault()
+
+    if (!canSubmitSelection) {
+      setHasSubmitted(true)
+      return
+    }
+
+    onSubmit(selection)
+  }
+
+  const isMixedSelection = arg.inputType === 'selectionMixed'
+  const allowNoSelection = isMixedSelection && arg.allowNoSelection
+  const showSceneSelection =
+    isMixedSelection && arg.selectionSource?.allowSceneSelection
+
+  return (
+    <form id="arg-form" onSubmit={handleSubmit}>
+      <label
+        className={
+          'relative flex flex-col mx-4 my-4 ' +
+          (!hasSubmitted || canSubmitSelection || 'text-destroy-50')
+        }
+      >
+        {canSubmitSelection
+          ? 'Select objects in the scene'
+          : 'Select code or objects in the scene'}
+
+        {showSceneSelection && (
+          <div className="scene-selection mt-2">
+            <p className="text-sm text-chalkboard-60">
+              Select objects in the scene
+            </p>
+            {/* Scene selection UI will be handled by the parent component */}
+          </div>
+        )}
+
+        {allowNoSelection && (
+          <button
+            type="button"
+            onClick={() => onSubmit(null)}
+            className="mt-2 px-4 py-2 rounded border border-chalkboard-30 text-chalkboard-90 dark:text-chalkboard-10 hover:bg-chalkboard-10 dark:hover:bg-chalkboard-90 transition-colors"
+          >
+            Continue without selection
+          </button>
+        )}
+
+        <span data-testid="cmd-bar-arg-name" className="sr-only">
+          {arg.name}
+        </span>
+        <input
+          id="selection"
+          name="selection"
+          ref={inputRef}
+          required
+          data-testid="cmd-bar-arg-value"
+          placeholder="Select an entity with your mouse"
+          className="absolute inset-0 w-full h-full opacity-0 cursor-default"
+          onKeyDown={(event) => {
+            if (event.key === 'Backspace') {
+              stepBack()
+            } else if (event.key === 'Escape') {
+              commandBarActor.send({ type: 'Close' })
+            }
+          }}
+          onChange={handleChange}
+          value={JSON.stringify(selection || {})}
+        />
+      </label>
+    </form>
+  )
+}

src/lib/commandBarConfigs/modelingCommandConfig.ts

@@ -666,7 +666,7 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
     icon: 'chat',
     args: {
       selection: {
-        inputType: 'selection',
+        inputType: 'selectionMixed',
         selectionTypes: [
           'solid2d',
           'segment',
@@ -678,6 +678,10 @@ export const modelingMachineCommandConfig: StateMachineCommandSetConfig<
         ],
         multiple: true,
         required: true,
+        selectionSource: {
+          allowSceneSelection: true,
+          allowCodeSelection: true,
+        },
         skip: true,
       },
       prompt: {

src/lib/commandTypes.ts

@@ -16,6 +16,7 @@ const INPUT_TYPES = [
   'text',
   'kcl',
   'selection',
+  'selectionMixed',
   'boolean',
 ] as const
 export interface KclExpression {
@@ -156,6 +157,23 @@ export type CommandArgumentConfig<
         context: CommandBarContext
       }) => Promise<boolean | string>
     }
+  | {
+      inputType: 'selectionMixed'
+      selectionTypes: Artifact['type'][]
+      multiple: boolean
+      allowNoSelection?: boolean
+      validation?: ({
+        data,
+        context,
+      }: {
+        data: any
+        context: CommandBarContext
+      }) => Promise<boolean | string>
+      selectionSource?: {
+        allowSceneSelection?: boolean
+        allowCodeSelection?: boolean
+      }
+    }
   | {
       inputType: 'kcl'
       createVariableByDefault?: boolean
@@ -252,6 +270,23 @@ export type CommandArgument<
         context: CommandBarContext
       }) => Promise<boolean | string>
     }
+  | {
+      inputType: 'selectionMixed'
+      selectionTypes: Artifact['type'][]
+      multiple: boolean
+      allowNoSelection?: boolean
+      validation?: ({
+        data,
+        context,
+      }: {
+        data: any
+        context: CommandBarContext
+      }) => Promise<boolean | string>
+      selectionSource?: {
+        allowSceneSelection?: boolean
+        allowCodeSelection?: boolean
+      }
+    }
   | {
       inputType: 'kcl'
       createVariableByDefault?: boolean

src/lib/createMachineCommand.ts

@@ -187,6 +187,16 @@ export function buildCommandArgument<
       selectionTypes: arg.selectionTypes,
       validation: arg.validation,
     } satisfies CommandArgument<O, T> & { inputType: 'selection' }
+  } else if (arg.inputType === 'selectionMixed') {
+    return {
+      inputType: arg.inputType,
+      ...baseCommandArgument,
+      multiple: arg.multiple,
+      selectionTypes: arg.selectionTypes,
+      validation: arg.validation,
+      allowNoSelection: arg.allowNoSelection,
+      selectionSource: arg.selectionSource,
+    } satisfies CommandArgument<O, T> & { inputType: 'selectionMixed' }
   } else if (arg.inputType === 'kcl') {
     return {
       inputType: arg.inputType,

src/lib/promptToEdit.ts

@@ -43,15 +43,33 @@ export async function submitPromptToEditToQueue({
   projectName,
 }: {
   prompt: string
-  selections: Selections
+  selections: Selections | null
   code: string
   projectName: string
   token?: string
   artifactGraph: ArtifactGraph
 }): Promise<Models['TextToCadIteration_type'] | Error> {
+  // If no selection, use whole file
+  if (selections === null) {
+    const body: Models['TextToCadIterationBody_type'] = {
+      original_source_code: code,
+      prompt,
+      source_ranges: [], // Empty ranges indicates whole file
+      project_name:
+        projectName !== '' && projectName !== 'browser'
+          ? projectName
+          : undefined,
+      kcl_version: kclManager.kclVersion,
+    }
+    return submitToApi(body, token)
+  }
+
+  // Handle manual code selections and artifact selections differently
   const ranges: Models['TextToCadIterationBody_type']['source_ranges'] =
     selections.graphSelections.flatMap((selection) => {
       const artifact = selection.artifact
+
+      // For artifact selections, add context
       const prompts: Models['TextToCadIterationBody_type']['source_ranges'] = []
 
       if (artifact?.type === 'cap') {
@@ -153,8 +171,17 @@ See later source ranges for more context. about the sweep`,
           }
         }
       }
+      if (!artifact) {
+        // manually selected code is more likely to not have an artifact
+        // an example might be highlighting the variable name only in a variable declaration
+        prompts.push({
+          prompt: '',
+          range: convertAppRangeToApiRange(selection.codeRef.range, code),
+        })
+      }
       return prompts
     })
+
   const body: Models['TextToCadIterationBody_type'] = {
     original_source_code: code,
     prompt,
@@ -163,6 +190,15 @@ See later source ranges for more context. about the sweep`,
       projectName !== '' && projectName !== 'browser' ? projectName : undefined,
     kcl_version: kclManager.kclVersion,
   }
+
+  return submitToApi(body, token)
+}
+
+// Helper function to handle API submission
+async function submitToApi(
+  body: Models['TextToCadIterationBody_type'],
+  token?: string
+): Promise<Models['TextToCadIteration_type'] | Error> {
   const url = VITE_KC_API_BASE_URL + '/ml/text-to-cad/iteration'
   const data: Models['TextToCadIteration_type'] | Error =
     await crossPlatformFetch(

src/lib/selections.ts

@@ -481,7 +481,9 @@ export function getSelectionTypeDisplayText(
 
 export function canSubmitSelectionArg(
   selectionsByType: 'none' | Map<ResolvedSelectionType, number>,
-  argument: CommandArgument<unknown> & { inputType: 'selection' }
+  argument: CommandArgument<unknown> & {
+    inputType: 'selection' | 'selectionMixed'
+  }
 ) {
   return (
     selectionsByType !== 'none' &&

src/machines/commandBarMachine.ts

@@ -295,7 +295,8 @@ export const commandBarMachine = setup({
           if (
             context.currentArgument &&
             context.selectedCommand &&
-            argConfig?.inputType === 'selection' &&
+            (argConfig?.inputType === 'selection' ||
+              argConfig?.inputType === 'selectionMixed') &&
             argConfig?.validation
           ) {
             argConfig