feat: support sending requests with the 'external_idp' type (aws#2247)

Author: rli

package-lock.json

@@ -236,23 +236,24 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
             this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities
                 ?.textDocument?.inlineCompletionWithReferences?.endpointOverride
 
-        // Connection type changed to 'builderId'
-
-        if (newConnectionType === 'builderId') {
-            this.log('Detected New connection type: builderId')
+        // Connection type changed to 'builderId' | 'external_idp'
+        // for now pretend External IdP is just a special case of Builder ID where the subscription has already been established
+        // and user does not need a profile
+        if (newConnectionType === 'builderId' || newConnectionType === 'external_idp') {
+            this.log(`Detected New connection type: ${newConnectionType}`)
             this.resetCodewhispererService()
 
             // For the builderId connection type regional endpoint discovery chain is:
             // region set by client -> runtime region -> default region
             const clientParams = this.features.lsp.getClientInitializeParams()
 
             this.createCodewhispererServiceInstances(
-                'builderId',
+                newConnectionType,
                 clientParams?.initializationOptions?.aws?.region,
                 endpointOverride
             )
             this.state = 'INITIALIZED'
-            this.log('Initialized Amazon Q service with builderId connection')
+            this.log(`Initialized Amazon Q service with ${newConnectionType} connection`)
 
             // Emit auth success event
             ProfileStatusMonitor.emitAuthSuccess()
@@ -506,7 +507,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
     }
 
     private createCodewhispererServiceInstances(
-        connectionType: 'builderId' | 'identityCenter',
+        connectionType: Exclude<SsoConnectionType, 'none'>,
         clientOrProfileRegion: string | undefined,
         endpointOverride: string | undefined
     ) {

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

@@ -1,11 +1,11 @@
 import * as assert from 'assert'
 import sinon, { StubbedInstance, stubInterface } from 'ts-sinon'
 import { CodeWhispererServiceToken } from '../codeWhispererService'
-import { SsoConnectionType } from '../utils'
 import {
     AWSInitializationOptions,
     CancellationTokenSource,
     Logging,
+    SsoConnectionType,
 } from '@aws/language-server-runtimes/server-interface'
 import {
     AmazonQDeveloperProfile,

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.test.ts

@@ -4,8 +4,9 @@ import {
     Logging,
     LSPErrorCodes,
     ResponseError,
+    SsoConnectionType,
 } from '@aws/language-server-runtimes/server-interface'
-import { isBool, isObject, SsoConnectionType } from '../utils'
+import { isBool, isObject } from '../utils'
 import { AWS_Q_ENDPOINTS } from '../../shared/constants'
 import { CodeWhispererServiceToken } from '../codeWhispererService'
 import { AmazonQServiceProfileThrottlingError } from './errors'

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

@@ -344,6 +344,10 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
                             if (!creds?.token) {
                                 throw new Error('Authorization failed, bearer token is not set')
                             }
+                            if (credentialsProvider.getConnectionType() === 'external_idp') {
+                                httpRequest.headers['TokenType'] = 'EXTERNAL_IDP'
+                            }
+
                             httpRequest.headers['Authorization'] = `Bearer ${creds.token}`
                             httpRequest.headers['x-amzn-codewhisperer-optout'] =
                                 `${!this.shareCodeWhispererContentWithAWS}`

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

@@ -88,6 +88,19 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase {
             }),
             customUserAgent: customUserAgent,
         })
+
+        this.client.middlewareStack.add(
+            (next, context) => args => {
+                if (credentialsProvider.getConnectionType() === 'external_idp') {
+                    // @ts-ignore
+                    args.request.headers['TokenType'] = 'EXTERNAL_IDP'
+                }
+                return next(args)
+            },
+            {
+                step: 'build',
+            }
+        )
     }
 
     public async sendMessage(

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

@@ -1,11 +1,11 @@
 import { TextDocument } from 'vscode-languageserver-textdocument'
 import { CodeWhispererServiceBase, ResponseContext, Suggestion } from './codeWhispererService'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
-import { SsoConnectionType } from './utils'
 import { stubInterface } from 'ts-sinon'
 import { StreamingClientServiceBase } from './streamingClientService'
 import { SessionData } from '../language-server/inline-completion/session/sessionManager'
 import { WorkspaceFolder } from '@aws/language-server-runtimes/protocol'
+import { SsoConnectionType } from '@aws/language-server-runtimes/server-interface'
 
 export const HELLO_WORLD_IN_CSHARP = `class HelloWorld
 {

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

@@ -3,6 +3,7 @@ import {
     BearerCredentials,
     CredentialsProvider,
     Position,
+    SsoConnectionType,
 } from '@aws/language-server-runtimes/server-interface'
 import { AWSError, Credentials } from 'aws-sdk'
 import { distance } from 'fastest-levenshtein'
@@ -31,7 +32,6 @@ import { getAuthFollowUpType } from '../language-server/chat/utils'
 import ignore = require('ignore')
 import { InitializeParams } from '@aws/language-server-runtimes/server-interface'
 import { QClientCapabilities } from '../language-server/configuration/qConfigurationServer'
-export type SsoConnectionType = 'builderId' | 'identityCenter' | 'none'
 
 export function isAwsError(error: unknown): error is AWSError {
     if (error === undefined) {
@@ -436,9 +436,7 @@ export const flattenMetric = (obj: any, prefix = '') => {
 }
 
 export function getSsoConnectionType(credentialsProvider: CredentialsProvider): SsoConnectionType {
-    const connectionMetadata = credentialsProvider.getConnectionMetadata()
-    const startUrl = connectionMetadata?.sso?.startUrl
-    return !startUrl ? 'none' : startUrl.includes(BUILDER_ID_START_URL) ? 'builderId' : 'identityCenter'
+    return credentialsProvider.getConnectionType()
 }
 
 // Port of implementation in AWS Toolkit for VSCode