fix(amazonq): Seperate inline completion from regular codewhisperer activation

jpinkney-aws · jpinkney-aws · commit bdc49ebb0fe9 · 2025-01-28T09:53:54.000-05:00
diff --git a/packages/amazonq/src/extension.ts b/packages/amazonq/src/extension.ts
@@ -4,7 +4,11 @@
  */
 
 import { AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
-import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
+import {
+    activate as activateCodeWhisperer,
+    shutdown as shutdownCodeWhisperer,
+    activateInlineCompletion,
+} from 'aws-core-vscode/codewhisperer'
 import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
 import { CommonAuthWebview } from 'aws-core-vscode/login'
 import {
@@ -31,6 +35,7 @@ import {
     setContext,
     setupUninstallHandler,
     maybeShowMinVscodeWarning,
+    Experiments,
 } from 'aws-core-vscode/shared'
 import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -39,6 +44,7 @@ import * as semver from 'semver'
 import * as vscode from 'vscode'
 import { registerCommands } from './commands'
 import { focusAmazonQPanel } from 'aws-core-vscode/codewhispererChat'
+import { activate as activateAmazonqLsp } from './lsp/activation'
 
 export const amazonQContextPrefix = 'amazonq'
 
@@ -113,7 +119,13 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     const extContext = {
         extensionContext: context,
     }
+    // This contains every lsp agnostic things (auth, security scan, code scan)
     await activateCodeWhisperer(extContext as ExtContext)
+    if (Experiments.instance.get('amazonqLSP', false)) {
+        await activateAmazonqLsp(context)
+    } else {
+        await activateInlineCompletion()
+    }
 
     // Generic extension commands
     registerGenericCommands(context, amazonQContextPrefix)
diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts
@@ -7,15 +7,7 @@ import * as vscode from 'vscode'
 import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension'
 import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
 import { activate as activateQGumby } from 'aws-core-vscode/amazonqGumby'
-import {
-    ExtContext,
-    globals,
-    CrashMonitoring,
-    getLogger,
-    isNetworkError,
-    isSageMaker,
-    Experiments,
-} from 'aws-core-vscode/shared'
+import { ExtContext, globals, CrashMonitoring, getLogger, isNetworkError, isSageMaker } from 'aws-core-vscode/shared'
 import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode'
 import { updateDevMode } from 'aws-core-vscode/dev'
 import { CommonAuthViewProvider } from 'aws-core-vscode/login'
@@ -29,7 +21,6 @@ import { beta } from 'aws-core-vscode/dev'
 import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
 import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
-import { activate as activateAmazonqLsp } from './lsp/activation'
 
 export async function activate(context: vscode.ExtensionContext) {
     // IMPORTANT: No other code should be added to this function. Place it in one of the following 2 functions where appropriate.
@@ -52,12 +43,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
         extensionContext: context,
     }
 
-    if (Experiments.instance.get('amazonqLSP', false)) {
-        await activateAmazonqLsp(context)
-    } else {
-        await activateCWChat(context)
-        await activateQGumby(extContext as ExtContext)
-    }
+    await activateCWChat(context)
+    await activateQGumby(extContext as ExtContext)
 
     const authProvider = new CommonAuthViewProvider(
         context,
diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts
@@ -98,12 +98,12 @@ import { setContext } from '../shared/vscode/setContext'
 import { syncSecurityIssueWebview } from './views/securityIssue/securityIssueWebview'
 import { detectCommentAboveLine } from '../shared/utilities/commentUtils'
 import { UserWrittenCodeTracker } from './tracker/userWrittenCodeTracker'
+import globals from '../shared/extensionGlobals'
 
 let localize: nls.LocalizeFunc
 
 export async function activate(context: ExtContext): Promise<void> {
     localize = nls.loadMessageBundle()
-    const codewhispererSettings = CodeWhispererSettings.instance
 
     // initialize AuthUtil earlier to make sure it can listen to connection change events.
     const auth = AuthUtil.instance
@@ -303,16 +303,6 @@ export async function activate(context: ExtContext): Promise<void> {
                 SecurityIssueProvider.instance.issues.some((group) => group.issues.some((issue) => issue.visible))
             void setContext('aws.amazonq.security.noMatches', noMatches)
         }),
-        // manual trigger
-        Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-            invokeRecommendation(
-                vscode.window.activeTextEditor as vscode.TextEditor,
-                client,
-                await getConfigEntry()
-            ).catch((e) => {
-                getLogger().error('invokeRecommendation failed: %s', (e as Error).message)
-            })
-        }),
         // select customization
         selectCustomizationPrompt.register(),
         // notify new customizations
@@ -480,6 +470,52 @@ export async function activate(context: ExtContext): Promise<void> {
         })
     }
 
