chore: merge agentic version 1.36.0 (#2371)

* chore: bump agentic version: 1.36.0

* fix(amazonq): fix to add opt-out header to streaming client (#2365) (#2369)

---------

Co-authored-by: aws-toolkit-automation <>
Co-authored-by: invictus <[email protected]>

Author: XiaoxuanLu

app/aws-lsp-codewhisperer-runtimes/src/version.json

@@ -1,3 +1,3 @@
 {
-    "agenticChat": "1.35.0"
+    "agenticChat": "1.36.0"
 }

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

@@ -38,7 +38,6 @@ type DeferredHandler = {
     reject: (err: Error) => void
 }
 export class ChatSessionService {
-    public shareCodeWhispererContentWithAWS = false
     public pairProgrammingMode: boolean = true
     public contextListSent: boolean = false
     public modelId: string | undefined

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

@@ -74,6 +74,9 @@ export class AmazonQIAMServiceManager extends BaseAmazonQServiceManager<
                 this.region,
                 this.endpoint
             )
+            this.cachedStreamingClient.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
+                'shareCodeWhispererContentWithAWS'
+            )
         }
         return this.cachedStreamingClient
     }

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

@@ -584,6 +584,9 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
             this.getCustomUserAgent()
         )
         streamingClient.profileArn = this.activeIdcProfile?.arn
+        streamingClient.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
+            'shareCodeWhispererContentWithAWS'
+        )
 
         this.logging.debug(`Created streaming client instance region=${region}, endpoint=${endpoint}`)
         return streamingClient

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

@@ -157,6 +157,17 @@ export abstract class BaseAmazonQServiceManager<
             )
             this.cachedCodewhispererService.shareCodeWhispererContentWithAWS = shareCodeWhispererContentWithAWS
         }
+
+        if (this.cachedStreamingClient) {
+            const shareCodeWhispererContentWithAWS = this.configurationCache.getProperty(
+                'shareCodeWhispererContentWithAWS'
+            )
+            this.logging.debug(
+                'Update shareCodeWhispererContentWithAWS setting on cachedStreamingClient to ' +
+                    shareCodeWhispererContentWithAWS
+            )
+            this.cachedStreamingClient.shareCodeWhispererContentWithAWS = shareCodeWhispererContentWithAWS
+        }
     }
 
     private async notifyDidChangeConfigurationListeners(): Promise<void> {

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

@@ -12,6 +12,8 @@ import {
 } from '@amzn/codewhisperer-streaming'
 import { QDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client'
 import { rejects } from 'assert'
+import { initBaseTestServiceManager, TestAmazonQServiceManager } from './amazonQServiceManager/testUtils'
+import { stubCodeWhispererService } from './testUtils'
 
 const TIME_TO_ADVANCE_MS = 100
 
@@ -113,6 +115,33 @@ describe('StreamingClientServiceToken', () => {
         sinon.assert.match(sendMessageStub.firstCall.firstArg, expectedRequest)
     })
 
+    it('creates client with shareCodeWhispererContentWithAWS parameter', () => {
+        const streamingClientServiceWithOptout = new StreamingClientServiceToken(
+            features.credentialsProvider,
+            features.sdkInitializator,
+            features.logging,
+            DEFAULT_AWS_Q_REGION,
+            DEFAULT_AWS_Q_ENDPOINT_URL,
+            'some-user-agent'
+        )
+        streamingClientServiceWithOptout.shareCodeWhispererContentWithAWS = false
+
+        expect(streamingClientServiceWithOptout['shareCodeWhispererContentWithAWS']).to.equal(false)
+    })
+
+    it('creates client without shareCodeWhispererContentWithAWS parameter', () => {
+        const streamingClientServiceDefault = new StreamingClientServiceToken(
+            features.credentialsProvider,
+            features.sdkInitializator,
+            features.logging,
+            DEFAULT_AWS_Q_REGION,
+            DEFAULT_AWS_Q_ENDPOINT_URL,
+            'some-user-agent'
+        )
+
+        expect(streamingClientServiceDefault['shareCodeWhispererContentWithAWS']).to.be.undefined
+    })
+
     describe('generateAssistantResponse', () => {
         const MOCKED_GENERATE_RESPONSE_REQUEST = {
             conversationState: {
@@ -318,4 +347,79 @@ describe('StreamingClientServiceIAM', () => {
         expect(credentials.expiration).to.be.instanceOf(Date)
         expect(credentials.expiration.getTime()).to.be.closeTo(Date.now(), 1000)
     })
+
+    it('creates client with shareCodeWhispererContentWithAWS parameter', () => {
+        const streamingClientServiceWithOptout = new StreamingClientServiceIAM(
+            features.credentialsProvider,
+            features.sdkInitializator,
+            features.logging,
+            DEFAULT_AWS_Q_REGION,
+            DEFAULT_AWS_Q_ENDPOINT_URL
+        )
+        streamingClientServiceWithOptout.shareCodeWhispererContentWithAWS = false
+
+        expect(streamingClientServiceWithOptout['shareCodeWhispererContentWithAWS']).to.equal(false)
+    })
+
+    it('creates client without shareCodeWhispererContentWithAWS parameter', () => {
+        const streamingClientServiceDefault = new StreamingClientServiceIAM(
+            features.credentialsProvider,
+            features.sdkInitializator,
+            features.logging,
+            DEFAULT_AWS_Q_REGION,
+            DEFAULT_AWS_Q_ENDPOINT_URL
+        )
+
+        expect(streamingClientServiceDefault['shareCodeWhispererContentWithAWS']).to.be.undefined
+    })
+})
+
+describe('BaseAmazonQServiceManager streaming client cache updates', () => {
+    let features: TestFeatures
+    let serviceManager: TestAmazonQServiceManager
+    let streamingClientMock: StreamingClientServiceToken
+
+    beforeEach(() => {
+        features = new TestFeatures()
+        const serviceStub = stubCodeWhispererService()
+
+        streamingClientMock = Object.assign(sinon.createStubInstance(StreamingClientServiceToken), {
+            region: DEFAULT_AWS_Q_REGION,
+            endpoint: DEFAULT_AWS_Q_ENDPOINT_URL,
+        }) as unknown as StreamingClientServiceToken
+        serviceManager = initBaseTestServiceManager(features, serviceStub, streamingClientMock)
+    })
+
+    afterEach(() => {
+        sinon.restore()
+        TestAmazonQServiceManager.resetInstance()
+    })
+
+    it('updates shareCodeWhispererContentWithAWS on cached streaming client when configuration changes', async () => {
+        // Set initial configuration
+        features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: true })
+
+        await serviceManager.handleDidChangeConfiguration()
+
+        expect(streamingClientMock.shareCodeWhispererContentWithAWS).to.equal(true)
+
+        // Change configuration
+        features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: false })
+
+        await serviceManager.handleDidChangeConfiguration()
+
+        expect(streamingClientMock.shareCodeWhispererContentWithAWS).to.equal(false)
+    })
+
+    it('does not update streaming client when no cached client exists', async () => {
+        TestAmazonQServiceManager.resetInstance()
+        const serviceManagerWithoutClient = initBaseTestServiceManager(features, stubCodeWhispererService())
+
+        features.lsp.workspace.getConfiguration.resolves({ shareCodeWhispererContentWithAWS: false })
+
+        // Should not throw when no cached streaming client exists
+        await serviceManagerWithoutClient.handleDidChangeConfiguration()
+
+        expect(serviceManagerWithoutClient['cachedStreamingClient']).to.be.undefined
+    })
 })

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

