Skip to content

Commit 172a170

Browse files
authored
refactor: cleanup codewhispererServer.ts (#2028)
1 parent 18bbc2c commit 172a170

File tree

8 files changed

+78
-99
lines changed

8 files changed

+78
-99
lines changed

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeDiffTracker.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { distance } from 'fastest-levenshtein'
22
import { Position } from '@aws/language-server-runtimes/server-interface'
33
import { Features } from '../types'
44
import { getErrorMessage, getUnmodifiedAcceptedTokens } from '../../shared/utils'
5+
import { CodewhispererLanguage } from '../../shared/languageDetection'
56

67
export interface AcceptedSuggestionEntry {
78
fileUrl: string
@@ -12,6 +13,15 @@ export interface AcceptedSuggestionEntry {
1213
customizationArn?: string
1314
}
1415

16+
export interface AcceptedInlineSuggestionEntry extends AcceptedSuggestionEntry {
17+
sessionId: string
18+
requestId: string
19+
languageId: CodewhispererLanguage
20+
completionType: string
21+
triggerType: string
22+
credentialStartUrl?: string | undefined
23+
}
24+
1525
export interface CodeDiffTrackerOptions {
1626
flushInterval?: number
1727
timeElapsedThreshold?: number

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { TestFeatures } from '@aws/language-server-runtimes/testing'
1212
import * as assert from 'assert'
1313
import { AWSError } from 'aws-sdk'
1414
import sinon, { StubbedInstance } from 'ts-sinon'
15-
import { CONTEXT_CHARACTERS_LIMIT, CodewhispererServerFactory } from './codeWhispererServer'
15+
import { CodewhispererServerFactory } from './codeWhispererServer'
1616
import {
1717
CodeWhispererServiceBase,
1818
CodeWhispererServiceToken,
@@ -58,8 +58,9 @@ import { initBaseTestServiceManager, TestAmazonQServiceManager } from '../../sha
5858
import { LocalProjectContextController } from '../../shared/localProjectContextController'
5959
import { URI } from 'vscode-uri'
6060
import { INVALID_TOKEN } from '../../shared/constants'
61-
import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../shared/amazonQServiceManager/errors'
61+
import { AmazonQError } from '../../shared/amazonQServiceManager/errors'
6262
import * as path from 'path'
63+
import { CONTEXT_CHARACTERS_LIMIT } from './constants'
6364

6465
const updateConfiguration = async (
6566
features: TestFeatures,

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts

Lines changed: 3 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,17 @@ import {
1212
TextDocument,
1313
ResponseError,
1414
LSPErrorCodes,
15-
WorkspaceFolder,
1615
} from '@aws/language-server-runtimes/server-interface'
1716
import { autoTrigger, getAutoTriggerType, getNormalizeOsName, triggerType } from './auto-trigger/autoTrigger'
1817
import {
1918
CodeWhispererServiceToken,
2019
GenerateSuggestionsRequest,
2120
GenerateSuggestionsResponse,
21+
getFileContext,
2222
Suggestion,
2323
SuggestionType,
2424
} from '../../shared/codeWhispererService'
25-
import { CodewhispererLanguage, getRuntimeLanguage, getSupportedLanguageId } from '../../shared/languageDetection'
25+
import { getSupportedLanguageId } from '../../shared/languageDetection'
2626
import { mergeEditSuggestionsWithFileContext, truncateOverlapWithRightContext } from './mergeRightUtils'
2727
import { CodeWhispererSession, SessionManager } from './session/sessionManager'
2828
import { CodePercentageTracker } from './codePercentage'
@@ -31,7 +31,7 @@ import { getIdeCategory, makeUserContextObject } from '../../shared/telemetryUti
3131
import { fetchSupplementalContext } from '../../shared/supplementalContextUtil/supplementalContextUtil'
3232
import { textUtils } from '@aws/lsp-core'
3333
import { TelemetryService } from '../../shared/telemetry/telemetryService'
34-
import { AcceptedSuggestionEntry, CodeDiffTracker } from './codeDiffTracker'
34+
import { AcceptedInlineSuggestionEntry, CodeDiffTracker } from './codeDiffTracker'
3535
import {
3636
AmazonQError,
3737
AmazonQServiceConnectionExpiredError,
@@ -44,7 +44,6 @@ import { hasConnectionExpired } from '../../shared/utils'
4444
import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager'
4545
import { WorkspaceFolderManager } from '../workspaceContext/workspaceFolderManager'
4646
import path = require('path')
47-
import { getRelativePath } from '../workspaceContext/util'
4847
import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
4948
import { RecentEditTracker, RecentEditTrackerDefaultConfig } from './tracker/codeEditTracker'
5049
import { CursorTracker } from './tracker/cursorTracker'
@@ -59,48 +58,6 @@ import {
5958
const { editPredictionAutoTrigger } = require('./auto-trigger/editPredictionAutoTrigger')
6059

6160
const EMPTY_RESULT = { sessionId: '', items: [] }
62-
export const FILE_URI_CHARS_LIMIT = 1024
63-
export const FILENAME_CHARS_LIMIT = 1024
64-
export const CONTEXT_CHARACTERS_LIMIT = 10240
65-
66-
// Both clients (token, sigv4) define their own types, this return value needs to match both of them.
67-
const getFileContext = (params: {
68-
textDocument: TextDocument
69-
position: Position
70-
inferredLanguageId: CodewhispererLanguage
71-
workspaceFolder: WorkspaceFolder | null | undefined
72-
}): {
73-
fileUri: string
74-
filename: string
75-
programmingLanguage: {
76-
languageName: CodewhispererLanguage
77-
}
78-
leftFileContent: string
79-
rightFileContent: string
80-
} => {
81-
const left = params.textDocument.getText({
82-
start: { line: 0, character: 0 },
83-
end: params.position,
84-
})
85-
const right = params.textDocument.getText({
86-
start: params.position,
87-
end: params.textDocument.positionAt(params.textDocument.getText().length),
88-
})
89-
90-
const relativeFilePath = params.workspaceFolder
91-
? getRelativePath(params.workspaceFolder, params.textDocument.uri)
92-
: path.basename(params.textDocument.uri)
93-
94-
return {
95-
fileUri: params.textDocument.uri.substring(0, FILE_URI_CHARS_LIMIT),
96-
filename: relativeFilePath.substring(0, FILENAME_CHARS_LIMIT),
97-
programmingLanguage: {
98-
languageName: getRuntimeLanguage(params.inferredLanguageId),
99-
},
100-
leftFileContent: left,
101-
rightFileContent: right,
102-
}
103-
}
10461

10562
const mergeSuggestionsWithRightContext = (
10663
rightFileContext: string,
@@ -143,16 +100,6 @@ const mergeSuggestionsWithRightContext = (
143100
})
144101
}
145102

146-
interface AcceptedInlineSuggestionEntry extends AcceptedSuggestionEntry {
147-
sessionId: string
148-
requestId: string
149-
languageId: CodewhispererLanguage
150-
customizationArn?: string
151-
completionType: string
152-
triggerType: string
153-
credentialStartUrl?: string | undefined
154-
}
155-
156103
export const CodewhispererServerFactory =
157104
(serviceManager: () => AmazonQBaseServiceManager): Server =>
158105
({ credentialsProvider, lsp, workspace, telemetry, logging, runtime, sdkInitializator }) => {
@@ -219,12 +166,6 @@ export const CodewhispererServerFactory =
219166
...currentSession.requestContext,
220167
fileContext: {
221168
...currentSession.requestContext.fileContext,
222-
leftFileContent: currentSession.requestContext.fileContext.leftFileContent
223-
.slice(-CONTEXT_CHARACTERS_LIMIT)
224-
.replaceAll('\r\n', '\n'),
225-
rightFileContent: currentSession.requestContext.fileContext.rightFileContent
226-
.slice(0, CONTEXT_CHARACTERS_LIMIT)
227-
.replaceAll('\r\n', '\n'),
228169
},
229170
nextToken: `${params.partialResultToken}`,
230171
})
@@ -338,7 +279,6 @@ export const CodewhispererServerFactory =
338279
workspace,
339280
logging,
340281
token,
341-
amazonQServiceManager,
342282
params.openTabFilepaths
343283
)
344284
: Promise.resolve(undefined)
@@ -517,15 +457,6 @@ export const CodewhispererServerFactory =
517457

518458
const generateCompletionReq = {
519459
...requestContext,
520-
fileContext: {
521-
...requestContext.fileContext,
522-
leftFileContent: requestContext.fileContext.leftFileContent
523-
.slice(-CONTEXT_CHARACTERS_LIMIT)
524-
.replaceAll('\r\n', '\n'),
525-
rightFileContent: requestContext.fileContext.rightFileContent
526-
.slice(0, CONTEXT_CHARACTERS_LIMIT)
527-
.replaceAll('\r\n', '\n'),
528-
},
529460
...(workspaceId ? { workspaceId: workspaceId } : {}),
530461
}
531462
try {
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
export const FILE_URI_CHARS_LIMIT = 1024
2+
export const FILENAME_CHARS_LIMIT = 1024
3+
export const CONTEXT_CHARACTERS_LIMIT = 10240

server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ import {
77
SDKInitializator,
88
CancellationToken,
99
CancellationTokenSource,
10+
TextDocument,
11+
Position,
12+
WorkspaceFolder,
1013
} from '@aws/language-server-runtimes/server-interface'
1114
import { waitUntil } from '@aws/lsp-core/out/util/timeoutUtils'
1215
import { AWSError, ConfigurationOptions, CredentialProviderChain, Credentials } from 'aws-sdk'
@@ -26,6 +29,14 @@ import CodeWhispererSigv4Client = require('../client/sigv4/codewhisperersigv4cli
2629
import CodeWhispererTokenClient = require('../client/token/codewhispererbearertokenclient')
2730
import { getErrorId } from './utils'
2831
import { GenerateCompletionsResponse } from '../client/token/codewhispererbearertokenclient'
32+
import { CodewhispererLanguage, getRuntimeLanguage } from './languageDetection'
33+
import { getRelativePath } from '../language-server/workspaceContext/util'
34+
import * as path from 'path'
35+
import {
36+
CONTEXT_CHARACTERS_LIMIT,
37+
FILE_URI_CHARS_LIMIT,
38+
FILENAME_CHARS_LIMIT,
39+
} from '../language-server/inline-completion/constants'
2940

3041
export interface Suggestion extends CodeWhispererTokenClient.Completion, CodeWhispererSigv4Client.Recommendation {
3142
itemId: string
@@ -56,6 +67,47 @@ export interface GenerateSuggestionsResponse {
5667
responseContext: ResponseContext
5768
}
5869

70+
export function getFileContext(params: {
71+
textDocument: TextDocument
72+
position: Position
73+
inferredLanguageId: CodewhispererLanguage
74+
workspaceFolder: WorkspaceFolder | null | undefined
75+
}): {
76+
fileUri: string
77+
filename: string
78+
programmingLanguage: {
79+
languageName: CodewhispererLanguage
80+
}
81+
leftFileContent: string
82+
rightFileContent: string
83+
} {
84+
const left = params.textDocument.getText({
85+
start: { line: 0, character: 0 },
86+
end: params.position,
87+
})
88+
const trimmedLeft = left.slice(-CONTEXT_CHARACTERS_LIMIT).replaceAll('\r\n', '\n')
89+
90+
const right = params.textDocument.getText({
91+
start: params.position,
92+
end: params.textDocument.positionAt(params.textDocument.getText().length),
93+
})
94+
const trimmedRight = right.slice(0, CONTEXT_CHARACTERS_LIMIT).replaceAll('\r\n', '\n')
95+
96+
const relativeFilePath = params.workspaceFolder
97+
? getRelativePath(params.workspaceFolder, params.textDocument.uri)
98+
: path.basename(params.textDocument.uri)
99+
100+
return {
101+
fileUri: params.textDocument.uri.substring(0, FILE_URI_CHARS_LIMIT),
102+
filename: relativeFilePath.substring(0, FILENAME_CHARS_LIMIT),
103+
programmingLanguage: {
104+
languageName: getRuntimeLanguage(params.inferredLanguageId),
105+
},
106+
leftFileContent: trimmedLeft,
107+
rightFileContent: trimmedRight,
108+
}
109+
}
110+
59111
// This abstract class can grow in the future to account for any additional changes across the clients
60112
export abstract class CodeWhispererServiceBase {
61113
protected readonly codeWhispererRegion

server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.test.ts

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -279,8 +279,7 @@ describe('crossFileContextUtil', function () {
279279
document,
280280
{ line: 0, character: 0 },
281281
features.workspace,
282-
fakeCancellationToken,
283-
amazonQServiceManager
282+
fakeCancellationToken
284283
)
285284
assert.deepStrictEqual(result, {
286285
supplementalContextItems: [],
@@ -308,8 +307,7 @@ describe('crossFileContextUtil', function () {
308307
document,
309308
{ line: 0, character: 0 },
310309
features.workspace,
311-
fakeCancellationToken,
312-
amazonQServiceManager
310+
fakeCancellationToken
313311
)
314312

315313
sinon.assert.notCalled(instanceStub)
@@ -338,8 +336,7 @@ describe('crossFileContextUtil', function () {
338336
document,
339337
{ line: 0, character: 0 },
340338
features.workspace,
341-
fakeCancellationToken,
342-
amazonQServiceManager
339+
fakeCancellationToken
343340
)
344341
assert.deepStrictEqual(result, {
345342
supplementalContextItems: [{ content: 'someOtherContet', filePath: '/path/', score: 29.879 }],
@@ -364,8 +361,7 @@ describe('crossFileContextUtil', function () {
364361
document,
365362
{ line: 0, character: 0 },
366363
features.workspace,
367-
fakeCancellationToken,
368-
amazonQServiceManager
364+
fakeCancellationToken
369365
)
370366
assert.deepStrictEqual(result, {
371367
supplementalContextItems: [

server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.ts

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ import { LocalProjectContextController } from '../../shared/localProjectContextC
2929
import { QueryInlineProjectContextRequestV2 } from 'local-indexing'
3030
import { URI } from 'vscode-uri'
3131
import { waitUntil } from '@aws/lsp-core/out/util/timeoutUtils'
32-
import { AmazonQBaseServiceManager } from '../amazonQServiceManager/BaseAmazonQServiceManager'
3332

3433
type CrossFileSupportedLanguage =
3534
| 'java'
@@ -67,7 +66,6 @@ export async function fetchSupplementalContextForSrc(
6766
position: Position,
6867
workspace: Workspace,
6968
cancellationToken: CancellationToken,
70-
amazonQServiceManager?: AmazonQBaseServiceManager,
7169
openTabFiles?: string[]
7270
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
7371
const supplementalContextConfig = getSupplementalContextConfig(document.languageId)
@@ -77,14 +75,7 @@ export async function fetchSupplementalContextForSrc(
7775
}
7876
//TODO: add logic for other strategies once available
7977
if (supplementalContextConfig === 'codemap') {
80-
return await codemapContext(
81-
document,
82-
position,
83-
workspace,
84-
cancellationToken,
85-
amazonQServiceManager,
86-
openTabFiles
87-
)
78+
return await codemapContext(document, position, workspace, cancellationToken, openTabFiles)
8879
}
8980
return { supplementalContextItems: [], strategy: 'Empty' }
9081
}
@@ -94,7 +85,6 @@ export async function codemapContext(
9485
position: Position,
9586
workspace: Workspace,
9687
cancellationToken: CancellationToken,
97-
amazonQServiceManager?: AmazonQBaseServiceManager,
9888
openTabFiles?: string[]
9989
): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
10090
let strategy: SupplementalContextStrategy = 'Empty'
@@ -108,7 +98,7 @@ export async function codemapContext(
10898

10999
const projectContextPromise = waitUntil(
110100
async function () {
111-
return await fetchProjectContext(document, position, 'codemap', amazonQServiceManager)
101+
return await fetchProjectContext(document, position, 'codemap')
112102
},
113103
{ timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
114104
)
@@ -146,8 +136,7 @@ export async function codemapContext(
146136
export async function fetchProjectContext(
147137
document: TextDocument,
148138
position: Position,
149-
target: 'default' | 'codemap' | 'bm25',
150-
amazonQServiceManager?: AmazonQBaseServiceManager
139+
target: 'default' | 'codemap' | 'bm25'
151140
): Promise<CodeWhispererSupplementalContextItem[]> {
152141
const inputChunk: Chunk = getInputChunk(document, position, crossFileContextConfig.numberOfLinesEachChunk)
153142
const fsPath = URI.parse(document.uri).fsPath

server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.ts

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import {
1212
} from '@aws/language-server-runtimes/server-interface'
1313
import { crossFileContextConfig, supplementalContextTimeoutInMs } from '../models/constants'
1414
import * as os from 'os'
15-
import { AmazonQBaseServiceManager } from '../amazonQServiceManager/BaseAmazonQServiceManager'
1615
import { TestIntentDetector } from './unitTestIntentDetection'
1716
import { FocalFileResolver } from './focalFileResolution'
1817
import * as fs from 'fs'
@@ -29,7 +28,6 @@ export async function fetchSupplementalContext(
2928
workspace: Workspace,
3029
logging: Logging,
3130
cancellationToken: CancellationToken,
32-
amazonQServiceManager?: AmazonQBaseServiceManager,
3331
openTabFiles?: string[]
3432
): Promise<CodeWhispererSupplementalContext | undefined> {
3533
const timesBeforeFetching = performance.now()
@@ -74,7 +72,6 @@ export async function fetchSupplementalContext(
7472
position,
7573
workspace,
7674
cancellationToken,
77-
amazonQServiceManager,
7875
openTabFiles
7976
)
8077
}

0 commit comments

Comments
 (0)