+    void FeatureConfigProvider.instance.fetchFeatureConfigs().catch((error) => {
+        getLogger().error('Failed to fetch feature configs - %s', error)
+    })
+
+    await Commands.tryExecute('aws.amazonq.refreshConnectionCallback')
+    container.ready()
+
+    function setSubscriptionsForCodeIssues() {
+        context.extensionContext.subscriptions.push(
+            vscode.workspace.onDidChangeTextDocument(async (e) => {
+                if (e.document.uri.scheme !== 'file') {
+                    return
+                }
+                const diagnostics = securityScanRender.securityDiagnosticCollection?.get(e.document.uri)
+                if (!diagnostics || diagnostics.length === 0) {
+                    return
+                }
+                disposeSecurityDiagnostic(e)
+
+                SecurityIssueProvider.instance.handleDocumentChange(e)
+                SecurityIssueTreeViewProvider.instance.refresh()
+                await syncSecurityIssueWebview(context)
+
+                toggleIssuesVisibility((issue, filePath) =>
+                    filePath !== e.document.uri.fsPath
+                        ? issue.visible
+                        : !detectCommentAboveLine(
+                              e.document,
+                              issue.startLine,
+                              CodeWhispererConstants.amazonqIgnoreNextLine
+                          )
+                )
+            })
+        )
+    }
+}
+
+export async function activateInlineCompletion() {
+    const codewhispererSettings = CodeWhispererSettings.instance
+    const client = new codewhispererClient.DefaultCodeWhispererClient()
+
+    if (isInlineCompletionEnabled()) {
+        await setSubscriptionsforInlineCompletion()
+        await AuthUtil.instance.setVscodeContextProps()
+    }
+
     function getAutoTriggerStatus(): boolean {
         return CodeSuggestionsState.instance.isSuggestionsEnabled()
     }
@@ -500,17 +536,12 @@ export async function activate(context: ExtContext): Promise<void> {
         }
     }
 
-    if (isInlineCompletionEnabled()) {
-        await setSubscriptionsforInlineCompletion()
-        await AuthUtil.instance.setVscodeContextProps()
-    }
-
     async function setSubscriptionsforInlineCompletion() {
         RecommendationHandler.instance.subscribeSuggestionCommands()
         /**
          * Automated trigger
          */
-        context.extensionContext.subscriptions.push(
+        globals.context.subscriptions.push(
             vscode.window.onDidChangeActiveTextEditor(async (editor) => {
                 await RecommendationHandler.instance.onEditorChange()
             }),
@@ -560,42 +591,16 @@ export async function activate(context: ExtContext): Promise<void> {
                 if (!RecommendationHandler.instance.isSuggestionVisible()) {
                     await KeyStrokeHandler.instance.processKeyStroke(e, editor, client, await getConfigEntry())
                 }
-            })
-        )
-    }
-
-    void FeatureConfigProvider.instance.fetchFeatureConfigs().catch((error) => {
-        getLogger().error('Failed to fetch feature configs - %s', error)
-    })
-
-    await Commands.tryExecute('aws.amazonq.refreshConnectionCallback')
-    container.ready()
-
-    function setSubscriptionsForCodeIssues() {
-        context.extensionContext.subscriptions.push(
-            vscode.workspace.onDidChangeTextDocument(async (e) => {
-                if (e.document.uri.scheme !== 'file') {
-                    return
-                }
-                const diagnostics = securityScanRender.securityDiagnosticCollection?.get(e.document.uri)
-                if (!diagnostics || diagnostics.length === 0) {
-                    return
-                }
-                disposeSecurityDiagnostic(e)
-
-                SecurityIssueProvider.instance.handleDocumentChange(e)
-                SecurityIssueTreeViewProvider.instance.refresh()
-                await syncSecurityIssueWebview(context)
-
-                toggleIssuesVisibility((issue, filePath) =>
-                    filePath !== e.document.uri.fsPath
-                        ? issue.visible
-                        : !detectCommentAboveLine(
-                              e.document,
-                              issue.startLine,
-                              CodeWhispererConstants.amazonqIgnoreNextLine
-                          )
-                )
+            }),
+            // manual trigger
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                invokeRecommendation(
+                    vscode.window.activeTextEditor as vscode.TextEditor,
+                    client,
+                    await getConfigEntry()
+                ).catch((e) => {
+                    getLogger().error('invokeRecommendation failed: %s', (e as Error).message)
+                })
             })
         )
     }
diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-export { activate, shutdown } from './activation'
+export { activate, shutdown, activateInlineCompletion } from './activation'
 export * from './util/authUtil'
 export * from './models/model'
 export * from './models/constants'
