Merge master into feature/agentic-chat

docs/lsp.md

@@ -45,8 +45,16 @@ sequenceDiagram
     npm run package
     ```
     to get the project setup
-3. Uncomment the `__AMAZONQLSP_PATH` variable in `amazonq/.vscode/launch.json` Extension configuration
-4. Use the `Launch LSP with Debugging` configuration and set breakpoints in VSCode or the language server
+3. Enable the lsp experiment:
+    ```
+    "aws.experiments": {
+      "amazonqLSP": true,
+      "amazonqLSPInline": true, // optional: enables inline completion from flare
+      "amazonqLSPChat": true // optional: enables chat from flare
+    }
+    ```
+4. Uncomment the `__AMAZONQLSP_PATH` variable in `amazonq/.vscode/launch.json` Extension configuration
+5. Use the `Launch LSP with Debugging` configuration and set breakpoints in VSCode or the language server
 
 ## Amazon Q Inline Activation
 

packages/amazonq/.changes/next-release/Bug Fix-ccbc5e50-0a90-4340-b6d5-e57537949898.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Amazon Q chat: `@prompts` not added to context"
+}

packages/amazonq/.changes/next-release/Feature-390d2096-ea1e-43b9-a889-2494d6e9e6d3.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "SageMaker Unified Studio: Disable Sign out"
+}

packages/amazonq/.changes/next-release/Feature-3c42f632-46b8-4e59-b5fa-75d638f3f08d.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "SageMaker Unified Studio: Update Q Chat Introduction message"
+}

packages/amazonq/.changes/next-release/Feature-e5ed7aac-542d-41ee-8abe-d32db5de73ae.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Amazon Q chat: Code blocks in chat messages have a max-height of 21 lines and can be scrolled inside"
+}

packages/amazonq/package.json

@@ -119,6 +119,14 @@
                         "ssoCacheError": {
                             "type": "boolean",
                             "default": false
+                        },
+                        "amazonQLspManifestMessage": {
+                            "type": "boolean",
+                            "default": false
+                        },
+                        "amazonQWorkspaceLspManifestMessage": {
+                            "type": "boolean",
+                            "default": false
                         }
                     },
                     "additionalProperties": false
@@ -351,7 +359,7 @@
                 },
                 {
                     "command": "aws.amazonq.signout",
-                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected",
+                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",
                     "group": "2_amazonQ@4"
                 },
                 {
@@ -627,7 +635,7 @@
                 "title": "%AWS.command.codewhisperer.signout%",
                 "category": "%AWS.amazonq.title%",
                 "icon": "$(debug-disconnect)",
-                "enablement": "aws.codewhisperer.connected"
+                "enablement": "aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio"
             },
             {
                 "command": "aws.amazonq.learnMore",

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

@@ -29,6 +29,7 @@ import {
     ReferenceHoverProvider,
     ReferenceInlineProvider,
     ReferenceLogViewProvider,
+    ImportAdderProvider,
 } from 'aws-core-vscode/codewhisperer'
 
 export class InlineCompletionManager implements Disposable {
@@ -66,6 +67,7 @@ export class InlineCompletionManager implements Disposable {
             item: InlineCompletionItemWithReferences,
             editor: TextEditor,
             requestStartTime: number,
+            startLine: number,
             firstCompletionDisplayLatency?: number
         ) => {
             // TODO: also log the seen state for other suggestions in session
@@ -96,6 +98,9 @@ export class InlineCompletionManager implements Disposable {
                 ReferenceLogViewProvider.instance.addReferenceLog(referenceLog)
                 ReferenceHoverProvider.instance.addCodeReferences(item.insertText as string, item.references)
             }
+            if (item.mostRelevantMissingImports?.length) {
+                await ImportAdderProvider.instance.onAcceptRecommendation(editor, item, startLine)
+            }
         }
         commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
 
@@ -200,6 +205,7 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
                     item,
                     editor,
                     session.requestStartTime,
+                    position.line,
                     session.firstCompletionDisplayLatency,
                 ],
             }
@@ -208,6 +214,7 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
                 item.insertText as string,
                 item.references
             )
+            ImportAdderProvider.instance.onShowRecommendation(document, position.line, item)
         }
         return items as InlineCompletionItem[]
     }

packages/amazonq/src/extension.ts

@@ -32,6 +32,7 @@ import {
     setupUninstallHandler,
     maybeShowMinVscodeWarning,
     Experiments,
+    isSageMaker,
 } from 'aws-core-vscode/shared'
 import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -120,7 +121,9 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     await activateCodeWhisperer(extContext as ExtContext)
     if (Experiments.instance.get('amazonqLSP', false)) {
         await activateAmazonqLsp(context)
-    } else {
+    }
+
+    if (!Experiments.instance.get('amazonqLSPInline', false)) {
         await activateInlineCompletion()
     }
 
@@ -144,6 +147,9 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // Hide the Amazon Q tree in toolkit explorer
     await setContext('aws.toolkit.amazonq.dismissed', true)
 
+    // set context var to check if its SageMaker Unified Studio or not
+    await setContext('aws.isSageMakerUnifiedStudio', isSageMaker('SMUS'))
+
     // reload webviews
     await vscode.commands.executeCommand('workbench.action.webview.reloadWebviewAction')
 
@@ -157,7 +163,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
 
     context.subscriptions.push(
         Experiments.instance.onDidChange(async (event) => {
-            if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP') {
+            if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') {
                 await vscode.window
                     .showInformationMessage(
                         'Amazon Q LSP setting has changed. Reload VS Code for the changes to take effect.',

packages/amazonq/src/lsp/activation.ts

@@ -6,22 +6,14 @@
 import vscode from 'vscode'
 import { startLanguageServer } from './client'
 import { AmazonQLspInstaller } from './lspInstaller'
-import { Commands, lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
+import { lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
 
 export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
     try {
         await lspSetupStage('all', async () => {
             const installResult = await new AmazonQLspInstaller().resolve()
             await lspSetupStage('launch', async () => await startLanguageServer(ctx, installResult.resourcePaths))
         })
-        ctx.subscriptions.push(
-            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-                await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
-            }),
-            vscode.workspace.onDidCloseTextDocument(async () => {
-                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
-            })
-        )
     } catch (err) {
         const e = err as ToolkitError
         void vscode.window.showInformationMessage(`Unable to launch amazonq language server: ${e.message}`)

packages/amazonq/src/lsp/chat/messages.ts

@@ -9,6 +9,7 @@ import {
     AUTH_FOLLOW_UP_CLICKED,
     CHAT_OPTIONS,
     COPY_TO_CLIPBOARD,
+    AuthFollowUpType,
 } from '@aws/chat-client-ui-types'
 import {
     ChatResult,
@@ -25,6 +26,7 @@ import { window } from 'vscode'
 import { Disposable, LanguageClient, Position, State, TextDocumentIdentifier } from 'vscode-languageclient'
 import * as jose from 'jose'
 import { AmazonQChatViewProvider } from './webviewProvider'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 
 export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) {
     languageClient.onDidChangeState(({ oldState, newState }) => {
@@ -77,10 +79,33 @@ export function registerMessageListeners(
                 })
                 break
             }
-            case AUTH_FOLLOW_UP_CLICKED:
-                // TODO hook this into auth
+            case AUTH_FOLLOW_UP_CLICKED: {
                 languageClient.info('[VSCode Client] AuthFollowUp clicked')
+                const authType = message.params.authFollowupType
+                const reAuthTypes: AuthFollowUpType[] = ['re-auth', 'missing_scopes']
+                const fullAuthTypes: AuthFollowUpType[] = ['full-auth', 'use-supported-auth']
+
+                if (reAuthTypes.includes(authType)) {
+                    try {
+                        await AuthUtil.instance.reauthenticate()
+                    } catch (e) {
+                        languageClient.error(
+                            `[VSCode Client] Failed to re-authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}`
+                        )
+                    }
+                }
+
+                if (fullAuthTypes.includes(authType)) {
+                    try {
+                        await AuthUtil.instance.secondaryAuth.deleteConnection()
+                    } catch (e) {
+                        languageClient.error(
+                            `[VSCode Client] Failed to authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}`
+                        )
+                    }
+                }
                 break
+            }
             case chatRequestType.method: {
                 const partialResultToken = uuidv4()
                 const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) =>