Skip to content

Commit b29478f

Browse files
authored
revert: fix to add opt-out header to streaming client (#2365) (#2370)
This reverts commit 692e77b.
1 parent 692e77b commit b29478f

File tree

6 files changed

+3
-142
lines changed

6 files changed

+3
-142
lines changed

server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ type DeferredHandler = {
3838
reject: (err: Error) => void
3939
}
4040
export class ChatSessionService {
41+
public shareCodeWhispererContentWithAWS = false
4142
public pairProgrammingMode: boolean = true
4243
public contextListSent: boolean = false
4344
public modelId: string | undefined

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,6 @@ export class AmazonQIAMServiceManager extends BaseAmazonQServiceManager<
7474
this.region,
7575
this.endpoint
7676
)
77-
this.cachedStreamingClient.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
78-
'shareCodeWhispererContentWithAWS'
79-
)
8077
}
8178
return this.cachedStreamingClient
8279
}

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -584,9 +584,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
584584
this.getCustomUserAgent()
585585
)
586586
streamingClient.profileArn = this.activeIdcProfile?.arn
587-
streamingClient.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
588-
'shareCodeWhispererContentWithAWS'
589-
)
590587

591588
this.logging.debug(`Created streaming client instance region=${region}, endpoint=${endpoint}`)
592589
return streamingClient

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -157,17 +157,6 @@ export abstract class BaseAmazonQServiceManager<
157157
)
158158
this.cachedCodewhispererService.shareCodeWhispererContentWithAWS = shareCodeWhispererContentWithAWS
159159
}
160-
161-
if (this.cachedStreamingClient) {
162-
const shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
163-
'shareCodeWhispererContentWithAWS'
164-
)
165-
this.logging.debug(
166-
'Update shareCodeWhispererContentWithAWS setting on cachedStreamingClient to ' +
167-
shareCodeWhispererContentWithAWS
168-
)
169-
this.cachedStreamingClient.shareCodeWhispererContentWithAWS = shareCodeWhispererContentWithAWS
170-
}
171160
}
172161

173162
private async notifyDidChangeConfigurationListeners(): Promise<void> {

server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts

Lines changed: 0 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@ import {
1212
} from '@amzn/codewhisperer-streaming'
1313
import { QDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client'
1414
import { rejects } from 'assert'
15-
import { initBaseTestServiceManager, TestAmazonQServiceManager } from './amazonQServiceManager/testUtils'
16-
import { stubCodeWhispererService } from './testUtils'
1715

1816
const TIME_TO_ADVANCE_MS = 100
1917

@@ -115,33 +113,6 @@ describe('StreamingClientServiceToken', () => {
115113
sinon.assert.match(sendMessageStub.firstCall.firstArg, expectedRequest)
116114
})
117115

118-
it('creates client with shareCodeWhispererContentWithAWS parameter', () => {
119-
const streamingClientServiceWithOptout = new StreamingClientServiceToken(
120-
features.credentialsProvider,
121-
features.sdkInitializator,
122-
features.logging,
123-
DEFAULT_AWS_Q_REGION,
124-
DEFAULT_AWS_Q_ENDPOINT_URL,
125-
'some-user-agent'
126-
)
127-
streamingClientServiceWithOptout.shareCodeWhispererContentWithAWS = false
128-
129-
expect(streamingClientServiceWithOptout['shareCodeWhispererContentWithAWS']).to.equal(false)
130-
})
131-
132-
it('creates client without shareCodeWhispererContentWithAWS parameter', () => {
133-
const streamingClientServiceDefault = new StreamingClientServiceToken(
134-
features.credentialsProvider,
135-
features.sdkInitializator,
136-
features.logging,
137-
DEFAULT_AWS_Q_REGION,
138-
DEFAULT_AWS_Q_ENDPOINT_URL,
139-
'some-user-agent'
140-
)
141-
142-
expect(streamingClientServiceDefault['shareCodeWhispererContentWithAWS']).to.be.undefined
143-
})
144-
145116
describe('generateAssistantResponse', () => {
146117
const MOCKED_GENERATE_RESPONSE_REQUEST = {
147118
conversationState: {
@@ -347,79 +318,4 @@ describe('StreamingClientServiceIAM', () => {
347318
expect(credentials.expiration).to.be.instanceOf(Date)
348319
expect(credentials.expiration.getTime()).to.be.closeTo(Date.now(), 1000)
349320
})
350-
351-
it('creates client with shareCodeWhispererContentWithAWS parameter', () => {
352-
const streamingClientServiceWithOptout = new StreamingClientServiceIAM(
353-
features.credentialsProvider,
354-
features.sdkInitializator,
355-
features.logging,
356-
DEFAULT_AWS_Q_REGION,
357-
DEFAULT_AWS_Q_ENDPOINT_URL
358-
)
359-
streamingClientServiceWithOptout.shareCodeWhispererContentWithAWS = false
360-
361-
expect(streamingClientServiceWithOptout['shareCodeWhispererContentWithAWS']).to.equal(false)
362-
})
363-
364-
it('creates client without shareCodeWhispererContentWithAWS parameter', () => {
365-
const streamingClientServiceDefault = new StreamingClientServiceIAM(
366-
features.credentialsProvider,
367-
features.sdkInitializator,
368-
features.logging,
369-
DEFAULT_AWS_Q_REGION,
370-
DEFAULT_AWS_Q_ENDPOINT_URL
371-
)
372-
373-
expect(streamingClientServiceDefault['shareCodeWhispererContentWithAWS']).to.be.undefined
374-
})
375-
})
376-
377-
describe('BaseAmazonQServiceManager streaming client cache updates', () => {
378-
let features: TestFeatures
379-
let serviceManager: TestAmazonQServiceManager
380-
let streamingClientMock: StreamingClientServiceToken
381-
382-
beforeEach(() => {
383-
features = new TestFeatures()
384-
const serviceStub = stubCodeWhispererService()
385-
386-
streamingClientMock = Object.assign(sinon.createStubInstance(StreamingClientServiceToken), {
387-
region: DEFAULT_AWS_Q_REGION,
388-
endpoint: DEFAULT_AWS_Q_ENDPOINT_URL,
389-
}) as unknown as StreamingClientServiceToken
390-
serviceManager = initBaseTestServiceManager(features, serviceStub, streamingClientMock)
391-
})
392-
393-
afterEach(() => {
394-
sinon.restore()
395-
TestAmazonQServiceManager.resetInstance()
396-
})
397-
398-
it('updates shareCodeWhispererContentWithAWS on cached streaming client when configuration changes', async () => {
399-
// Set initial configuration
400-
features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: true })
401-
402-
await serviceManager.handleDidChangeConfiguration()
403-
404-
expect(streamingClientMock.shareCodeWhispererContentWithAWS).to.equal(true)
405-
406-
// Change configuration
407-
features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: false })
408-
409-
await serviceManager.handleDidChangeConfiguration()
410-
411-
expect(streamingClientMock.shareCodeWhispererContentWithAWS).to.equal(false)
412-
})
413-
414-
it('does not update streaming client when no cached client exists', async () => {
415-
TestAmazonQServiceManager.resetInstance()
416-
const serviceManagerWithoutClient = initBaseTestServiceManager(features, stubCodeWhispererService())
417-
418-
features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: false })
419-
420-
// Should not throw when no cached streaming client exists
421-
await serviceManagerWithoutClient.handleDidChangeConfiguration()
422-
423-
expect(serviceManagerWithoutClient['cachedStreamingClient']).to.be.undefined
424-
})
425321
})

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

