Skip to content

Commit

Permalink
feat(amazonq): handle virtual spaces when inserting code to cursor po…
Browse files Browse the repository at this point in the history 
…sition
  • Loading branch information
Ege Ozcan committed Jan 2, 2025
1 parent f74263c commit 7391990
Show file tree
Hide file tree
Showing 2 changed files with 356 additions and 2 deletions.
336 changes: 336 additions & 0 deletions server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ import {
Telemetry,
Logging,
Workspace,
Position,
InsertToCursorPositionParams,
TextDocumentEdit,
} from '@aws/language-server-runtimes/server-interface'
import { TestFeatures } from '@aws/language-server-runtimes/testing'
import * as assert from 'assert'
Expand All @@ -21,6 +24,7 @@ import { DocumentContextExtractor } from './contexts/documentContext'
import * as utils from './utils'
import { DEFAULT_HELP_FOLLOW_UP_PROMPT, HELP_MESSAGE } from './constants'
import { TelemetryService } from '../telemetryService'
import { TextEdit } from 'vscode-languageserver-textdocument'

describe('ChatController', () => {
const mockTabId = 'tab-1'
Expand Down Expand Up @@ -469,4 +473,336 @@ describe('ChatController', () => {
})
})
})

describe('onCodeInsertToCursorPosition', () => {
beforeEach(() => {
chatController.onTabAdd({ tabId: mockTabId })
testFeatures.lsp.workspace.applyWorkspaceEdit.resolves({ applied: true })
testFeatures.workspace.getTextDocument = sinon.stub()
})

it('handles regular insertion correctly', async () => {
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, ' ')
testFeatures.workspace.getTextDocument.resolves(document)

const cursorPosition = Position.create(0, 0)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'const x = 1\n const y = 2',
tabId: mockTabId,
messageId: 'XXX',
}
await chatController.onCodeInsertToCursorPosition(params)

assert.deepStrictEqual(testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0], {
edit: {
documentChanges: [
{
textDocument: { uri: 'test.ts', version: 0 },
edits: [
{
range: {
start: cursorPosition,
end: cursorPosition,
},
newText: params.code,
},
],
},
],
},
})
})

it('handles tab-based indentation correctly', async () => {
const documentContent = 'function test() {\n\tif (true) {\n\t\t// cursor here\n\t}'
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, documentContent)
testFeatures.workspace.getTextDocument.resolves(document)

const cursorPosition = Position.create(2, 2)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'console.log("test");\nconsole.log("test2")',
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

const documentChanges = testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0].edit.documentChanges
assert(documentChanges)
const insertedText = (documentChanges[0] as TextDocumentEdit).edits[0].newText
// Should maintain tab-based indentation
assert.deepStrictEqual(insertedText, 'console.log("test");\n\t\tconsole.log("test2")')
})

it('handles insertion at mixed indentation levels correctly', async () => {
const documentContent = `function test() {
if (true) {
// cursor here
console.log("test");
}
}`
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, documentContent)
testFeatures.workspace.getTextDocument.resolves(document)

const cursorPosition = Position.create(2, 12)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'const x = 1;\nconst y = 2;',
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

// Verify that the inserted code maintains the indentation level of the insertion point
const documentChanges = testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0].edit.documentChanges
assert(documentChanges)
const insertedText = (documentChanges[0] as TextDocumentEdit).edits[0].newText
assert.deepStrictEqual(insertedText, `const x = 1;\n${' '.repeat(12)}const y = 2;`)
})

it('handles indentation correctly when inserting after an indent', async () => {
// Text document contains 8 space characters
const documentContent = ' '.repeat(8)
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, documentContent)
testFeatures.workspace.getTextDocument.resolves(document)

// Cursor is positioned at the end of the first line, after the 8 spaces
const cursorPosition = Position.create(0, documentContent.length)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'const x = 1\nconst y = 2',
tabId: mockTabId,
messageId: 'XXX',
}
await chatController.onCodeInsertToCursorPosition(params)

assert.deepStrictEqual(testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0], {
edit: {
documentChanges: [
{
textDocument: { uri: 'test.ts', version: 0 },
edits: [
{
range: {
start: cursorPosition,
end: cursorPosition,
},
// We expect new text to be added to the end of the existing line and also apply indentation on the next line
newText: `const x = 1\n${' '.repeat(8)}const y = 2`,
},
],
},
],
},
})
})

it('handles indentation correctly when inserting at the end of a single line that does not have any indentation', async () => {
const documentContent = 'console.log("Hello world")'
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, documentContent)
testFeatures.workspace.getTextDocument.resolves(document)

const forLoop = `for (let i = 2; i <= n; i++) {
const next = prev + current;
prev = current;
current = next;
}`

const cursorPosition = Position.create(0, documentContent.length)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: forLoop,
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

const documentChanges = testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0].edit.documentChanges
assert(documentChanges)
const insertedText = (documentChanges[0] as TextDocumentEdit).edits[0].newText
// For loop should be inserted as is in this case
assert.deepStrictEqual(insertedText, forLoop)
})

