Truncate left and right file contents for suggestions requests

aws · Nov 15, 2023 · 3079840 · 3079840
1 parent 883a749
commit 3079840
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.test.ts
@@ -172,6 +172,35 @@ class HelloWorld
             sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
         })
 
+        it('should truncate left and right context', async () => {
+            const BIG_FILE_CONTENT = '123456789\n'.repeat(5000)
+            const BIG_FILE = TextDocument.create('file:///big_file.cs', 'csharp', 1, BIG_FILE_CONTENT)
+            const cutOffLine = 2000
+            features.openDocument(BIG_FILE)
+
+            await features.doInlineCompletionWithReferences(
+                {
+                    textDocument: { uri: BIG_FILE.uri },
+                    position: { line: cutOffLine, character: 1 },
+                    context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+                },
+                CancellationToken.None
+            )
+            const leftContentChecker = (leftContent: string) =>
+                leftContent.length == codeWhispererServer.CONTEXT_CHARACTERS_LIMIT && leftContent.endsWith('\n1')
+            const rightContentChecker = (rightContent: string) =>
+                rightContent.length == codeWhispererServer.CONTEXT_CHARACTERS_LIMIT && rightContent.startsWith('234')
+
+            sinon.assert.calledWith(
+                service.generateSuggestions,
+                sinon.match.hasNested('fileContext.leftFileContent', sinon.match(leftContentChecker))
+            )
+            sinon.assert.calledWith(
+                service.generateSuggestions,
+                sinon.match.hasNested('fileContext.rightFileContent', sinon.match(rightContentChecker))
+            )
+        })
+
         it('should return recommendations when using a different languageId casing', async () => {
             const result = await features.doInlineCompletionWithReferences(
                 {

diff --git a/server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.ts b/server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.ts
@@ -37,6 +37,7 @@ interface InvocationContext {
 }
 
 const EMPTY_RESULT = { sessionId: '', items: [] }
+export const CONTEXT_CHARACTERS_LIMIT = 10240
 
 // Both clients (token, sigv4) define their own types, this return value needs to match both of them.
 const getFileContext = (params: {
@@ -241,8 +242,12 @@ export const CodewhispererServerFactory =
                         ...requestContext,
                         fileContext: {
                             ...requestContext.fileContext,
-                            leftFileContent: requestContext.fileContext.leftFileContent.replaceAll('\r\n', '\n'),
-                            rightFileContent: requestContext.fileContext.rightFileContent.replaceAll('\r\n', '\n'),
+                            leftFileContent: requestContext.fileContext.leftFileContent
+                                .slice(-CONTEXT_CHARACTERS_LIMIT)
+                                .replaceAll('\r\n', '\n'),
+                            rightFileContent: requestContext.fileContext.rightFileContent
+                                .slice(0, CONTEXT_CHARACTERS_LIMIT)
+                                .replaceAll('\r\n', '\n'),
                         },
                     })
                     .catch(emitServiceInvocationFailure({ telemetry, invocationContext }))