Implement invocation count selection and support for multiple parallel AI requests per prompt

Author: robertpiosik

packages/ui/src/components/editor/common/DropdownMenu/DropdownMenu.module.scss

@@ -6,7 +6,7 @@
   z-index: 2;
   width: max-content;
   text-align: left;
-  min-width: 100%;
+  min-width: max(100%, 80px);
 
   &__inner {
     background-color: var(--vscode-quickInput-background);
@@ -35,6 +35,12 @@
     white-space: nowrap;
   }
 
+  &__right {
+    display: flex;
+    align-items: center;
+    gap: 8px;
+  }
+
   &:hover,
   &[data-selected='true'] {
     background-color: var(--vscode-list-activeSelectionBackground);

packages/ui/src/components/editor/common/DropdownMenu/DropdownMenu.tsx

@@ -6,6 +6,7 @@ export namespace DropdownMenu {
     label: string
     shortcut?: string
     on_click: () => void
+    checked?: boolean
     is_selected?: boolean
   }
 
@@ -50,8 +51,13 @@ export const DropdownMenu: React.FC<DropdownMenu.Props> = (props) => {
                   item.label
                 )}
               </span>
-              {item.shortcut && (
-                <span className={styles.shortcut}>{item.shortcut}</span>
+              {(item.shortcut || item.checked) && (
+                <div className={styles.item__right}>
+                  {item.checked && <span className="codicon codicon-check" />}
+                  {item.shortcut && (
+                    <span className={styles.shortcut}>{item.shortcut}</span>
+                  )}
+                </div>
               )}
             </div>
           )

packages/ui/src/components/editor/panel/PromptField/PromptField.module.scss

@@ -274,7 +274,7 @@
     display: flex;
     align-items: center;
     justify-content: flex-end;
-    gap: 8px;
+    gap: 5px;
     pointer-events: none;
 
     &__edit-format {
@@ -321,6 +321,18 @@
       }
     }
 
+    &__invocation-count {
+      position: relative;
+      pointer-events: all;
+      display: flex;
+      align-items: center;
+
+      &__button {
+        font-size: 11px !important;
+        padding: 0 4px !important;
+      }
+    }
+
     &__submit {
       display: flex;
       align-items: center;

packages/ui/src/components/editor/panel/PromptField/PromptField.stories.tsx

@@ -24,6 +24,10 @@ export const Empty = () => (
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
     on_curly_braces_click={() => {}}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
   />
 )
 
@@ -46,6 +50,10 @@ export const WithText = () => (
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
     on_curly_braces_click={() => {}}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
   />
 )
 
@@ -68,6 +76,10 @@ export const LongText = () => (
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
     on_curly_braces_click={() => {}}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
   />
 )
 
@@ -89,6 +101,10 @@ export const WithPlaceholderSavedContext = () => (
     is_web_mode={false}
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
     on_curly_braces_click={() => {}}
   />
 )
@@ -111,6 +127,10 @@ export const WithPlaceholderSelection = () => (
     is_web_mode={false}
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
     on_curly_braces_click={() => {}}
   />
 )
@@ -133,6 +153,10 @@ export const WithCommit = () => (
     is_web_mode={false}
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
     on_curly_braces_click={() => {}}
   />
 )
@@ -155,6 +179,10 @@ export const WithCommitWithQuotes = () => (
     is_web_mode={false}
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
     on_curly_braces_click={() => {}}
   />
 )
@@ -176,6 +204,10 @@ export const WithContextAtCommit = () => (
     is_web_mode={false}
     on_submit_with_control={() => console.log('Submitted with control')}
     context_file_paths={[]}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
     on_curly_braces_click={() => {}}
   />
 )
@@ -210,6 +242,8 @@ export const WithEditFormatSelector = () => {
         diff: 'Use diff...'
       }}
       context_file_paths={[]}
+      invocation_count={1}
+      on_invocation_count_change={set_edit_format as any} // Mocking for demo
     />
   )
 }
@@ -234,5 +268,9 @@ export const WithFilePaths = () => (
     on_curly_braces_click={() => {}}
     on_go_to_file={(path) => console.log('Go to file:', path)}
     context_file_paths={['path/to/my/file.ts']}
+    invocation_count={1}
+    on_invocation_count_change={(count) =>
+      console.log('Invocation count changed:', count)
+    }
   />
 )

packages/ui/src/components/editor/panel/PromptField/PromptField.tsx

@@ -52,6 +52,8 @@ export type PromptFieldProps = {
   edit_format_instructions?: Record<EditFormat, string>
   context_file_paths?: string[]
   currently_open_file_text?: string
+  invocation_count: number
+  on_invocation_count_change: (count: number) => void
   on_go_to_file?: (file_path: string) => void
 }
 