Lines changed: 2 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ export type ChatCommandOutput = SendMessageCommandOutput | GenerateAssistantResp
3939
export abstract class StreamingClientServiceBase {
4040
protected readonly region
4141
protected readonly endpoint
42-
public shareCodeWhispererContentWithAWS?: boolean
4342

4443
inflightRequests: Set<AbortController> = new Set()
4544

@@ -66,7 +65,6 @@ export abstract class StreamingClientServiceBase {
6665
export class StreamingClientServiceToken extends StreamingClientServiceBase {
6766
client: CodeWhispererStreaming
6867
public profileArn?: string
69-
7068
constructor(
7169
credentialsProvider: CredentialsProvider,
7270
sdkInitializator: SDKInitializator,
@@ -76,7 +74,6 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
7674
customUserAgent: string
7775
) {
7876
super(region, endpoint)
79-
8077
const tokenProvider = async () => {
8178
const token = getBearerTokenFromProvider(credentialsProvider)
8279
// without setting expiration, the tokenProvider will only be called once
@@ -98,15 +95,11 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
9895
})
9996

10097
this.client.middlewareStack.add(
101-
(next, context) => (args: any) => {
98+
(next, context) => args => {
10299
if (credentialsProvider.getConnectionType() === 'external_idp') {
100+
// @ts-ignore
103101
args.request.headers['TokenType'] = 'EXTERNAL_IDP'
104102
}
105-
if (this.shareCodeWhispererContentWithAWS !== undefined) {
106-
args.request.headers['x-amzn-codewhisperer-optout'] = `${!this.shareCodeWhispererContentWithAWS}`
107-
}
108-
// Log headers for debugging
109-
logging.debug(`StreamingClient headers: ${JSON.stringify(args.request.headers)}`)
110103
return next(args)
111104
},
112105
{
@@ -220,18 +213,6 @@ export class StreamingClientServiceIAM extends StreamingClientServiceBase {
220213
credentials: iamCredentialProvider,
221214
retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),
222215
})
223-
224-
this.client.middlewareStack.add(
225-
(next, context) => (args: any) => {
226-
if (this.shareCodeWhispererContentWithAWS !== undefined) {
227-
args.request.headers['x-amzn-codewhisperer-optout'] = `${!this.shareCodeWhispererContentWithAWS}`
228-
}
229-
return next(args)
230-
},
231-
{
232-
step: 'build',
233-
}
234-
)
235216
}
236217

237218
public async sendMessage(

0 commit comments

Comments
 (0)