Allow SSO users to login and use pro tier

Author: ahusseinali

packages/core/src/auth/activation.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import * as vscode from 'vscode'
 import { Auth } from './auth'
 import { LoginManager } from './deprecated/loginManager'
 import { fromString } from './providers/credentials'
@@ -17,9 +18,16 @@ import { Ec2CredentialsProvider } from './providers/ec2CredentialsProvider'
 import { EcsCredentialsProvider } from './providers/ecsCredentialsProvider'
 import { EnvVarsCredentialsProvider } from './providers/envVarsCredentialsProvider'
 
+interface SagemakerCookie {
+    authMode?: string
+}
+
 export async function initialize(loginManager: LoginManager): Promise<void> {
     if (isAmazonQ() && isSageMaker()) {
-        initializeCredentialsProviderManager()
+        const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie
+        if (result.authMode !== 'Sso') {
+            initializeCredentialsProviderManager()
+        }
     }
     Auth.instance.onDidChangeActiveConnection(async (conn) => {
         // This logic needs to be moved to `Auth.useConnection` to correctly record `passive`

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

@@ -14,7 +14,7 @@ import { AWSError } from 'aws-sdk'
 import { isAwsError } from '../../shared/errors'
 import { TelemetryHelper } from '../util/telemetryHelper'
 import { getLogger } from '../../shared/logger'
-import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
+import { isCloud9 } from '../../shared/extensionUtilities'
 import { hasVendedIamCredentials } from '../../auth/auth'
 import {
     asyncCallWithTimeout,
@@ -157,9 +157,7 @@ export class RecommendationHandler {
         config: ConfigurationEntry,
         autoTriggerType?: CodewhispererAutomatedTriggerType,
         pagination: boolean = true,
-        page: number = 0,
-        isSM: boolean = isSageMaker(),
-        retry: boolean = false
+        page: number = 0
     ): Promise<GetRecommendationsResponse> {
         let invocationResult: 'Succeeded' | 'Failed' = 'Failed'
         let errorMessage: string | undefined = undefined
@@ -186,7 +184,7 @@ export class RecommendationHandler {
         ).language
         session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName)
 
-        if (pagination && !isSM) {
+        if (pagination) {
             if (page === 0) {
                 session.requestContext = await EditorContext.buildListRecommendationRequest(
                     editor as vscode.TextEditor,
@@ -240,8 +238,9 @@ export class RecommendationHandler {
             startTime = performance.now()
             this.lastInvocationTime = startTime
             const mappedReq = runtimeLanguageContext.mapToRuntimeLanguage(request)
-            const codewhispererPromise =
-                pagination && !isSM ? client.listRecommendations(mappedReq) : client.generateRecommendations(mappedReq)
+            const codewhispererPromise = pagination
+                ? 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
@@ -333,8 +332,7 @@ export class RecommendationHandler {
                         config,
                         autoTriggerType,
                         pagination,
-                        page,
-                        true
+                        page
                     )
                 }
             }