fix(chat): document code command now properly indents generated comments especially python docstrings

agent/src/TestClient.ts

@@ -357,6 +357,14 @@
         this.registerRequest('textDocument/show', () => {
             return Promise.resolve(true)
         })
+        this.registerRequest('textEditor/selection', () => {
+            // No-op handler for textEditor/selection requests in test environment
+            return Promise.resolve(null)
+        })
+        this.registerRequest('textEditor/revealRange', () => {
+            // No-op handler for textEditor/revealRange requests in test environment
+            return Promise.resolve(null)
+        })
         this.registerRequest('editTask/getUserInput', async params => {
             return {
                 instruction: this.userInput.instruction ?? params.instruction,
@@ -840,7 +848,7 @@
                     await this.printDiffAgainstClosestMatchingRecording(error)
                 }
                 const errorMessage = missingRecordingErrors[0].error?.split?.('\n')?.[0]
                 throw new Error(
                     dedent`${errorMessage}.
 
                            To fix this problem, run the following commands to update the HTTP recordings:

agent/src/__snapshots__/document-code.test.ts.snap

@@ -51,8 +51,8 @@ export class TestClass {
     constructor(private shouldGreet: boolean) {}
 
     /**
-         * Prints a greeting message if shouldGreet is true.
-         */
+     * Prints a greeting message if shouldGreet is true.
+     */
     public functionName() {
         if (this.shouldGreet) {
             console.log(/* CURSOR */ 'Hello World!')

agent/src/vscode-shim.test.ts

@@ -434,6 +434,82 @@ describe('vscode.workspace.getConfiguration', () => {
     })
 })
 
+describe('vscode.commands.executeCommand', () => {
+    it('handles null arguments without throwing TypeError', async () => {
+        // This test demonstrates the fix for the issue found in document-code.test.ts
+        // where args[0] could be null, causing a TypeError when checking Symbol.iterator
+
+        let callbackInvoked = false
+        let receivedArg: any
+
+        // Register a test command that accepts null arguments
+        const disposable = vscode.commands.registerCommand('test.command.with.null', arg => {
+            callbackInvoked = true
+            receivedArg = arg
+            return 'success'
+        })
+
+        try {
+            // Execute command with null argument - this should not throw
+            const result = await vscode.commands.executeCommand('test.command.with.null', null)
+
+            expect(callbackInvoked).toBe(true)
+            expect(receivedArg).toBe(null)
+            expect(result).toBe('success')
+        } finally {
+            disposable.dispose()
+        }
+    })
+
+    it('handles undefined arguments without throwing TypeError', async () => {
+        let callbackInvoked = false
+        let receivedArg: any
+
+        const disposable = vscode.commands.registerCommand('test.command.with.undefined', arg => {
+            callbackInvoked = true
+            receivedArg = arg
+            return 'success'
+        })
+
+        try {
+            // Execute command with undefined argument
+            const result = await vscode.commands.executeCommand('test.command.with.undefined', undefined)
+
+            expect(callbackInvoked).toBe(true)
+            expect(receivedArg).toBe(undefined)
+            expect(result).toBe('success')
+        } finally {
+            disposable.dispose()
+        }
+    })
+
+    it('handles iterable objects correctly', async () => {
+        let callbackInvoked = false
+        let receivedArgs: any[] = []
+
+        const disposable = vscode.commands.registerCommand('test.command.with.iterable', (...args) => {
+            callbackInvoked = true
+            receivedArgs = args
+            return 'success'
+        })
+
+        try {
+            // Execute command with an iterable array - should spread the arguments
+            const iterableArg = ['arg1', 'arg2', 'arg3']
+            const result = await vscode.commands.executeCommand(
+                'test.command.with.iterable',
+                iterableArg
+            )
+
+            expect(callbackInvoked).toBe(true)
+            expect(receivedArgs).toEqual(['arg1', 'arg2', 'arg3'])
+            expect(result).toBe('success')
+        } finally {
+            disposable.dispose()
+        }
+    })
+})
+
 describe('workspaces', () => {
     const workspaces = [vscode.Uri.parse('file:///a'), vscode.Uri.parse('file:///b')]
 

agent/src/vscode-shim.ts

@@ -1043,7 +1043,12 @@ const _commands: Partial<typeof vscode.commands> = {
             try {
                 // Handle the case where a single object is passed
                 if (args.length === 1) {
-                    if (typeof args[0] === 'object' && typeof args[0][Symbol.iterator] === 'function') {
+                    // Check for null before accessing Symbol.iterator to avoid TypeError
+                    if (
+                        typeof args[0] === 'object' &&
+                        args[0] !== null &&
+                        typeof args[0][Symbol.iterator] === 'function'
+                    ) {
                         return promisify(registered.callback(...args[0]))
                     }
                     return promisify(registered.callback(args[0]))

vscode/src/edit/prompt/models/claude.ts

@@ -22,7 +22,9 @@ export const claude: EditLLMInteraction = {
     getDoc(options) {
         const docStopSequences = [...SHARED_PARAMETERS.stopSequences]
         const firstLine = options.selectedText.toString().split('\n')[0]
-        if (firstLine.trim().length > 0) {
+        // when claude receives a def function line, it will return an empty response
+        // if we add that first line to the stop sequences
+        if (firstLine.trim().length > 0 && options.document.languageId !== 'python') {
             docStopSequences.push(firstLine)
         }
 

vscode/src/non-stop/FixupController.ts

@@ -24,15 +24,14 @@ import {
     type EditMode,
     EditModeTelemetryMetadataMapping,
 } from '../edit/types'
+import { isStreamedIntent } from '../edit/utils/edit-intent'
 import { getOverriddenModelForIntent } from '../edit/utils/edit-models'
 import type { ExtensionClient } from '../extension-client'
 import { isRunningInsideAgent } from '../jsonrpc/isRunningInsideAgent'
 import { logDebug } from '../output-channel-logger'
 import { charactersLogger } from '../services/CharactersLogger'
 import { splitSafeMetadata } from '../services/telemetry-v2'
 import { countCode } from '../services/utils/code-count'
-
-import { isStreamedIntent } from '../edit/utils/edit-intent'
 import { FixupDocumentEditObserver } from './FixupDocumentEditObserver'
 import type { FixupFile } from './FixupFile'
 import { FixupFileObserver } from './FixupFileObserver'
@@ -952,26 +951,83 @@ export class FixupController
             return false
         }
 
-        // Get the index of the first non-whitespace character on the line where the insertion point is.
-        const nonEmptyStartIndex = document.lineAt(
-            task.insertionPoint.line
-        ).firstNonWhitespaceCharacterIndex
-        // Split the text into lines and prepend each line with spaces to match the indentation level
-        // of the line where the insertion point is.
-        const textLines = text.split('\n').map(line => ' '.repeat(nonEmptyStartIndex) + line)
-        // Join the lines back into a single string with newline characters
-        // Remove any leading whitespace from the first line, as we are inserting at the insertionPoint
-        // Keep any trailing whitespace on the last line to preserve the original indentation.
-        const replacementText = textLines.join('\n').trimStart()
+        // Get the indentation context for the insertion point
+        const insertionLine = document.lineAt(task.insertionPoint.line)
+        let targetIndentSize = insertionLine.firstNonWhitespaceCharacterIndex
+        const textIndentSize = text.search(/\S/) || 0
+
+        const isPythonFile = task.document.languageId === 'python'
+        const isDocIntent = task.intent === 'doc'
+
+        // Special case for Python documentation
+        if (targetIndentSize === 0 && isPythonFile && isDocIntent) {
+            targetIndentSize = 4
+        }
 
-        // Insert the updated text at the specified insertionPoint.
+        // Calculate indentation adjustments
+        const needsIndentAdjustment = targetIndentSize > textIndentSize
+        const indentDifference = needsIndentAdjustment ? targetIndentSize - textIndentSize : 0
+
+        // Process the text lines
+        const textLines = text.split('\n')
+        const processedLines = textLines.map((line, index) => {
+            // Don't add extra indentation to empty lines
+            if (line.trim() === '') {
+                return line
+            }
+
+            // For the first line, only add indentation if we're at the start of a line
+            if (
+                index === 0 &&
+                task.insertionPoint.character > 0 &&
+                line.startsWith(' '.repeat(textIndentSize))
+            ) {
+                return line.trimStart()
+            }
+
+            // Add the calculated indentation difference
+            return ' '.repeat(indentDifference) + line
+        })
+
+        // Ensure proper line ending and handle insertion at line start
+        const proposedText =
+            processedLines.join('\n') +
+            (!isPythonFile
+                ? task.insertionPoint.character > 0
+                    ? ' '.repeat(targetIndentSize)
+                    : ''
+                : '')
+
+        const replacementText = proposedText
+        const startLine = task.insertionPoint.line > 0 ? task.insertionPoint.line - 1 : 0
+        const startLineText = document.lineAt(startLine).text
+
+        // Insert the updated text at the specified insertionPoint
         if (edit instanceof vscode.WorkspaceEdit) {
             edit.insert(document.uri, task.insertionPoint, replacementText)
             return vscode.workspace.applyEdit(edit)
         }
 
+        // If we have a doc intent with python file and no start line text,
+        // we want to insert the docstring after the insertion point.
+        // This is because we need the docstring to be on the next line
+        // after the function or class definition.
+        const insertionPoint = new vscode.Position(
+            isPythonFile && isDocIntent && !startLineText
+                ? task.insertionPoint.line + 1
+                : task.insertionPoint.line,
+            task.insertionPoint.character
+        )
+
         return edit(editBuilder => {
-            editBuilder.insert(task.insertionPoint, replacementText)
+            // Replace the code block if the start line matches the start of the text
+            // This happens sometimes with python document code action where instead
+            // of adding only the docstring, the LLM returns the entire code block
+            if (startLine > 0 && startLineText && text.startsWith(startLineText)) {
+                editBuilder.replace(task.originalRange, text)
+            } else {
+                editBuilder.insert(insertionPoint, replacementText)
+            }
         }, options)
     }
 

vscode/src/tree-sitter/query-sdk.ts

@@ -227,8 +227,10 @@ function getLanguageSpecificQueryWrappers(
              * For all other cases, docstrings should be attached above the symbol range, use this.
              */
             const docStringLine =
-                languageId === 'python' && insertionPoint
-                    ? insertionPoint.node.startPosition.row + 1
+                languageId === 'python'
+                    ? insertionPoint
+                        ? insertionPoint.node.startPosition.row + 1
+                        : start.row + 1
                     : start.row - 1
             const docstringCaptures = queries.documentableNodes.compiled
                 .captures(root, {