it('handles indentation correctly when inserting inside an indented block', async () => {
const fibonacci = `function fibonacci(n) {
if (n <= 1) return n;
let prev = 0,
let current = 1;
for (let i = 2; i <= n; i++) {
// Insertion will happen on the line below
const next = prev + current;
prev = current;
current = next;
}
return current;
}
`

// This test will insert an extra for loop inside the existing for loop in the fibonacci function above
const forLoop = `for (let i = 2; i <= n; i++) {
const next = prev + current;
prev = current;
current = next;
}
`
// Given the for loop is inside a function and we will be inserting a new for loop inside, the for loop to be inserted will have 8 prefix spaces
const twiceIndentedForLoop = `for (let i = 2; i <= n; i++) {
${' '.repeat(8)} const next = prev + current;
${' '.repeat(8)} prev = current;
${' '.repeat(8)} current = next;
${' '.repeat(8)}}
`

let document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, fibonacci)
testFeatures.workspace.getTextDocument.resolves(document)

const cursorPosition = Position.create(8, 8)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: forLoop,
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

const documentChanges = testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0].edit.documentChanges
assert(documentChanges)
const insertedText = (documentChanges[0] as TextDocumentEdit).edits[0].newText
assert.deepStrictEqual(insertedText, twiceIndentedForLoop)
})

it('handles virtual spaces when cursor is in empty line with virtual indent', async () => {
// Create an empty document
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, '')
testFeatures.workspace.getTextDocument.resolves(document)

// Position cursor at character 8 in an empty line (virtual space)
const cursorPosition = Position.create(0, 8)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'const x = 1\nconst y = 2',
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

// The code should be indented with 8 spaces for both lines
// and cursor should be moved to position 0
assert.deepStrictEqual(testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0], {
edit: {
documentChanges: [
{
textDocument: { uri: 'test.ts', version: 0 },
edits: [
{
range: {
start: Position.create(0, 0), // Note: cursor moved to start
end: Position.create(0, 0),
},
newText: `${' '.repeat(8)}const x = 1\n${' '.repeat(8)}const y = 2`,
},
],
},
],
},
})
})

it('handles virtual spaces with multiline code containing empty lines', async () => {
// Create an empty document
let document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, '')
testFeatures.workspace.getTextDocument.resolves(document)

// Position cursor at character 4 in an empty line (virtual space)
const cursorPosition = Position.create(0, 4)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
code: 'if (condition) {\n\n console.log("test");\n}',
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

// The code should be indented with 4 spaces, empty lines should remain empty
assert.deepStrictEqual(testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0], {
edit: {
documentChanges: [
{
textDocument: { uri: 'test.ts', version: 0 },
edits: [
{
range: {
start: Position.create(0, 0), // Note: cursor moved to start
end: Position.create(0, 0),
},
newText: `${' '.repeat(4)}if (condition) {\n\n${' '.repeat(8)}console.log("test");\n }`,
},
],
},
],
},
})
})

it('handles virtual spaces correctly when code starts with empty line', async () => {
const document: TextDocument = TextDocument.create('test.ts', 'typescript', 1, '')
testFeatures.workspace.getTextDocument.resolves(document)

// Position cursor at character 6 in an empty line (virtual space)
const cursorPosition = Position.create(0, 6)
const params: InsertToCursorPositionParams = {
textDocument: { uri: 'test.ts' },
cursorPosition,
// Code starts with an empty line, followed by actual code
code: '\nfunction test() {\n console.log("test");\n}',
tabId: mockTabId,
messageId: 'XXX',
}

await chatController.onCodeInsertToCursorPosition(params)

// The first empty line should have no indentation
// Subsequent lines should be indented with 6 spaces
assert.deepStrictEqual(testFeatures.lsp.workspace.applyWorkspaceEdit.firstCall.args[0], {
edit: {
documentChanges: [
{
textDocument: { uri: 'test.ts', version: 0 },
edits: [
{
range: {
start: Position.create(0, 0), // Note: cursor moved to start
end: Position.create(0, 0),
},
// First line is empty (no indentation)
// Following lines get the virtual space indentation
newText: `\n${' '.repeat(6)}function test() {\n${' '.repeat(10)}console.log("test");\n${' '.repeat(6)}}`,
},
],
},
],
},
})
})
})
})
22 changes: 20 additions & 2 deletions server/aws-lsp-codewhisperer/src/language-server/chat/chatController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,24 +200,42 @@ export class ChatController implements ChatHandlers {
return
}

const cursorPosition = params.cursorPosition
let cursorPosition = params.cursorPosition

const indentRange = {
start: { line: cursorPosition.line, character: 0 },
end: cursorPosition,
}
const documentContent = await this.#features.workspace.getTextDocument(params.textDocument.uri)
let indent = documentContent?.getText(indentRange)

let hasVirtualSpace = false
if (indent && indent.trim().length !== 0) {
indent = ' '.repeat(indent.length - indent.trimStart().length)
} else if ((isNullish(indent) || !indent) && cursorPosition.character > 0) {
// When text content at indent range is undefined or empty but the cursor is not at position 0
// It means there are virtual spaces that is being rendered by the IDE
// In this case, the indentation is determined by the cursorPosition
this.#log('Indent is nullish and the cursor position is greater than zero while inserting code')
indent = ' '.repeat(cursorPosition.character)
hasVirtualSpace = true
cursorPosition.character = 0
}

let textWithIndent = ''
params.code.split('\n').forEach((line, index) => {
if (index === 0) {
if (hasVirtualSpace && line !== '') {
textWithIndent += indent
}
textWithIndent += line
} else {
textWithIndent += '\n' + indent + line
// Don't put indentation in empty lines
if (line == '') {
textWithIndent += '\n'
} else {
textWithIndent += '\n' + indent + line
}
}
})

Expand Down

0 comments on commit 7391990

Please sign in to comment.