Skip to content

Commit

Permalink
feat: pass supplemental contexts only for token client (#697)
Browse files Browse the repository at this point in the history 
* fix: supplementalContexts passed only with token authentication

* fix: pass supplementalMetadata to telemetry only with token authentication

* test: fix tests after new logic

* fix: rename test variable to prevent overriding variables and dispose features

* style: call fetchSupplementalContext only for token-based server

---------

Co-authored-by: Francesco Piccoli <[email protected]>
  • Loading branch information
@francescoopiccoli
francescoopiccoli and Francesco Piccoli authored Jan 7, 2025
1 parent f2949d4 commit 7242835
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 32 deletions.
49 changes: 35 additions & 14 deletions server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,12 @@ import * as assert from 'assert'
import { AWSError } from 'aws-sdk'
import sinon, { StubbedInstance, stubInterface } from 'ts-sinon'
import { CONTEXT_CHARACTERS_LIMIT, CodewhispererServerFactory } from './codeWhispererServer'
import { CodeWhispererServiceBase, ResponseContext, Suggestion } from './codeWhispererService'
import {
CodeWhispererServiceBase,
CodeWhispererServiceToken,
ResponseContext,
Suggestion,
} from './codeWhispererService'
import { CodeWhispererSession, SessionData, SessionManager } from './session/sessionManager'
import {
EMPTY_RESULT,
Expand Down Expand Up @@ -128,7 +133,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: HELLO_WORLD_IN_CSHARP,
},
maxResults: 5,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -159,7 +163,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: remainingLines,
},
maxResults: 5,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -214,7 +217,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: HELLO_WORLD_IN_CSHARP,
},
maxResults: 5,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -272,7 +274,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: HELLO_WORLD_IN_CSHARP,
},
maxResults: 5,
supplementalContexts: [],
}

// Check the service was called with the right parameters
Expand Down Expand Up @@ -356,7 +357,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: rightContext,
},
maxResults: 5,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -392,7 +392,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: modifiedRightContext,
},
maxResults: 5,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -478,13 +477,35 @@ describe('CodeWhisperer Server', () => {
})

describe('Supplemental Context', () => {
it('should send supplemental context', async () => {
// Open 3 files supporting cross-file context
features
it('should send supplemental context when using token authentication', async () => {
const test_service = sinon.createStubInstance(
CodeWhispererServiceToken
) as StubbedInstance<CodeWhispererServiceToken>

test_service.generateSuggestions.returns(
Promise.resolve({
suggestions: EXPECTED_SUGGESTION,
responseContext: EXPECTED_RESPONSE_CONTEXT,
})
)

const test_server = CodewhispererServerFactory(_auth => test_service)

// Initialize the features, but don't start server yet
const test_features = new TestFeatures()

// Return no specific configuration for CodeWhisperer
test_features.lsp.workspace.getConfiguration.returns(Promise.resolve({}))

// Start the server and open a document
await test_features.start(test_server)

// Open files supporting cross-file context
test_features
.openDocument(TextDocument.create('file:///SampleFile.java', 'java', 1, 'sample-content'))
.openDocument(TextDocument.create('file:///TargetFile.java', 'java', 1, ''))

await features.doInlineCompletionWithReferences(
await test_features.doInlineCompletionWithReferences(
{
textDocument: { uri: 'file:///TargetFile.java' },
position: Position.create(0, 0),
Expand All @@ -506,7 +527,9 @@ describe('CodeWhisperer Server', () => {
{ content: 'sample-content', filePath: '/SampleFile.java' },
],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
sinon.assert.calledOnceWithExactly(test_service.generateSuggestions, expectedGenerateSuggestionsRequest)

test_features.dispose()
})
})

Expand Down Expand Up @@ -946,7 +969,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: SPECIAL_CHARACTER_HELLO_WORLD.substring(1, SPECIAL_CHARACTER_HELLO_WORLD.length),
},
maxResults: 1,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down Expand Up @@ -979,7 +1001,6 @@ describe('CodeWhisperer Server', () => {
rightFileContent: RIGHT_FILE_CONTEXT,
},
maxResults: 1,
supplementalContexts: [],
}
sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
})
Expand Down
29 changes: 16 additions & 13 deletions server/aws-lsp-codewhisperer/src/language-server/codeWhispererServer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -365,23 +365,26 @@ export const CodewhispererServerFactory =
) {
return EMPTY_RESULT
}
const supplementalContext = await fetchSupplementalContext(
textDocument,
params.position,
workspace,
logging,
token
)

const requestContext: GenerateSuggestionsRequest = {
// supplementalContext available only via token authentication
const supplementalContextPromise =
codeWhispererService instanceof CodeWhispererServiceToken
? fetchSupplementalContext(textDocument, params.position, workspace, logging, token)
: Promise.resolve(undefined)

let requestContext: GenerateSuggestionsRequest = {
fileContext,
maxResults,
supplementalContexts: supplementalContext?.supplementalContextItems
}

const supplementalContext = await supplementalContextPromise
if (codeWhispererService instanceof CodeWhispererServiceToken) {
requestContext.supplementalContexts = supplementalContext?.supplementalContextItems
? supplementalContext.supplementalContextItems.map(v => ({
content: v.content,
filePath: v.filePath,
}))
: [],
: []
}

// Close ACTIVE session and record Discard trigger decision immediately
Expand Down Expand Up @@ -603,9 +606,9 @@ export const CodewhispererServerFactory =
`Inline completion configuration updated to use ${codeWhispererService.customizationArn}`
)
/*
The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination
configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true
*/
The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination
configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true
*/
// const enableTelemetryEventsToDestination = true
// telemetryService.updateEnableTelemetryEventsToDestination(enableTelemetryEventsToDestination)
const optOutTelemetryPreference = qConfig['optOutTelemetry'] === true ? 'OPTOUT' : 'OPTIN'
Expand Down
5 changes: 0 additions & 5 deletions server/aws-lsp-codewhisperer/src/language-server/telemetry/userTriggerDecision.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,6 @@ describe('Telemetry', () => {
rightFileContent: ')\n {\n Console.WriteLine("Hello World!");\n }\n}\n',
},
maxResults: 1,
supplementalContexts: [],
},
supplementalMetadata: undefined,
timeToFirstRecommendation: 2000,
Expand Down Expand Up @@ -438,7 +437,6 @@ describe('Telemetry', () => {
'}\n',
},
maxResults: 5,
supplementalContexts: [],
},
})
sinon.assert.calledWithMatch(telemetryServiceSpy, expectedUserTriggerDecisionMetric, 0)
Expand Down Expand Up @@ -878,7 +876,6 @@ describe('Telemetry', () => {
rightFileContent: ')\n {\n Console.WriteLine("Hello World!");\n }\n}\n',
},
maxResults: 1,
supplementalContexts: [],
},
timeToFirstRecommendation: 2000,
credentialStartUrl: 'teststarturl',
Expand Down Expand Up @@ -1031,7 +1028,6 @@ describe('Telemetry', () => {
rightFileContent: ')\n {\n Console.WriteLine("Hello World!");\n }\n}\n',
},
maxResults: 1,
supplementalContexts: [],
},
timeToFirstRecommendation: 2000,
credentialStartUrl: 'teststarturl',
Expand Down Expand Up @@ -1144,7 +1140,6 @@ describe('Telemetry', () => {
rightFileContent: ')\n {\n Console.WriteLine("Hello World!");\n }\n}\n',
},
maxResults: 1,
supplementalContexts: [],
},
timeToFirstRecommendation: 1260,
credentialStartUrl: 'teststarturl',
Expand Down

0 comments on commit 7242835

Please sign in to comment.