@@ -62,6 +64,8 @@ export const PromptField: React.FC<PromptFieldProps> = (props) => {
   const [show_submit_tooltip, set_show_submit_tooltip] = useState(false)
   const [is_text_selecting, set_is_text_selecting] = useState(false)
   const [is_focused, set_is_focused] = useState(false)
+  const [is_invocation_dropdown_open, set_is_invocation_dropdown_open] =
+    useState(false)
   const [hovered_edit_format, set_hovered_edit_format] =
     useState<EditFormat | null>(null)
   const { is_alt_pressed, handle_container_key_down } = use_keyboard_shortcuts({
@@ -73,10 +77,12 @@ export const PromptField: React.FC<PromptFieldProps> = (props) => {
   const {
     is_dropdown_open,
     toggle_dropdown,
+    close_dropdown,
     dropdown_ref,
     handle_copy_click,
     handle_select_click
   } = use_dropdown(props)
+  const invocation_dropdown_ref = useRef<HTMLDivElement>(null)
   const { ghost_text, handle_accept_ghost_text } = use_ghost_text({
     value: props.value,
     input_ref,
@@ -207,6 +213,26 @@ export const PromptField: React.FC<PromptFieldProps> = (props) => {
     }
   }, [props.focus_key])
 
+  useEffect(() => {
+    if (input_ref.current) {
+      const handle_click_outside = (event: MouseEvent) => {
+        if (
+          invocation_dropdown_ref.current &&
+          !invocation_dropdown_ref.current.contains(event.target as Node)
+        ) {
+          set_is_invocation_dropdown_open(false)
+        }
+      }
+
+      if (is_invocation_dropdown_open) {
+        document.addEventListener('mousedown', handle_click_outside)
+      }
+      return () => {
+        document.removeEventListener('mousedown', handle_click_outside)
+      }
+    }
+  }, [is_invocation_dropdown_open])
+
   useEffect(() => {
     if (input_ref.current) {
       input_ref.current.focus()
@@ -493,6 +519,37 @@ export const PromptField: React.FC<PromptFieldProps> = (props) => {
               {(!props.is_web_mode ||
                 (props.is_web_mode && props.is_connected)) && (
                 <>
+                  <div
+                    className={styles['footer__right__invocation-count']}
+                    ref={invocation_dropdown_ref}
+                  >
+                    <button
+                      className={cn(
+                        styles['footer__right__submit__button'],
+                        styles['footer__right__invocation-count__button']
+                      )}
+                      onClick={(e) => {
+                        e.stopPropagation()
+                        set_is_invocation_dropdown_open((prev) => !prev)
+                        close_dropdown()
+                      }}
+                      title="Invocation count"
+                    >
+                      {props.invocation_count}×
+                    </button>
+                    {is_invocation_dropdown_open && (
+                      <DropdownMenu
+                        items={[1, 2, 3, 4, 5].map((count) => ({
+                          label: `${count}×`,
+                          checked: count == props.invocation_count,
+                          on_click: () => {
+                            props.on_invocation_count_change(count)
+                            set_is_invocation_dropdown_open(false)
+                          }
+                        }))}
+                      />
+                    )}
+                  </div>
                   <button
                     className={cn(
                       styles['footer__right__submit__button'],
@@ -509,7 +566,10 @@ export const PromptField: React.FC<PromptFieldProps> = (props) => {
                       styles['footer__right__submit__button'],
                       styles['footer__right__submit__button--chevron']
                     )}
-                    onClick={toggle_dropdown}
+                    onClick={() => {
+                      toggle_dropdown()
+                      set_is_invocation_dropdown_open(false)
+                    }}
                     title="More actions"
                   >
                     <span

packages/ui/src/components/editor/panel/PromptField/hooks/use-dropdown.ts

@@ -50,6 +50,7 @@ export const use_dropdown = ({
   return {
     is_dropdown_open,
     toggle_dropdown,
+    close_dropdown,
     dropdown_ref,
     handle_copy_click,
     handle_select_click

packages/vscode/src/commands/apply-chat-response-command/response-processor.ts

@@ -6,7 +6,6 @@ import { parse_response, FileItem, DiffItem } from './utils/clipboard-parser'
 import { create_safe_path } from '@/utils/path-sanitizer'
 import { dictionary } from '@shared/constants/dictionary'
 import { Logger } from '@shared/utils/logger'
-import { Diff } from './utils/clipboard-parser/extract-diff-patches'
 import { apply_git_patch } from './handlers/diff-handler'
 import { apply_file_relocations } from './utils/file-operations'
 import { ModelProvidersManager } from '@/services/model-providers-manager'
@@ -23,6 +22,7 @@ import {
   handle_truncated_edit,
   is_truncation_line
 } from './handlers/truncated-handler'
+import { Diff } from 'diff'
 
 export type PreviewData = {
   original_states: OriginalFileState[]
@@ -197,7 +197,7 @@ export const process_chat_response = async (
 
     let all_original_states: OriginalFileState[] = []
     const applied_patches: {
-      patch: Diff
+      patch: DiffItem
       original_states: OriginalFileState[]
       diff_application_method?: 'recount' | 'search_and_replace'
     }[] = []