Address review comments

Author: ahusseinali

docs/build.md

@@ -1,11 +1,23 @@
 # Build
 
-The AmazonQ features rely on the `codewhisperer-streaming` service, who's client
-is generated from the service's smithy models and placed in
-`src.gen/@amzn/codewhisperer-streaming` (For more
-information about this client and how it is generated, please see this
+The AmazonQ features rely on the `codewhisperer-streaming` service, to support both Sigv4 and Bearer token modes of this service,
+two clients are generated from the service's smithy models and placed in
+`src.gen/@amzn/amazon-q-developer-streaming-client` and `src.gen/@amzn/codewhisperer-streaming` respectively (For more
+information about these clients and how they are generated, please see this
 [quip document](https://quip-amazon.com/2dAWAvTIYXXr/Build-instructions-for-AWS-CodeWhisperer-Streaming-Typescript-client)).
 
+## @amzn/amazon-q-developer-streaming client
+
+This client is a standalone npm project in typescript, and it is added to
+the project as a workspace in the project's root `package.json` with the line `"workspaces": [ ..., "src.gen/@amzn/amazon-q-developer-streaming" ]`.
+The client may be manually built using `npm run build -w @amzn/amazon-q-developer-streaming"`.
+The `generateClients` run script ensures that this dependency is
+built before the toolkit project itself. Workspaces are automatically ready to
+be imported in the root toolkit project by their declared package.json name,
+(`@amzn/amazon-q-developer-streaming` in this case).
+
+## @amzn/codewhisperer-streaming client
+
 This client is a standalone npm project in typescript, and it is added to
 the project as a workspace in the project's root `package.json` with the line `"workspaces": [ ..., "src.gen/@amzn/codewhisperer-streaming" ]`.
 The client may be manually built using `npm run build -w @amzn/codewhisperer-streaming"`.

packages/core/src/auth/activation.ts

@@ -8,22 +8,18 @@ import { Auth } from './auth'
 import { LoginManager } from './deprecated/loginManager'
 import { fromString } from './providers/credentials'
 import { getLogger } from '../shared/logger'
-import { ExtensionUse } from './utils'
+import { ExtensionUse, initializeCredentialsProviderManager } from './utils'
 import { isAmazonQ, isCloud9, isSageMaker } from '../shared/extensionUtilities'
 import { isInDevEnv } from '../shared/vscode/env'
 import { isWeb } from '../shared/extensionGlobals'
-import { CredentialsProviderManager } from './providers/credentialsProviderManager'
-import { SharedCredentialsProviderFactory } from './providers/sharedCredentialsProviderFactory'
-import { Ec2CredentialsProvider } from './providers/ec2CredentialsProvider'
-import { EcsCredentialsProvider } from './providers/ecsCredentialsProvider'
-import { EnvVarsCredentialsProvider } from './providers/envVarsCredentialsProvider'
 
 interface SagemakerCookie {
-    authMode?: string
+    authMode?: 'Sso' | 'Iam'
 }
 
 export async function initialize(loginManager: LoginManager): Promise<void> {
     if (isAmazonQ() && isSageMaker()) {
+        // The command `sagemaker.parseCookies` is registered in VS Code Sagemaker environment.
         const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie
         if (result.authMode !== 'Sso') {
             initializeCredentialsProviderManager()
@@ -41,12 +37,6 @@ export async function initialize(loginManager: LoginManager): Promise<void> {
     await showManageConnectionsOnStartup()
 }
 
-function initializeCredentialsProviderManager() {
-    const manager = CredentialsProviderManager.getInstance()
-    manager.addProviderFactory(new SharedCredentialsProviderFactory())
-    manager.addProviders(new Ec2CredentialsProvider(), new EcsCredentialsProvider(), new EnvVarsCredentialsProvider())
-}
-
 /**
  * Show the Manage Connections page when the extension starts up, if it should be shown.
  */

packages/core/src/auth/utils.ts

@@ -52,6 +52,11 @@ import { ExtStartUpSources } from '../shared/telemetry'
 import { CommonAuthWebview } from '../login/webview/vue/backend'
 import { AuthSource } from '../login/webview/util'
 import { setContext } from '../shared/vscode/setContext'
+import { CredentialsProviderManager } from './providers/credentialsProviderManager'
+import { SharedCredentialsProviderFactory } from './providers/sharedCredentialsProviderFactory'
+import { Ec2CredentialsProvider } from './providers/ec2CredentialsProvider'
+import { EcsCredentialsProvider } from './providers/ecsCredentialsProvider'
+import { EnvVarsCredentialsProvider } from './providers/envVarsCredentialsProvider'
 
 // iam-only excludes Builder ID and IAM Identity Center from the list of valid connections
 // TODO: Understand if "iam" should include these from the list at all
@@ -731,3 +736,9 @@ export function getAuthFormIdsFromConnection(conn?: Connection): AuthFormId[] {
 
     return authIds
 }
+
+export function initializeCredentialsProviderManager() {
+    const manager = CredentialsProviderManager.getInstance()
+    manager.addProviderFactory(new SharedCredentialsProviderFactory())
+    manager.addProviders(new Ec2CredentialsProvider(), new EcsCredentialsProvider(), new EnvVarsCredentialsProvider())
+}

packages/core/src/codewhisperer/service/recommendationHandler.ts

@@ -43,6 +43,7 @@ import { application } from '../util/codeWhispererApplication'
 import { openUrl } from '../../shared/utilities/vsCodeUtils'
 import { indent } from '../../shared/utilities/textUtilities'
 import path from 'path'
+import { isIamConnection } from '../../auth/connection'
 
 /**
  * This class is for getRecommendation/listRecommendation API calls and its states
@@ -157,7 +158,8 @@ export class RecommendationHandler {
         config: ConfigurationEntry,
         autoTriggerType?: CodewhispererAutomatedTriggerType,
         pagination: boolean = true,
-        page: number = 0
+        page: number = 0,
+        generate: boolean = isIamConnection(AuthUtil.instance.conn)
     ): Promise<GetRecommendationsResponse> {
         let invocationResult: 'Succeeded' | 'Failed' = 'Failed'
         let errorMessage: string | undefined = undefined
@@ -184,7 +186,7 @@ export class RecommendationHandler {
         ).language
         session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName)
 
-        if (pagination) {
+        if (pagination && !generate) {
             if (page === 0) {
                 session.requestContext = await EditorContext.buildListRecommendationRequest(
                     editor as vscode.TextEditor,
@@ -238,9 +240,10 @@ export class RecommendationHandler {
             startTime = performance.now()
             this.lastInvocationTime = startTime
             const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request)
-            const codewhispererPromise = pagination
-                ? client.listRecommendations(mappedReq)
-                : client.generateRecommendations(mappedReq)
+            const codewhispererPromise =
+                pagination && !generate
+                    ? client.listRecommendations(mappedReq)
+                    : client.generateRecommendations(mappedReq)
             const resp = await this.getServerResponse(triggerType, config.isManualTriggerEnabled, codewhispererPromise)
             TelemetryHelper.instance.setSdkApiCallEndTime()
             latency = startTime !== 0 ? performance.now() - startTime : 0
@@ -332,7 +335,8 @@ export class RecommendationHandler {
                         config,
                         autoTriggerType,
                         pagination,
-                        page
+                        page,
+                        true
                     )
                 }
             }

packages/core/src/codewhisperer/util/authUtil.ts

@@ -8,7 +8,7 @@ import * as localizedText from '../../shared/localizedText'
 import { Auth } from '../../auth/auth'
 import { ToolkitError, isNetworkError, tryRun } from '../../shared/errors'
 import { getSecondaryAuth, setScopes } from '../../auth/secondaryAuth'
-import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
+import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { AmazonQPromptSettings } from '../../shared/settings'
 import {
     scopesCodeWhispererCore,
@@ -50,6 +50,9 @@ export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
 export const codeWhispererChatScopes = [...codeWhispererCoreScopes, ...scopesCodeWhispererChat]
 export const amazonQScopes = [...codeWhispererChatScopes, ...scopesGumby, ...scopesFeatureDev]
 
+function isValidSageMakerConnection(conn?: Connection) {
+    return isIamConnection(conn) || isSsoConnection(conn)
+}
 /**
  * "Core" are the CW scopes that existed before the addition of new scopes
  * for Amazon Q.
@@ -59,25 +62,19 @@ export const isValidCodeWhispererCoreConnection = (conn?: Connection): conn is C
         return isIamConnection(conn)
     }
 
-    if (isSageMaker()) {
-        return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))
-    }
-
     return (
+        isValidSageMakerConnection(conn) ||
         (isCloud9('codecatalyst') && isIamConnection(conn)) ||
         (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))
     )
 }
 /** Superset that includes all of CodeWhisperer + Amazon Q */
 export const isValidAmazonQConnection = (conn?: Connection): conn is Connection => {
-    if (isSageMaker() && isAmazonQ()) {
-        return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, amazonQScopes))
-    }
-
     return (
-        (isSsoConnection(conn) || isBuilderIdConnection(conn)) &&
-        isValidCodeWhispererCoreConnection(conn) &&
-        hasScopes(conn, amazonQScopes)
+        isValidSageMakerConnection(conn) ||
+        ((isSsoConnection(conn) || isBuilderIdConnection(conn)) &&
+            isValidCodeWhispererCoreConnection(conn) &&
+            hasScopes(conn, amazonQScopes))
     )
 }
 

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -35,7 +35,7 @@ import { EditorContextCommand } from '../../commands/registerCommands'
 import { PromptsGenerator } from './prompts/promptsGenerator'
 import { TriggerEventsStorage } from '../../storages/triggerEvents'
 import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
-import { CodeWhispererStreamingServiceException, __MetadataBearer } from '@amzn/codewhisperer-streaming'
+import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
 import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
 import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
 import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -18,7 +18,6 @@ import {
     ChatResponseStream as cwChatResponseStream,
     CodeWhispererStreamingServiceException,
     SupplementaryWebLink,
-    __MetadataBearer,
 } from '@amzn/codewhisperer-streaming'
 import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/connector/connector'
 import { ChatSession } from '../../../clients/chat/v0/chat'
@@ -39,7 +38,8 @@ import { extractAuthFollowUp } from '../../../../amazonq/util/authUtils'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
-export type MessengerResponseType = __MetadataBearer & {
+export type MessengerResponseType = {
+    $metadata: { requestId?: string; httpStatusCode?: number }
     message?: AsyncIterable<cwChatResponseStream | qdevChatResponseStream>
 }
 

packages/core/src/extensionNode.ts

@@ -9,8 +9,6 @@ import * as nls from 'vscode-nls'
 import * as codecatalyst from './codecatalyst/activation'
 import { activate as activateAwsExplorer } from './awsexplorer/activation'
 import { activate as activateCloudWatchLogs } from './awsService/cloudWatchLogs/activation'
-import { CredentialsProviderManager } from './auth/providers/credentialsProviderManager'
-import { SharedCredentialsProviderFactory } from './auth/providers/sharedCredentialsProviderFactory'
 import { activate as activateSchemas } from './eventSchemas/activation'
 import { activate as activateLambda } from './lambda/activation'
 import { activate as activateCloudFormationTemplateRegistry } from './shared/cloudformation/activation'
@@ -39,9 +37,6 @@ import { activate as activateDev } from './dev/activation'
 import { activate as activateApplicationComposer } from './applicationcomposer/activation'
 import { activate as activateRedshift } from './awsService/redshift/activation'
 import { activate as activateIamPolicyChecks } from './awsService/accessanalyzer/activation'
-import { Ec2CredentialsProvider } from './auth/providers/ec2CredentialsProvider'
-import { EnvVarsCredentialsProvider } from './auth/providers/envVarsCredentialsProvider'
-import { EcsCredentialsProvider } from './auth/providers/ecsCredentialsProvider'
 import { SchemaService } from './shared/schemas'
 import { AwsResourceManager } from './dynamicResources/awsResourceManager'
 import globals from './shared/extensionGlobals'
@@ -55,7 +50,7 @@ import { learnMoreAmazonQCommand, qExtensionPageCommand, dismissQTree } from './
 import { AuthUtil, codeWhispererCoreScopes, isPreviousQUser } from './codewhisperer/util/authUtil'
 import { installAmazonQExtension } from './codewhisperer/commands/basicCommands'
 import { isExtensionInstalled, VSCODE_EXTENSION_ID } from './shared/utilities'
-import { ExtensionUse } from './auth/utils'
+import { ExtensionUse, initializeCredentialsProviderManager } from './auth/utils'
 import { ExtStartUpSources } from './shared/telemetry'
 import { activate as activateThreatComposerEditor } from './threatComposer/activation'
 import { isSsoConnection, hasScopes } from './auth/connection'
@@ -310,12 +305,6 @@ async function handleAmazonQInstall() {
     })
 }
 
-function initializeCredentialsProviderManager() {
-    const manager = CredentialsProviderManager.getInstance()
-    manager.addProviderFactory(new SharedCredentialsProviderFactory())
-    manager.addProviders(new Ec2CredentialsProvider(), new EcsCredentialsProvider(), new EnvVarsCredentialsProvider())
-}
-
 function recordToolkitInitialization(activationStartedOn: number, settingsValid: boolean, logger?: Logger) {
     try {
         const activationFinishedOn = Date.now()

packages/toolkit/package.json

@@ -709,7 +709,7 @@
                 {
                     "id": "aws.amazonq.codewhisperer",
                     "name": "%AWS.amazonq.codewhisperer.title%",
-                    "when": "!isCloud9 && !isSageMaker && !aws.toolkit.amazonq.dismissed && !aws.explorer.showAuthView"
+                    "when": "!isCloud9 && !aws.isSageMaker && !aws.toolkit.amazonq.dismissed && !aws.explorer.showAuthView"
                 },
                 {
                     "id": "aws.explorer",