Skip to content

Commit b896e4f

Browse files
feat(amazonq): Use new auth class using flare identity server for Amazon Q (#7062)
## Problem The business logic of auth-related features is complex and implemented separately for all IDEs. Flare/DEXP LSP server has implemented auth that can be reused for all IDEs. The source code was introduced, but not yet instantiated upon activation. Follow up for #6958 ## Solution * Remove `"amazonqLSP"` experiment flag * Have [regionProfileManager.ts](https://github.com/aws/aws-toolkit-vscode/pull/7062/files#diff-f5716036a70a8b212e9dae829ec09e33731f181c173a3f2b9ad9d361d51af14b) depend on `authProvider` through `IAuthProvider` interface to remove circular dependency with `AuthUtil` * Deprecate `getChatAuthState()` API in favor of new `getAuthState()` * Instantiate and activate LSP client and new AuthUtil instance NOTE: CI is expected to fail, since the reference updates for AuthUtil in the rest of code and unit tests have not been committed yet --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Justin M. Keyes <[email protected]>
1 parent 713e27e commit b896e4f

File tree

15 files changed

+238
-147
lines changed

15 files changed

+238
-147
lines changed

docs/lsp.md

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@ sequenceDiagram
4848
3. Enable the lsp experiment:
4949
```
5050
"aws.experiments": {
51-
"amazonqLSP": true,
5251
"amazonqLSPInline": true, // optional: enables inline completion from flare
5352
"amazonqLSPChat": true // optional: enables chat from flare
5453
}

packages/amazonq/src/api.ts

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseReques
88
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
99
import { ChatSession } from 'aws-core-vscode/codewhispererChat'
1010
import { api } from 'aws-core-vscode/amazonq'
11+
import { getLogger } from 'aws-core-vscode/shared'
1112

1213
export default {
1314
chatApi: {
@@ -26,8 +27,25 @@ export default {
2627
await AuthUtil.instance.showReauthenticatePrompt()
2728
}
2829
},
30+
/**
31+
* @deprecated use getAuthState() instead
32+
*
33+
* Legacy function for callers who expect auth state to be granular amongst Q features.
34+
* Auth state is consistent between features, so getAuthState() can be consumed safely for all features.
35+
*
36+
*/
2937
async getChatAuthState() {
30-
return AuthUtil.instance.getChatAuthState()
38+
getLogger().warn('Warning: getChatAuthState() is deprecated. Use getAuthState() instead.')
39+
const state = AuthUtil.instance.getAuthState()
40+
const convertedState = state === 'notConnected' ? 'disconnected' : state
41+
return {
42+
codewhispererCore: convertedState,
43+
codewhispererChat: convertedState,
44+
amazonQ: convertedState,
45+
}
46+
},
47+
getAuthState() {
48+
return AuthUtil.instance.getAuthState()
3149
},
3250
},
3351
} satisfies api

packages/amazonq/src/app/chat/activation.ts

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ export async function activate(context: ExtensionContext) {
2727

2828
const setupLsp = funcUtil.debounce(async () => {
2929
void amazonq.LspController.instance.trySetupLsp(context, {
30-
startUrl: AuthUtil.instance.startUrl,
30+
startUrl: AuthUtil.instance.connection?.startUrl,
3131
maxIndexSize: CodeWhispererSettings.instance.getMaxIndexSize(),
3232
isVectorIndexEnabled: CodeWhispererSettings.instance.isLocalIndexEnabled(),
3333
})
@@ -44,8 +44,6 @@ export async function activate(context: ExtensionContext) {
4444
amazonq.walkthroughSecurityScanExample.register(),
4545
amazonq.openAmazonQWalkthrough.register(),
4646
amazonq.listCodeWhispererCommandsWalkthrough.register(),
47-
amazonq.focusAmazonQPanel.register(),
48-
amazonq.focusAmazonQPanelKeybinding.register(),
4947
amazonq.tryChatCodeLensCommand.register(),
5048
vscode.workspace.onDidChangeConfiguration(async (configurationChangeEvent) => {
5149
if (configurationChangeEvent.affectsConfiguration('amazonQ.workspaceIndex')) {

packages/amazonq/src/extension.ts

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
import { AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
6+
import { CredentialsStore, LoginManager, authUtils, initializeAuth } from 'aws-core-vscode/auth'
77
import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
88
import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
99
import { CommonAuthWebview } from 'aws-core-vscode/login'
@@ -33,6 +33,7 @@ import {
3333
maybeShowMinVscodeWarning,
3434
Experiments,
3535
isSageMaker,
36+
Commands,
3637
} from 'aws-core-vscode/shared'
3738
import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
3839
import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -117,11 +118,11 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
117118
const extContext = {
118119
extensionContext: context,
119120
}
121+
122+
await activateAmazonqLsp(context)
123+
120124
// This contains every lsp agnostic things (auth, security scan, code scan)
121125
await activateCodeWhisperer(extContext as ExtContext)
122-
if (Experiments.instance.get('amazonqLSP', false)) {
123-
await activateAmazonqLsp(context)
124-
}
125126

126127
if (!Experiments.instance.get('amazonqLSPInline', false)) {
127128
await activateInlineCompletion()
@@ -130,6 +131,10 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
130131
// Generic extension commands
131132
registerGenericCommands(context, amazonQContextPrefix)
132133

134+
// Create status bar and reference log UI elements
135+
void Commands.tryExecute('aws.amazonq.refreshStatusBar')
136+
void Commands.tryExecute('aws.amazonq.updateReferenceLog')
137+
133138
// Amazon Q specific commands
134139
registerCommands(context)
135140

@@ -156,7 +161,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
156161
// reload webviews
157162
await vscode.commands.executeCommand('workbench.action.webview.reloadWebviewAction')
158163

159-
if (AuthUtils.ExtensionUse.instance.isFirstUse()) {
164+
if (authUtils.ExtensionUse.instance.isFirstUse()) {
160165
// Give time for the extension to finish initializing.
161166
globals.clock.setTimeout(async () => {
162167
CommonAuthWebview.authSource = ExtStartUpSources.firstStartUp
@@ -166,7 +171,7 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
166171

167172
context.subscriptions.push(
168173
Experiments.instance.onDidChange(async (event) => {
169-
if (event.key === 'amazonqLSP' || event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') {
174+
if (event.key === 'amazonqChatLSP' || event.key === 'amazonqLSPInline') {
170175
await vscode.window
171176
.showInformationMessage(
172177
'Amazon Q LSP setting has changed. Reload VS Code for the changes to take effect.',

packages/amazonq/src/extensionNode.ts

Lines changed: 33 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -7,29 +7,22 @@ import * as vscode from 'vscode'
77
import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension'
88
import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
99
import { activate as activateQGumby } from 'aws-core-vscode/amazonqGumby'
10-
import {
11-
ExtContext,
12-
globals,
13-
CrashMonitoring,
14-
getLogger,
15-
isNetworkError,
16-
isSageMaker,
17-
Experiments,
18-
} from 'aws-core-vscode/shared'
10+
import { ExtContext, globals, CrashMonitoring /* getLogger, isSageMaker,*/, Experiments } from 'aws-core-vscode/shared'
1911
import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode'
2012
import { updateDevMode } from 'aws-core-vscode/dev'
2113
import { CommonAuthViewProvider } from 'aws-core-vscode/login'
2214
import { isExtensionActive, VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
2315
import { registerSubmitFeedback } from 'aws-core-vscode/feedback'
2416
import { DevOptions } from 'aws-core-vscode/dev'
25-
import { Auth, AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection } from 'aws-core-vscode/auth'
17+
import { Auth /* , AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection*/ } from 'aws-core-vscode/auth'
2618
import api from './api'
2719
import { activate as activateCWChat } from './app/chat/activation'
2820
import { activate as activateInlineChat } from './inlineChat/activation'
2921
import { beta } from 'aws-core-vscode/dev'
30-
import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
31-
import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer'
32-
import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
22+
import * as amazonq from 'aws-core-vscode/amazonq'
23+
import { /* activate as activateNotifications,*/ NotificationsController } from 'aws-core-vscode/notifications'
24+
// import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer'
25+
// import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
3326

3427
export async function activate(context: vscode.ExtensionContext) {
3528
// IMPORTANT: No other code should be added to this function. Place it in one of the following 2 functions where appropriate.
@@ -58,6 +51,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
5851
}
5952
activateInlineChat(context)
6053

54+
context.subscriptions.push(amazonq.focusAmazonQPanel.register(), amazonq.focusAmazonQPanelKeybinding.register())
55+
6156
const authProvider = new CommonAuthViewProvider(
6257
context,
6358
amazonQContextPrefix,
@@ -78,49 +73,34 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
7873
await setupDevMode(context)
7974
await beta.activate(context)
8075

81-
// TODO: Should probably emit for web as well.
76+
// TODO: @opieter Fix telemetry
8277
// Will the web metric look the same?
83-
telemetry.auth_userState.emit({
84-
passive: true,
85-
result: 'Succeeded',
86-
source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
87-
...(await getAuthState()),
88-
})
89-
90-
void activateNotifications(context, getAuthState)
78+
// telemetry.auth_userState.emit({
79+
// passive: true,
80+
// result: 'Succeeded',
81+
// source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
82+
// ...(await getAuthState()),
83+
// })
84+
85+
// void activateNotifications(context, getAuthState)
9186
}
9287

93-
async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
94-
let authState: AuthState = 'disconnected'
95-
try {
96-
// May call connection validate functions that try to refresh the token.
97-
// This could result in network errors.
98-
authState = (await AuthUtil.instance._getChatAuthState(false)).codewhispererChat
99-
} catch (err) {
100-
if (
101-
isNetworkError(err) &&
102-
AuthUtil.instance.conn &&
103-
AuthUtil.instance.auth.getConnectionState(AuthUtil.instance.conn) === 'valid'
104-
) {
105-
authState = 'connectedWithNetworkError'
106-
} else {
107-
throw err
108-
}
109-
}
110-
const currConn = AuthUtil.instance.conn
111-
if (currConn !== undefined && !(isAnySsoConnection(currConn) || isSageMaker())) {
112-
getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
113-
}
114-
115-
return {
116-
authStatus:
117-
authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
118-
? authState
119-
: 'notConnected',
120-
authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
121-
...(await getTelemetryMetadataForConn(currConn)),
122-
}
123-
}
88+
// async function getAuthState(): Promise<Omit<auth2.AuthState, 'source'>> {
89+
// let authState: AuthState = 'disconnected'
90+
// authState = AuthUtil.instance.getAuthState()
91+
92+
// if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
93+
// getLogger().error('Current Amazon Q connection is not SSO')
94+
// }
95+
96+
// return {
97+
// authStatus:
98+
// authState === 'connected' || authState === 'expired'
99+
// ? authState
100+
// : 'notConnected',
101+
// ...(await getTelemetryMetadataForConn(currConn)),
102+
// }
103+
// }
124104

125105
/**
126106
* Some parts of this do not work in Web mode so we need to set Dev Mode up here.

packages/amazonq/src/lsp/activation.ts

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,20 @@
44
*/
55

66
import vscode from 'vscode'
7-
import { startLanguageServer } from './client'
7+
import { clientId, encryptionKey, startLanguageServer } from './client'
88
import { AmazonQLspInstaller } from './lspInstaller'
9-
import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared'
9+
import { lspSetupStage, ToolkitError } from 'aws-core-vscode/shared'
10+
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
11+
import { auth2 } from 'aws-core-vscode/auth'
1012

11-
export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
13+
export async function activate(ctx: vscode.ExtensionContext) {
1214
try {
13-
await lspSetupStage('all', async () => {
15+
const client = await lspSetupStage('all', async () => {
1416
const installResult = await new AmazonQLspInstaller().resolve()
15-
await lspSetupStage('launch', async () => await startLanguageServer(ctx, installResult.resourcePaths))
17+
return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths))
1618
})
19+
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
20+
await AuthUtil.instance.restore()
1721
} catch (err) {
1822
const e = err as ToolkitError
1923
void messages.showViewLogsMessage(`Failed to launch Amazon Q language server: ${e.message}`)

0 commit comments

Comments
 (0)