Refactor the code completion parser to correctly handle and return surrounding text alongside the completion item

Author: robertpiosik

packages/vscode/src/commands/apply-chat-response-command/utils/clipboard-parser/clipboard-parser.spec.ts

@@ -573,6 +573,32 @@ describe('clipboard-parser', () => {
       })
     })
 
+    it('parses code completion format with surrounding text', () => {
+      const test_case = 'code-completion-with-text'
+      const text = load_test_case_file(test_case, `${test_case}.txt`)
+      const result = parse_response({
+        response: text,
+        is_single_root_folder_workspace: true
+      })
+
+      expect(result).toHaveLength(3)
+      expect(result[0]).toMatchObject({
+        type: 'text',
+        content: 'Some text before.'
+      })
+      expect(result[1]).toMatchObject({
+        type: 'completion',
+        file_path: 'src/index.ts',
+        line: 25,
+        character: 5,
+        content: 'console.log("completion");'
+      })
+      expect(result[2]).toMatchObject({
+        type: 'text',
+        content: 'Some text after.'
+      })
+    })
+
     it('parses diff format with git header but no markdown code block', () => {
       const test_case = 'diff-with-git-header-no-markdown'
       const text = load_test_case_file(test_case, `${test_case}.txt`)

packages/vscode/src/commands/apply-chat-response-command/utils/clipboard-parser/clipboard-parser.ts

@@ -79,12 +79,12 @@ export const parse_response = (params: {
 }): ClipboardItem[] => {
   const is_single_root_folder_workspace =
     params.is_single_root_folder_workspace ?? true
-  const code_completion = parse_code_completion({
+  const code_completion_items = parse_code_completion({
     response: params.response,
     is_single_root_folder_workspace
   })
-  if (code_completion) {
-    return [code_completion]
+  if (code_completion_items && code_completion_items.length > 0) {
+    return code_completion_items
   }
 
   const processed_response = params.response.replace(/``````/g, '```\n```')

packages/vscode/src/commands/apply-chat-response-command/utils/clipboard-parser/parsers/code-completion-parser.ts

@@ -1,5 +1,9 @@
 import { cleanup_api_response } from '@/utils/cleanup-api-response'
-import { CompletionItem, extract_workspace_and_path } from '../clipboard-parser'
+import {
+  CompletionItem,
+  TextItem,
+  extract_workspace_and_path
+} from '../clipboard-parser'
 
 const extract_path_and_position = (
   line: string
@@ -28,61 +32,71 @@ const extract_path_and_position = (
 export const parse_code_completion = (params: {
   response: string
   is_single_root_folder_workspace: boolean
-}): CompletionItem | null => {
+}): (CompletionItem | TextItem)[] | null => {
   const lines = params.response.split('\n')
-  let in_code_block = false
-  let first_line_of_block: string | null = null
-  let content_lines: string[] = []
-
-  for (const line of lines) {
-    if (line.trim().startsWith('```')) {
-      if (in_code_block) {
-        if (first_line_of_block) {
-          const completion_info = extract_path_and_position(first_line_of_block)
-          if (completion_info) {
-            const { workspace_name, relative_path } =
-              extract_workspace_and_path({
-                raw_file_path: completion_info.path,
-                is_single_root_folder_workspace:
-                  params.is_single_root_folder_workspace
-              })
-            return {
-              type: 'completion',
-              file_path: relative_path,
-              content: cleanup_api_response({
-                content: content_lines.join('\n')
-              }),
-              line: completion_info.line,
-              character: completion_info.character,
-              workspace_name
-            }
-          }
+  let code_block_start_index = -1
+  let code_block_end_index = -1
+
+  for (let i = 0; i < lines.length; i++) {
+    if (lines[i].trim().startsWith('```')) {
+      code_block_start_index = i
+      for (let j = i + 1; j < lines.length; j++) {
+        if (lines[j].trim().startsWith('```')) {
+          code_block_end_index = j
+          break
         }
-        in_code_block = false
-        first_line_of_block = null
-        content_lines = []
-      } else {
-        in_code_block = true
       }
-      continue
+      break
     }
+  }
 
-    if (in_code_block) {
-      if (first_line_of_block === null) {
-        first_line_of_block = line
-      } else {
-        content_lines.push(line)
-      }
-    }
+  if (code_block_start_index === -1) {
+    return null
   }
 
-  // Handle unclosed block at end of file
-  if (in_code_block && first_line_of_block) {
-    const completion_info = extract_path_and_position(first_line_of_block)
-    if (completion_info) {
-      // This case is not handled and will result in returning null.
-    }
+  if (code_block_end_index === -1) {
+    code_block_end_index = lines.length - 1
   }
 
-  return null
+  const first_line_of_block = lines[code_block_start_index + 1]
+  const completion_info = extract_path_and_position(first_line_of_block)
+
+  if (!completion_info) {
+    return null
+  }
+
+  const results: (CompletionItem | TextItem)[] = []
+  const text_before = lines.slice(0, code_block_start_index).join('\n').trim()
+  if (text_before) {
+    results.push({ type: 'text', content: text_before })
+  }
+
+  const { workspace_name, relative_path } = extract_workspace_and_path({
+    raw_file_path: completion_info.path,
+    is_single_root_folder_workspace: params.is_single_root_folder_workspace
+  })
+
+  const content_lines = lines.slice(
+    code_block_start_index + 2,
+    code_block_end_index
+  )
+
+  results.push({
+    type: 'completion',
+    file_path: relative_path,
+    content: cleanup_api_response({ content: content_lines.join('\n') }),
+    line: completion_info.line,
+    character: completion_info.character,
+    workspace_name
+  })
+
+  const text_after = lines
+    .slice(code_block_end_index + 1)
+    .join('\n')
+    .trim()
+  if (text_after) {
+    results.push({ type: 'text', content: text_after })
+  }
+
+  return results
 }

packages/vscode/src/commands/apply-chat-response-command/utils/clipboard-parser/parsers/multiple-files-parser.ts

@@ -294,7 +294,7 @@ export const parse_multiple_files = (params: {
         }
         backtick_nesting_level++
         if (
-          (current_language === 'markdown' || current_language === 'md') &&
+          (current_language == 'markdown' || current_language === 'md') &&
           backtick_nesting_level > 1
         ) {
           is_markdown_container_block = true
@@ -318,7 +318,7 @@ export const parse_multiple_files = (params: {
           return false
         })()
       ) {
-        if (current_language === 'markdown' || current_language === 'md') {
+        if (current_language == 'markdown' || current_language === 'md') {
           is_markdown_container_block = true
         }
         backtick_nesting_level++