fix: add re-auth when auto trigger access denied

Author: leigaol

packages/amazonq/src/app/inline/completion.ts

@@ -319,6 +319,7 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
                 position,
                 context,
                 token,
+                isAutoTrigger,
                 getAllRecommendationsOptions
             )
             // get active item from session for displaying

packages/amazonq/src/app/inline/recommendationService.ts

@@ -12,7 +12,7 @@ import { CancellationToken, InlineCompletionContext, Position, TextDocument } fr
 import { LanguageClient } from 'vscode-languageclient'
 import { SessionManager } from './sessionManager'
 import { InlineGeneratingMessage } from './inlineGeneratingMessage'
-import { CodeWhispererStatusBarManager } from 'aws-core-vscode/codewhisperer'
+import { AuthUtil, CodeWhispererStatusBarManager } from 'aws-core-vscode/codewhisperer'
 import { TelemetryHelper } from './telemetryHelper'
 import { ICursorUpdateRecorder } from './cursorUpdateManager'
 import { globals, getLogger } from 'aws-core-vscode/shared'
@@ -28,7 +28,6 @@ export class RecommendationService {
         private readonly inlineGeneratingMessage: InlineGeneratingMessage,
         private cursorUpdateRecorder?: ICursorUpdateRecorder
     ) {}
-
     /**
      * Set the recommendation service
      */
@@ -42,6 +41,7 @@ export class RecommendationService {
         position: Position,
         context: InlineCompletionContext,
         token: CancellationToken,
+        isAutoTrigger: boolean,
         options: GetAllRecommendationsOptions = { emitTelemetry: true, showUi: true }
     ) {
         // Record that a regular request is being made
@@ -131,8 +131,20 @@ export class RecommendationService {
             this.sessionManager.closeSession()
             TelemetryHelper.instance.setAllPaginationEndTime()
             options.emitTelemetry && TelemetryHelper.instance.tryRecordClientComponentLatency()
-        } catch (error) {
+        } catch (error: any) {
             getLogger().error('Error getting recommendations: %O', error)
+            // bearer token expired
+            if (error.data && error.data.awsErrorCode === 'E_AMAZON_Q_CONNECTION_EXPIRED') {
+                // ref: https://github.com/aws/aws-toolkit-vscode/blob/amazonq/v1.74.0/packages/core/src/codewhisperer/service/inlineCompletionService.ts#L104
+                // show re-auth once if connection expired
+                if (AuthUtil.instance.isConnectionExpired()) {
+                    await AuthUtil.instance.notifyReauthenticate(isAutoTrigger)
+                } else {
+                    // get a new bearer token, if this failed, the connection will be marked as expired.
+                    // new tokens will be synced per 10 seconds in auth.startTokenRefreshInterval
+                    await AuthUtil.instance.getBearerToken()
+                }
+            }
             return []
         } finally {
             // Remove all UI indicators if UI is enabled

packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts

@@ -138,7 +138,14 @@ describe('RecommendationService', () => {
 
             sendRequestStub.resolves(mockFirstResult)
 
-            await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true
+            )
 
             // Verify sendRequest was called with correct parameters
             assert(sendRequestStub.calledOnce)
@@ -172,7 +179,14 @@ describe('RecommendationService', () => {
             sendRequestStub.onFirstCall().resolves(mockFirstResult)
             sendRequestStub.onSecondCall().resolves(mockSecondResult)
 
-            await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true
+            )
 
             // Verify sendRequest was called with correct parameters
             assert(sendRequestStub.calledTwice)
@@ -204,7 +218,14 @@ describe('RecommendationService', () => {
 
             sendRequestStub.resolves(mockFirstResult)
 
-            await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true
+            )
 
             // Verify recordCompletionRequest was called
             // eslint-disable-next-line @typescript-eslint/unbound-method
@@ -232,10 +253,18 @@ describe('RecommendationService', () => {
             const { showGeneratingStub, hideGeneratingStub } = setupUITest()
 
             // Call with showUi: false option
-            await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken, {
-                showUi: false,
-                emitTelemetry: true,
-            })
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true,
+                {
+                    showUi: false,
+                    emitTelemetry: true,
+                }
+            )
 
             // Verify UI methods were not called
             sinon.assert.notCalled(showGeneratingStub)
@@ -248,7 +277,14 @@ describe('RecommendationService', () => {
             const { showGeneratingStub, hideGeneratingStub } = setupUITest()
 
             // Call with default options (showUi: true)
-            await service.getAllRecommendations(languageClient, mockDocument, mockPosition, mockContext, mockToken)
+            await service.getAllRecommendations(
+                languageClient,
+                mockDocument,
+                mockPosition,
+                mockContext,
+                mockToken,
+                true
+            )
 
             // Verify UI methods were called
             sinon.assert.calledOnce(showGeneratingStub)
@@ -284,6 +320,7 @@ describe('RecommendationService', () => {
                     mockPosition,
                     mockContext,
                     mockToken,
+                    true,
                     options
                 )