@@ -39,6 +39,7 @@ export type ChatCommandOutput = SendMessageCommandOutput | GenerateAssistantResp
 export abstract class StreamingClientServiceBase {
     protected readonly region
     protected readonly endpoint
+    public shareCodeWhispererContentWithAWS?: boolean
 
     inflightRequests: Set<AbortController> = new Set()
 
@@ -65,6 +66,7 @@ export abstract class StreamingClientServiceBase {
 export class StreamingClientServiceToken extends StreamingClientServiceBase {
     client: CodeWhispererStreaming
     public profileArn?: string
+
     constructor(
         credentialsProvider: CredentialsProvider,
         sdkInitializator: SDKInitializator,
@@ -74,6 +76,7 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
         customUserAgent: string
     ) {
         super(region, endpoint)
+
         const tokenProvider = async () => {
             const token = getBearerTokenFromProvider(credentialsProvider)
             // without setting expiration, the tokenProvider will only be called once
@@ -95,11 +98,15 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
         })
 
         this.client.middlewareStack.add(
-            (next, context) => args => {
+            (next, context) => (args: any) => {
                 if (credentialsProvider.getConnectionType() === 'external_idp') {
-                    // @ts-ignore
                     args.request.headers['TokenType'] = 'EXTERNAL_IDP'
                 }
+                if (this.shareCodeWhispererContentWithAWS !== undefined) {
+                    args.request.headers['x-amzn-codewhisperer-optout'] = `${!this.shareCodeWhispererContentWithAWS}`
+                }
+                // Log headers for debugging
+                logging.debug(`StreamingClient headers: ${JSON.stringify(args.request.headers)}`)
                 return next(args)
             },
             {
@@ -213,6 +220,18 @@ export class StreamingClientServiceIAM extends StreamingClientServiceBase {
             credentials: iamCredentialProvider,
             retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),
         })
+
+        this.client.middlewareStack.add(
+            (next, context) => (args: any) => {
+                if (this.shareCodeWhispererContentWithAWS !== undefined) {
+                    args.request.headers['x-amzn-codewhisperer-optout'] = `${!this.shareCodeWhispererContentWithAWS}`
+                }
+                return next(args)
+            },
+            {
+                step: 'build',
+            }
+        )
     }
 
     public async sendMessage(